[論文レビュー] 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のインstructチューニングでは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}.

研究の動機と目的

• 大規模なデータセット上で大規模な深層学習モデルをトレーニングする際の高い計算コストに対処すること。
• 静的および動的データプルーニング手法が導入する勾配期待値バイアスを克服すること。
• 損失なしの学習高速化を実現する、プラグアンドプレイ可能でアーキテクチャに依存しないフレームワークを開発すること。
• 既存のデータ圧縮技術（コアセット選択やデータセット量子化など）と互換性を持つこと。

提案手法

• InfoBatchは、順伝搬中に各サンプルの累積損失スコアを維持し、損失が低い（よく学習された）サンプルを特定する。
• 各エポックにおいて、損失が最小のサンプルの固定比率をランダムにプルーニングし、ソートやヒューリスティックスコア計算のオーバーヘッドを回避する。
• 残りのサンプルの勾配を再スケーリングして、元の勾配期待値を保持する。これにより最適化が偏りなく保たれる。
• 期待値再スケーリングを用いて、元のデータセットにおけるフルバッチ学習と同一の勾配期待値の方向を維持する。
• トレーニング中にリアルタイムの損失値に基づき動的にプルーニングを調整し、オンラインで低オーバーヘッドの適応が可能になる。
• 最終数エポックでは、収束安定性の向上と分散の低減を目的に、フルデータセットでの学習を回復する。

実験結果

リサーチクエスチョン

• RQ1 動的データプルーニングは、勾配期待値バイアスを導入せずに損失なしの学習高速化を達成できるか？
• RQ2 単純で低オーバーヘッドのプルーニング戦略は、ヒューリスティックまたは反復的手法を上回る速度向上と精度を達成できるか？
• RQ3 サンプルの削除にもかかわらず、勾配再スケーリングが元の最適化軌道を効果的に保持できるか？
• RQ4 InfoBatchは、視覚と自然言語処理を含む多様な深層学習タスクに、アーキテクチャの変更なしに普遍的に適用可能か？
• RQ5 トレーニングコスト削減と計算オーバーヘッドの観点で、InfoBatchは最先端手法と比べてどのように差をつけるか？

主な発見

• InfoBatchは、CIFAR-10/100、ImageNet-1K、ADE20Kの各データセットで、フルバッチ学習と同一の精度を維持しながら、全体のトレーニングコストを40%削減した。
• マスクされた自己符号化（MAE）事前学習では、24.8%のコスト削減が達成され、画像再構築品質に劣化は見られなかった。
• 拡散モデルの事前学習では、27%のトレーニングコスト削減が達成され、生成画像の品質に影響はなかった。
• LLaMA-7Bのインstructチューニングでは、20%のトレーニングコスト削減が達成され、BBH、DROP、MMLU、HumanEvalのベンチマークスコアは最先端水準を維持した。
• EL2Nのような最先端手法（スコア推定にO(MNT)時間が必要）と比較して、サンプルスコア計算のオーバーヘッドを10倍以上削減した。
• このフレームワークはコアセット選択およびデータセット量子化と完全に互換性があり、既存の圧縮技術と組み合わせてさらなるコスト削減が可能である。