AIの分業学習で使う新手法、最初の担当分けで一部に仕事が偏る弱点が見つかる

AIの混合エキスパート（MoE）モデルの学習において、最適化手法「Muon」を採用すると、最初の層で負荷の偏りが起きやすくなることが研究で明らかになった。

実験では、5500万および1億800万の有効パラメータを持つモデルを検証した。短い学習で調整した学習率において、Muonを適用したところ、他の層に比べて第1層のバイアスの広がりが平均で約2倍に拡大した。これにより最初の振り分け役において、最多処理を担うエキスパートに公平な配分の4〜7倍の負荷が集中する現象が確認された。

この過負荷に対し、最初の3,000ステップで入力値からバッチ平均を差し引く処理を施すことで、損失をほぼ変えずに偏りを解消できたという。また、小型モデルではMuonの学習率を半減させることでも偏りの大部分が抑えられ、損失のわずかな改善が見られたとしている。