LLMにおけるアクティベーション圧縮：理論解析と効率的アルゴリズム

要旨: 大規模言語モデル（LLM）の学習は、メモリ消費が非常に大きい。学習では、重みやオプティマイザ状態だけでなく、逆伝播のための中間活性も保存する必要があるためである。既存のメモリ効率化手法は主として勾配およびオプティマイザ状態に着目しているが、活性圧縮は、LLMに特化した理論や保証が欠如していることから、十分に確立されていない。本研究では、活性圧縮が「活性圧縮が偏りのない（unbiased）」場合に限り、線形作用素に対して安全である一方、非線形作用素に対しては問題が生じることを示す理論的枠組みを構築する。さらに、勾配分散の上界を導出し、標準的なL-滑らかさ（smoothness）仮定のもとで、全ての線形作用素に活性圧縮を適用する際の収束保証を確立し、これが収束率を変えないことを示す。理論に導かれて、低ランクの活性因子を再利用し、追加の計算や追加の勾配誤差なしに線形層の勾配を圧縮する、活性勾配の共同圧縮（activation-gradient co-compression）手法を提案する。事前学習ベンチマークおよび複数のファインチューニングベンチマークを用いて、QwenおよびLLaMAモデルに関する大規模な実験を実施し、理論を検証するとともに、精度と圧縮効率の両面で本手法が競争力のある性能を示すことを実証する。再現性のため、コードは補足資料に提供する。