[论文解读] InfoBatch: Lossless Training Speed Up by Unbiased Dynamic Data Pruning

Ziheng Qin, Kai Wang+9 more|arXiv (Cornell University)|Mar 8, 2023

Advanced Neural Network Applications被引用 7

一句话总结

InfoBatch 提出了一种新颖的无偏动态数据剪枝框架，通过随机剪枝低损失（已充分学习）样本并重新缩放梯度以保持原始梯度期望，实现在不损失性能的前提下加速深度学习训练。该方法在 ImageNet-1K 上实现高达 40% 的无损训练加速，在扩散模型上实现 27% 的加速，在 LLaMA 指令微调上实现 20% 的加速，且开销极低，适用于视觉与自然语言处理任务。

ABSTRACT

Data pruning aims to obtain lossless performances with less overall cost. A common approach is to filter out samples that make less contribution to the training. This could lead to gradient expectation bias compared to the original data. To solve this problem, we propose extbf{InfoBatch}, a novel framework aiming to achieve lossless training acceleration by unbiased dynamic data pruning. Specifically, InfoBatch randomly prunes a portion of less informative samples based on the loss distribution and rescales the gradients of the remaining samples to approximate the original gradient. As a plug-and-play and architecture-agnostic framework, InfoBatch consistently obtains lossless training results on classification, semantic segmentation, vision pertaining, and instruction fine-tuning tasks. On CIFAR10/100, ImageNet-1K, and ADE20K, InfoBatch losslessly saves 40\% overall cost. For pertaining MAE and diffusion model, InfoBatch can respectively save 24.8\% and 27\% cost. For LLaMA instruction fine-tuning, InfoBatch is also able to save 20\% cost and is compatible with coreset selection methods. The code is publicly available at \href{https://github.com/henryqin1997/InfoBatch}{github.com/NUS-HPC-AI-Lab/InfoBatch}.

研究动机与目标

• 为解决大规模深度学习模型在海量数据集上训练的高计算成本问题。
• 克服静态与动态数据剪枝方法因丢弃样本而引入的梯度期望偏差问题。
• 开发一种即插即用、与架构无关的框架，实现无损训练加速且附加开销极低。
• 确保与现有数据压缩技术（如 coreset 选择与数据集量化）的兼容性。

提出的方法

• InfoBatch 在前向传播过程中为每个样本维护一个运行损失分数，以识别低损失（已充分学习）样本。
• 在每个训练周期中，它随机剪枝最低损失样本的固定比例，避免排序或启发式评分的开销。
• 对剩余样本的梯度进行重缩放，以保持总梯度更新的期望值，确保优化无偏。
• 该框架使用期望重缩放，使优化方向与原始数据集上的全批量训练保持一致。
• 它根据训练期间的实时损失值动态调整剪枝策略，实现在线、低开销的自适应调整。
• 在最后几个周期中，恢复全量数据训练，以提升收敛稳定性并降低方差。

实验结果

研究问题

• RQ1 动态数据剪枝能否在不引入梯度期望偏差的前提下实现无损训练加速？
• RQ2 一种简单、低开销的剪枝策略能否在速度提升与准确率方面超越现有启发式或迭代方法？
• RQ3 梯度重缩放是否能有效保持样本移除后的原始优化轨迹？
• RQ4 InfoBatch 是否能无需架构修改地广泛应用于包括视觉与自然语言处理在内的多样化深度学习任务？
• RQ5 与最先进方法相比，InfoBatch 在训练成本降低与计算开销方面表现如何？

主要发现

• InfoBatch 在 CIFAR-10/100、ImageNet-1K 和 ADE20K 数据集上实现整体训练成本降低 40%，同时保持与全批量训练相同的准确率。
• 在掩码自编码（MAE）预训练中，InfoBatch 实现 24.8% 的成本降低，且图像重建质量无下降。
• 在扩散模型预训练中，InfoBatch 节省 27% 的训练成本，且不影响生成图像质量。
• 在 LLaMA-7B 的指令微调中，InfoBatch 实现 20% 的训练成本降低，同时在 BBH、DROP、MMLU 和 HumanEval 基准测试中保持最先进性能。
• 与最先进方法（如 EL2N）相比，InfoBatch 将样本评分开销降低至少 10 倍，后者需要 O(MNT) 时间进行评分估计。
• 该框架与 coreset 选择和数据集量化完全兼容，可与现有压缩技术结合使用，实现进一步的成本降低。