3行のコードで損失なく40%高速化、You YangチームのAIトレーニングアクセラレータがICLR口頭発表論文に選出

プルーニングを使用すると、AI トレーニングを高速化し、ロスレス操作を実現できます。わずか 3 行のコードで実行できます。

今年のトップクラスのディープラーニングカンファレンス ICLR で、シンガポール国立大学の You Yang 教授のチームの成果が口頭発表されました。

このテクノロジーを使用することで、損失なくトレーニング コストを最大 40% 節約できます。

この成果は InfoBatch と呼ばれ、サンプルをプルーニングする加速方式を引き続き使用します。

しかし、InfoBatch は、プルーニングの内容を動的に調整することで、加速によって生じるトレーニング損失の問題を解決します。

プラグアンドプレイで、アーキテクチャに制限されません。CNN ネットワークと Transformer モデルの両方を最適化できます。

現在、このアルゴリズムは多くのクラウドコンピューティング企業の注目を集めています。

では、InfoBatch はどのような加速効果を実現できるのでしょうか?

ロスレスによりトレーニングコストが40%削減

研究チームが複数のデータセットに対して実施した実験。これらすべてが InfoBatch の有効性を検証しました。

これらの実験では、画像の分類、セグメンテーション、生成などのタスクや、言語モデルの命令の微調整などがカバーされます。

画像分類タスクでは、研究チームは CIFAR10 および CIFAR100 データセットを使用して ResNet-18 をトレーニングしました。

結果は、剪定率が 30%、50%、70% の場合、InfoBatch の精度がランダム剪定やその他のベースライン方法の精度を上回り、剪定率が 30% の場合、精度が低下しないことを示しています。

剪定率が 30% から 70% に増加すると、InfoBatch の精度損失も他の方法に比べて大幅に低くなります。

ImageNet-1K データセットでトレーニングされた ResNet-50 を、プルーニング率 40%、エポック 90 で使用すると、InfoBatch は UCB と同じトレーニング時間を実現しながら、より高い精度を実現し、完全なデータ トレーニングを上回ることさえあります。

同時に、ImageNet の追加 (OverHead) 時間コストは他の方法に比べて大幅に低く、わずか 0.0028 時間、つまり 10 秒です。

Vit-Base モデルをトレーニングする場合(事前トレーニング フェーズで 300 エポック、微調整フェーズで 100 エポック) 、InfoBatch は24.8% のコスト削減率で、完全なトレーニングと同じ精度を維持できます。

クロスアーキテクチャテストの比較結果からも、InfoBatch はさまざまなモデルアーキテクチャに対して強力な堅牢性を発揮することが示されています。

さらに、InfoBatch は既存のオプティマイザーと互換性があり、さまざまなオプティマイザーと併用した場合に優れたロスレスの加速効果を発揮します。

これらの視覚タスクだけでなく、InfoBatch は言語モデルの教師あり微調整にも適用できます。

常識 (MMLU)、推論 (BBH、DROP) などの能力に大きな損失がなく、プログラミング能力 (HumanEval) がわずかに向上するだけで、InfoBatch は DQ に基づいて時間消費をさらに 20% 削減できます。

さらに、作者の最新のアップデートによると、InfoBatch は検出タスク (YOLOv8) で 30% のロスレス高速化も達成しており、コードは github で更新される予定です。

では、InfoBatch はどのようにしてロスレス加速を実現するのでしょうか?

剪定内容を動的に調整する

核となる秘密は、偏りのない動的なデータプルーニングです。

従来のプルーニング方法における勾配期待方向の偏差と総更新量の削減の問題を解決するために、InfoBatch は動的プルーニング方法を採用しています。

InfoBatch の順方向伝播中、各サンプルのスコア (損失) は維持され、平均をしきい値として、スコアの低いサンプルの一定割合がランダムに削除されます。

同時に、予想される勾配更新を維持するために、残りの低スコアのサンプルの勾配がそれに応じて増幅されます。

このようにして、InfoBatch トレーニング結果と元のデータのトレーニング結果間のパフォーマンス ギャップが、以前の方法と比較して改善されます。

具体的には、トレーニングのフォワードプロセス中に、InfoBatch はサンプルの損失値をサンプルスコアとして記録するため、基本的に追加のスコアリングオーバーヘッドは発生しません。

最初のエポックでは、InfoBatch はデフォルトですべてのサンプルを初期化して保持します。その後の各エポックの開始前に、InfoBatch はプルーニング確率 r に従って、平均よりも低いスコアを持つサンプルをランダムにプルーニングします。

確率の具体的な表現は次のようになります。

スコアが平均より低いがトレーニングを続行するために保持されるサンプルの場合、InfoBatch は再スケーリング メソッドを使用して対応する勾配を 1/(1-r) に増やし、全体的な更新をほぼ偏りのないものにしています。

さらに、InfoBatch は、トレーニングの後半の段階で完全なデータセットを使用して、段階的なプルーニング プロセスを使用します。

その理由は、理論上の期待更新は基本的に同じであるものの、上記の期待値には実際には時刻 t における複数の値が含まれているためです。

つまり、サンプルが途中の特定のラウンドで削減された場合、後でトレーニングされる可能性はまだありますが、残りの更新ラウンドが不十分な場合、この確率は大幅に低下し、残差勾配期待偏差が発生します。

したがって、最後の数回のトレーニング ラウンド (通常は約 12.5% ～ 17.5%) では、InfoBatch は完全な元のデータをトレーニングに使用します。

論文アドレス: https://arxiv.org/abs/2303.04947
GitHub ホームページ:
https://github.com/NUS-HPC-AI-Lab/InfoBatch