Dust：逆伝播を使わないTransformer事前学習への挑戦

導入

現代の深層学習は、ほぼ一貫して逆伝播を中心に発展してきた。ネットワークの微分可能性を利用して勾配を計算し、オプティマイザーで重みを更新する仕組みである。Dustは、この前提を将来的に緩められるかを検討する研究だ。十分な計算資源があるなら、解析的な勾配よりも広い探索に基づく学習が競争力を持つ可能性がある。

Dustの仕組み

Dustはゼロ次最適化アルゴリズムであり、損失を解析的に微分しない。内部の活性値に摂動を加え、その結果として損失がどのように変化したかを測定する。損失を下げる方向の摂動に大きな報酬を与え、複数の摂動を重み付きで平均することで、更新方向を推定する。

特徴は、探索対象が重みではなく活性値である点だ。さらに、摂動を各トークンに独立して加えることで、各トークンを「仮想的な集団」の一員として扱う。同じフォワード計算の中で多数の候補を並列評価でき、候補ごとにモデルを複製して保存する必要がない。Transformerの層の種類ごとにトークン単位の報酬規則を変え、摂動したモジュール同士の干渉を抑える工夫も組み込まれている。

主なポイント

• Transformer言語モデルの事前学習で、逆伝播に近い性能を示すゼロ次手法の一つだと報告されている。
• 仮想集団を大きくすると、Dustの勾配推定は逆伝播の勾配とより整合する。複数の設定では逆伝播を上回ったという。
• 重み空間で進化戦略を実行する方法に比べ、活性値空間の探索は評価とメモリの負担を抑えられる。論文の推定では、一定以上のトークン規模で、比較対象のEGGROLL実装より数桁効率的になり得る。
• 大きなモデルほど集団を効率よく利用するケースがあり、ゼロ次手法はモデルの大型化で必ずしも不利になるとは限らない。
• ただし、現時点で逆伝播を低コストに置き換えられるという主張ではない。著者らも、実用的な計算効率にはまだ達していないと位置づけている。

意義と限界

Dustの重要性は、すぐに新しい標準訓練法になることよりも、大規模ニューラルネットで試せる学習アルゴリズムの範囲を広げた点にある。活性値を探索すれば、内部表現や潜在的な推論過程を直接探索する方向も考えられる。また、外部プログラムや微分できない部品を含むシステムに、別の信用割当て方法を適用する余地も生まれる。

一方、逆伝播に近づくには大きな集団と追加の計算が必要になりやすい。今回の結果は実現可能性とスケーリングの傾向を示すものであり、コスト、データ効率、最終性能のすべてで逆伝播を上回ることを意味しない。今後は、さらに大きなモデルや、非微分可能な構成要素を本当に必要とするアーキテクチャで優位性が続くかが焦点になる。

出典：Hacker News