※本記事は、Pluralis社のResearch ScientistであるSameera Ramasinghe氏による講演「Communication-Efficient Model-Parallel Training」の内容を基に作成されています。本講演は、Pluralis Researchが主催し、ブラジル・リオデジャネイロで開催されたICLR 2026「Protocol Learning Workshop」にて収録されたものです。同ワークショップには、分散型かつ低帯域幅のネットワーク環境において大規模モデルを学習させる研究に取り組む研究者たちが集まりました。Protocol Learningに関する詳細情報は https://pluralis.ai でご覧いただけます。本記事では、講演の内容を要約しております。 なお、本記事の内容は登壇者の見解を正確に反映するよう努めていますが、要約や解釈による誤りがある可能性もありますので、正確な情報や文脈については、オリジナルの講演をご視聴いただくことをお勧めいたします。
1. 研究への向き合い方:理論と実践のギャップ
1.1 リソース制約下での意思決定と理論の役割
Sameera： 私が今日お話しする内容の大部分は、昨年私たちが取り組んだ主要な研究、特に圧縮アルゴリズムに関するものになります。ただ、単に手法の話にとどまらず、私たちがこうした問題にどう向き合うべきか、つまり会社としてどのようなアプローチを取っているかについてもお話ししたいと思っています。というのも、こうした問題への取り組み方そのものについて、ある種のメンタルモデルの転換が必要だと考えているからです。なぜなら、これらの技術は最終的に大規模な環境で機能しなければならないからです。結局のところ、私たちはたくさんの研究や理論的な検討を行うことはできますが、それらは実際の現場で機能しなければ意味がありません。そして、大規模な環境でアブレーションやグリッドサーチの実験を数多く行うということは、事実上不可能に近いのです。特にフロンティアラボの外で研究をしている場合はそうです。フロンティアラボには、いわゆる「golden ablation models」と呼ばれる贅沢な環境があります。彼らは小規模でグリッドサーチの実験を行い、それをスケールアップし、グリッドを絞り込み、さらに洗練させていくということを繰り返せるのです。それには何百万ドル、何十億ドルという計算資源が投じられています。しかし私たちにはそのような贅沢はありません。特に分散型(decentralized)の学習に取り組んでいる場合はなおさらです。要するに、リソースが制約された環境の中で仕事をしなければならないということです。ですから、こうした問題にどう向き合うかという点で、根本的に異なるアプローチが必要になってくるのです。
私が手法そのものよりも先にこの点についてお話ししたいのは、大規模にスケールアップする際に、こうした経験的な意思決定をどうすればより効果的に行えるかということが、昨年私たちが行った主要な研究とも密接につながっているからです。私は意図的に、理論から実践への「跳躍(leap)」という表現を使いました。というのも、これは単なる連続的な一歩ではないからです。そこには常に外挿(extrapolation)が伴います。連続的なものではないのですが、それでも私たちは根拠のある形でその跳躍を行わなければなりません。理論が果たす役割について少しお話ししますと、私の見方では、理論は私たちが大規模な環境で実験できる解の集合について、探索空間を絞り込む手段を与えてくれるものだと考えています。理論から実践への外挿について正しい直感を持ち、正しい外挿を行うことができれば、それによって多くの計算資源とお金を節約できる可能性があるのです。
1.2 分散学習の背景と通信ボトルネック
Sameera： これはもう、コンピュータビジョンにおけるレナ画像のような存在になりつつありますが、皆さんご存知の通り、モデルのスケールアップが進んでいます。現時点では、こうしたより大きなモデルを学習させ、良好な性能を発揮させられるのはフロンティアラボだけという状況です。なぜ分散型の学習が必要なのかという点については、この場にいる皆さんはすでに十分に馴染みがあると思いますので、あまり動機づけの説明はしないでおこうと思います。
主なボトルネックは帯域幅です。分散型の設定では、高速リンクで密に接続されたGPUクラスタを使えるという贅沢がありません。そのため、インターネット経由で通信を行わなければなりません。そうなると、圧縮アルゴリズムが極めて重要になってきます。そして、収束性を損なうことなくどうやってこれを実現するかが問題になります。このアプローチには大きく分けて2つの方法があります。1つはモデル並列、そしてもう1つはデータ並列です。
2. 分散学習における2つのアプローチとその課題
2.1 データ並列(DDP)の仕組みと限界
Sameera： モデル並列の考え方についてお話しします。基本的な発想としては、モデルのレプリカをノードごとに分割して配置し、この低帯域幅の接続を通じて通信を行いながらローカルな更新を行い、その後にオールリデュースのステップを実施するというものです。ローカルな更新を行った後、これをさらに効率化する方法としては、例えばPowerSGDのように勾配そのものを圧縮する手法や、DiLoCoのようなスパース化技術を用いる方法があります。データ並列、つまりDDPの問題点は、スケールできる量に限界があるということです。というのも、各ノードにモデルの完全なレプリカを持たせる必要があるからです。
とはいえ、この問題を解決することの重要性を決して過小評価しているわけではありません。実際、私たちはこの点について痛い経験を通じて学ぶことになりました。というのも、異種(heterogeneous)な環境ではこの方式には多くの問題が生じるからです。特に、デバイスやノードがデータや計算資源の面で異種混在している場合、素朴な通信アルゴリズムはそのままでは機能しません。ですから、専用のアルゴリズムを新たに考案する必要がありました。この問題を解決すること自体は非常に重要なことなのですが、限界としては、スケールアップが難しいという点が挙げられます。
一方、モデル並列では、モデルを異なる部分に分割し、それを複数のGPUに振り分けたうえで、活性化(activations)と活性化の勾配(activation gradients)を各デバイス間で通信します。
2.2 モデル並列における誤差蓄積問題
Sameera： この方式は、通信という観点でははるかに集約的なものになります。というのも、フォワードパスとバックワードパスのたびに、これらの活性化と活性化勾配をやり取りしなければならないからです。そのため、この通信を行うたびに誤差が蓄積していく可能性があります。そして、誤差がどれだけ蓄積しうるかについては、正確な上限を示すことができます。
具体的には、誤差は指数関数的に蓄積していく可能性があります。つまり、何らかの非可逆(lossy)な圧縮を用いている場合、モデルが深くなればなるほど誤差が積み重なっていき、最終的には信号が著しく劣化してしまい、収束性に非常に悪い影響を及ぼすことになります。ですから、私がここで少し主張しておきたいのは、モデル並列のために考案する圧縮アルゴリズムは、計算グラフの内部に組み込まれているべきだということです。例えば、スパース化や量子化、あるいはトップK選択のような手法が、フォワードパスやバックワードパスから切り離された形になっていて、計算グラフ側がその圧縮について何も認識していないとしたら、それは非常に良くない状態だということです。誤差が蓄積してしまうからです。つまり、活性化と活性化勾配に対してどのような圧縮を適用しているかについて、最適化のプロセス自体がそれを認識している必要があるのです。
では、この問題にどう取り組むべきでしょうか。ここで、こうした問題をどう捉えるべきかについて、少し話が逸れますが触れておきたいと思います。冒頭でも申し上げた通り、私たちは大規模な環境で多くのアブレーション実験を行うことはできません。しかし、理論がいつ、どのように機能するのか、そして理論から実践へどのように外挿すればよいのかについてしっかりとした理解を持っていれば、かなり良い、根拠のある意思決定や設計上の判断を下すことができます。次のいくつかのスライドでは、私たちが非常に気に入っていて、研究の多くの基盤とした重要な論文をいくつか紹介していきたいと思います。ただし、ここで一つ断っておきたいのは、これらの理論的な研究に含まれる強い仮定の一部を、私はかなり大胆に無視しているということです。それでも、これは単なる直感として有用なものです。私たちはこうした直感に基づいて、意思決定を行っているのです。
3. ニューラルネットワークの帰納バイアスに関する理論的考察
3.1 パラメータ空間の膨大さと収束の神秘性、低ランク構造という解の性質
Sameera： ここで少し話が脇道に逸れますが、この点について考えるたびに私はいつも驚かされます。皆さんもご存知の通りだとは思いますが、それでも私にとっては本当に不思議でなりません。300個のニューロンからなる二値ニューラルネットワークだけでも、観測可能な宇宙に存在する原子の数よりも100億倍も多くの配置の組み合わせを持っているという事実です。それにもかかわらず、私たちはこうしたニューラルネットワークを機能させ、有用なタスクをこなさせ、有用な解へと収束させることができています。もちろん、最適化から生じるバイアスなどによってなぜこれがうまくいくのかについて、私たちにはある程度の直感があります。しかし、それでもこのことを考えるたびに極めて驚かされますし、今でも私を驚かせ続けています。
理論的には、たとえ2層のニューラルネットワークであっても、線形しきい値活性化やReLUのようなものを用い、かつパラメータ数がデータに対して不足している(under-parameterized)状態であれば、最悪のケースにおいて誤差ゼロの大域最小値へ収束させることは計算量的に困難な(intractable)問題であることを示すことができます。しかし実際には、確率的勾配降下法のようなオプティマイザを使うと、ほとんど常に良い解へと収束します。これは何十年にもわたる研究がその理由を説明してきたとはいえ、今でも私にとっては非常に驚くべきことであり、素晴らしいことだと感じています。
これが可能になる唯一の理由は、深層学習における様々な側面から生じる帰納バイアス(inductive bias)によるものです。私がこの問題について考えている枠組みとしては、大規模な深層学習における任意の問題は、初期化、アーキテクチャ、最適化といった異なる側面から来る異なる帰納バイアスが存在しているというものです。これらの側面はそれぞれ異なる帰納バイアスを持っています。そして、もし私たちが端から端まで機能する完全なシステムを設計しようとするならば、これらを個別に切り離して見るべきではありません。もちろん、研究や科学的な調査を行う上ではそれぞれを個別に研究することは重要ですし、有用な知見も得られます。しかし、システムを端から端まで設計する際には、これらを統一的な視点から見ることが非常に有用なのです。
特に初期化について少し強調しておきたいと思います。というのも、ほとんどの場合、最適化とアーキテクチャについては議論されますが、初期化についてはそれほど議論されないからです。しかし私たちは数年前のICMLで、初期化がモデルの収束する解の性質、特に汎化性能を完全に決定しうることを示す論文を発表しました。そして今回のこのカンファレンスで偶然出会った最近の研究も、初期化が最終的なモデルの性質を完全に決定しうるというこの事実に言及しています。ニューラルネットワークのパラメータ空間は極めて広大であるにもかかわらず、それらが収束する解は極めて構造化されていて単純なものです。
例えば、低ランクという観点から見ると、収束時点でのモデルを観察すると、それらは構造的に単純であり、非常に構造化されていることが分かります。私はこれを、ある種の新しい形、あるいは現代版のオッカムの剃刀のようなものだと捉えています。つまり、ネットワークはデータに適合する最も単純な解へと収束するということです。そしてこれは異なる観点からも見ることができます。もう一つの見方はフーリエレンズからのものです。私たちはこれについても最近論文を書きましたが、モデルを学習させて訓練データに適合させると、それは最も滑らかな解、つまりフーリエの観点から見て訓練データに適合できる最も単純なモデルになるということを示しています。
そしてこれは、今日利用可能なほとんどの大規模モデルにも当てはまります。例えばLlamaやQwenのようなモデルを取り上げて重みを分析すると、そのほとんどが達成しうる最大のランク、つまりフルランクと比較して低ランクであることが分かります。そして、実際に動いている重みの割合、つまりフルパラメータ空間と比較して移動する重みの割合は、モデルの規模が大きくなるほど小さくなっていきます。
もう一つの証拠はモデルプルーニングです。私たちは、収束後のモデルを大幅にプルーニングしても、推論時にはうまく機能し続けることを知っています。つまり、驚くほど単純な構造を持つ解が存在するということが分かっているわけです。しかし問題は、最適化を通じてそこにたどり着けるのかということです。例えば、モデルに対して低ランクの制約を課すとした場合、最終的な解が低ランクであることは分かっていても、最初からその制約を課した状態で最適化を通じてそこに到達できるのでしょうか。射影勾配降下法の理論のようなものを考えると、一次定常点にはほぼ確実に収束すると言えますが、それが十分に良い最小値であるかどうかは分かりません。そこで問題になるのは、こうした制約を課しても良いモデルへ収束できるかどうかについて、どうすれば直感を得られるかということです。そしてその方向で、モデル内の低ランク構造を利用した圧縮アルゴリズムを考案し、通信の圧縮という利得を得ることができるのではないかということです。
これが可能になるのは、こうしたネットワークにおける帰納バイアスについて理解を持ち、それを統一的な視点で利用できた場合です。そして何らかの形で、これらの帰納バイアスを活用して活性化と活性化勾配を圧縮する方法を編み出せた場合です。
3.2 深さと幅がもたらす損失地形への影響
Sameera： その話に入る前に、どのような帰納バイアスに着目すべきか、そしてそれらがここでどう役立つのかについてお話ししたいと思います。まず一つ目は深さです。モデルのネットワークは非常に深く、多くの層を持っています。これは私が非常に気に入っている古典的な論文なのですが、そこで示されているのは、線形ネットワークにおいて層の数を増やしていくと、それはほとんど前提条件付け(preconditioner)のように働くということです。つまり、確率的勾配降下法やAdamを適用した場合、深さが損失地形を平滑化する演算子として機能するということです。そして最適化はより容易に、そしてより安定的になります。
ここで、私自身がこの現象の実践的な側面に偶然出会った経験をお話ししたいと思います。博士課程1年目の頃、私はこうしたバイアスについて何の知識も持っていませんでした。実は今でもそれほど詳しいわけではないのですが、当時ある特徴ベクトルを別の特徴ベクトルへ変換するモデルに取り組んでいました。機能上の制約があって、線形層を使わなければならなかったのです。そして厳しいニューロン予算の制約もありました。どうやっても、私がやりたいことがうまくいきませんでした。特徴変換が全くうまくいかず、勾配降下法は良い解に収束できなかったのです。
しかし、単なる好奇心から他のことを試していたときに、私は同じニューロン予算の中で複数の線形層を積み重ねてみました。ニューロンの数は同じで、複雑さも同じです。しかし今度は構造が異なります。そうすると、特徴変換が見事にうまくいったのです。私は「これはaha momentだ」と思いました。というのも、数学的には線形層を積み重ねても、それはやはり線形変換のままで等価だからです。しかし、深さから生じるバイアスと、それが最適化とどのように相互作用するかという点に、何か特有のものがあるということです。後になって、私はこの論文に出会い、まさにこれがそれだったのだと気づきました。
そしてもう一つ紹介したい論文があります。なぜこの話をしているかという文脈を説明しますと、私はモデルの深さや規模がすべてを容易にするということを動機づけたいのです。圧縮や、それに関する制約の設計は、損失地形が深さ方向や幅方向のスケーリングによってどう振る舞うかについて良い理解があれば、はるかに容易になります。それが分かっていれば、大規模において私たちが直面するあらゆる問題に対する解決策を設計するのがずっと簡単になるのです。
この論文が扱っているのは深さではなく幅についてです。簡単に言うと、ネットワークが十分に幅広く、ニューロンを追加して幅を広げ続けると、ある時点を境に一種の相転移が起こり、損失地形の中にもはや劣最適な局所解が存在しなくなるということが示されています。例えば、ネットワークが狭ければ、多くの劣最適な罠が存在します。しかし、ネットワークの幅を広げ続けると、property PTと呼ばれる性質が現れます。これは、パラメータ空間のある点にモデルがあるとして、それを少し摂動させても、モデルは弱い意味での大域最小値へと収束できるというものです。つまり要点は、ネットワークが十分に幅広ければ、損失地形はよく整った(well-conditioned)状態になり、劣最適な局所最小値がほとんど存在しなくなるということです。
これと似た論文もあります。そこで示されているのは、ネットワークが十分に幅広く、過剰パラメータ化されていれば、どの点からでも大域最小値への連続的な経路が存在するということです。これは本当に興味深い結果です。もちろん、こうした結果の多くは全結合ネットワークや特定の構造を前提とするなど、厳しい条件を仮定しています。
3.3 理論のTransformerへの外挿可能性
Sameera： しかし、こうした結果の多くは実際にはTransformerにおいてもさらによく成り立つと私は考えています。ただ、大規模なTransformerを数理的にモデル化するのが難しいというだけの話です。それがここでの問題です。しかし実際には、これが当てはまることを示す十分な証拠があると考えています。というのも、モデルマージや線形モード接続性(linear mode connectivity)を見ると分かるからです。これは、異なる2つの基底(basin)に存在する2つの異なるニューラルネットワークを組み合わせることができ、この2つの基底の間には損失障壁が全く存在しない経路が存在するというものです。単純な線形補間を行うだけで、良いモデルが得られるのです。そして、こうしたことは大規模になるほど確認するのが非常に困難になります。
それにもかかわらず、この規模において理論的な証明が存在するという事実は、大規模においてこうした直感がさらに妥当性を持つことを示しています。つまり、大規模なニューラルネットワークのTransformerにおいて、これらの結果が実際にこうしたケースへも外挿できるということについて、非常に説得力のある経験的な裏付けを私たちは持っているということです。
これらすべてを踏まえて、部分空間ネットワーク(subspace networks)についてお話ししたいと思います。そして、ここまでお話ししてきたそれぞれの知見や帰納バイアスが、私たちがインターネット越しにこうした大規模モデルを学習させることを可能にした圧縮アルゴリズムの手法とどのように結びついているのかを、具体的にお話ししていきます。
4. トランスフォーマーの再帰構造と部分空間仮説
4.1 数式的表現と勾配の部分空間へのスナップ現象
Sameera： それでは、私たちが部分空間ネットワークをどのように設計したのか、そして圧縮アルゴリズムの手法について具体的にお話ししていきます。これが基本的なTransformerブロックです。標準的なアテンションブロック、正規化ブロック、線形層、フィードフォワードネットワークから構成されており、皆さんもよくご存知の構造だと思います。ある入力Xを与えたときのこのTransformerブロックの出力を数式で書き下すことができます。
出力をX(l+1)としたとき、この出力を上段の式にきれいに書き下すことができます。そして2段目の式では、層を積み重ねていったときに、非常にきれいな再帰的構造が現れることが分かります。このX0は、3段目の式にある通り、位置埋め込みとトークン埋め込みを足し合わせたものにすぎません。そして次の式でX(l+1)を見ることができます。
ここで、X hat(l+1)という性質について述べることができます。これは、X(l+1)から位置埋め込みとトークン埋め込みを差し引いたものです。そしてこれは完全に、WP1とWP2の張る部分空間に含まれています。つまり、部分空間の線形結合を取ると、その結果として得られるベクトルもまたその部分空間に含まれるということです。ですから、この全体がWP1とWP2の張る空間(span)の中に存在すると言えます。もしこれらが同じ部分空間にあるならば、出力もまた同じ部分空間に含まれることになります。これが、私たちの圧縮手法が依拠している根本的な基盤です。これは非常に重要なアーキテクチャ上のバイアスです。つまり、数式を見て、Transformerがどう機能するかを理解することで、この再帰的構造が非常に強い帰納バイアスを与えてくれるということが分かります。そして、私たちはそれを活用できるかもしれないのです。
もしそれが低ランクであれば、私たちにとって好都合です。そうすれば、勾配を圧縮するための糸口が得られることになります。重み行列が低ランクであることについては少しお話ししましたが、なぜそれが起こるのかについても触れておきたいと思います。その理由は、学習が非常に小さな有効部分空間の中で展開していくからです。勾配についてもそうです。
そこには3つのフェーズがあります。初期化の時点では、勾配はかなりランダムで分散した状態にあります。そしてモデルを学習させ続けると、非常に初期の段階で、勾配は非常に低次元の部分空間へと「スナップ(snap)」します。そしてこの部分空間は、時間とともにゆっくりと滑るように移動していくことがあります。しかし、一度この部分空間にスナップされると、それ以降はかなり安定した状態になります。つまり、勾配が特定の部分空間にある場合、重み減衰(weight decay)によって何が起こるかというと、その部分空間の外にある重みはすべて減衰していきます。それらは勾配信号を全く受け取らなくなるのです。そのため、重みはその部分空間の中に留まるようになっていきます。これが、WP1とWP2が特定の部分空間へと収束していく理由であり、そのためX hat(l+1)もまた同じ部分空間に含まれると言えるのです。
このようなことが起こる理由は、大規模モデルのヘッシアン(Hessian)を見てみると、その曲率のほとんどが最初のいくつかの固有方向に集中しているからです。残りの方向についてはほとんど常にフラットな状態です。
4.2 ヘッシアンの曲率集中仮説に対する留保と部分空間外の勾配の重要性
Sameera： ただし、私はこれが本当にそうなのかについては、もはやそうは思っていません。つまり、これは事実ではあるのですが、それが厳密にフラットであるとまでは、もはや信じていないということです。私が思うに、部分空間の外側から出てくる勾配、つまり部分空間の外側から来るごく小さな情報や勾配は、大規模なLLMの学習において、きめ細かい情報を得るためにやはり重要なのではないかと考えています。この研究に取り組んでいた際には、こうした仮定によってかなり遠くまで到達することができたのですが、これはニューラルネットワークのスペクトルバイアス(spectral bias)、つまり高周波数の情報をより後になってから学習するという性質と何らかの関係があるのではないかと思っています。そして、その高周波数の情報もまた重要なのではないかということです。
とはいえ、この研究の目的のためには、私たちは支配的な部分空間の勾配だけで十分な性能が得られると仮定しました。これはほとんどのケースで当てはまります。しかし、部分空間の外側にあるこの微小な勾配もまた、学習の終盤においては重要なのではないかと私は考えています。そして、これは実際に自分たちで検証することができます。
WP1とWP2、つまり射影行列のランクを測定してみると、それらは時間とともに実際に低下していくことが分かります。そして、位置埋め込みとトークン埋め込みを差し引くと、ランクの劇的な低下が見られます。つまり、ランクの大部分は実は埋め込みの部分から来ているということです。活性化には高ランクであるかのような錯覚がありますが、これは活性化自体が高ランクだからそう見えるだけであって、実際には高ランクの成分の大部分はトークン埋め込みや位置埋め込みから来ているのです。ですから、層自体から来る内容は、実はかなり低ランクなのです。私たちはこの事実を利用して、層自体から出てくる本当に重要な情報を圧縮することができます。
5. 部分空間制約付き圧縮アルゴリズムの詳細
5.1 トークン埋め込みの高ランク・低ランク分解による圧縮設計
Sameera： では、これをどうやって実現するのかについてお話しします。位置埋め込みについては、それが単に位置に依存するものであるため、差し引くことができます。しかし、トークン埋め込みについては事情が異なります。トークン埋め込みは、そのトークンが何であるかに依存するものですが、ある層に着目したとき、その層はトークン埋め込みそのものにアクセスできません。トークン埋め込みへのアクセスは最初の層でしか行われないからです。つまり、他の層と通信することなしに、ある層の中でトークン埋め込みを計算することはできないのです。もしトークン埋め込みを層と層の間で通信しようとすれば、それは通信の圧縮という目的そのものを台無しにしてしまいます。層間でやり取りする情報がかえって増えてしまうことになるからです。
そこで私たちが考えたワークアラウンドは、トークン埋め込みを高ランクの成分と低ランクの成分に分解するというものです。高ランクの成分は固定したまま、低ランクの成分だけを学習します。つまり、いわば差分(デルタ)を学習しているわけです。そして、初期化された高ランクのトークン埋め込みからの差分は低ランクであると仮定しています。まずトークン埋め込みテーブルを初期化し、その初期のトークン埋め込みから学習される差分は低ランクであると仮定するのです。
この2つを足し合わせると、依然として高ランクなトークン埋め込みを得ることができますが、実際に学習する部分は低ランクのままです。これが私たちの採用した手法です。つまり、各層に固定された、学習不可能なトークン埋め込みを保持しておき、学習可能な部分だけを更新していくという方式です。
5.2 射影勾配降下法とロスレス圧縮の証明
Sameera： これを図解したものがこちらになります。私たちはWP1とWP2を、ある特定の部分空間に作用させ、それを部分空間へと制約します。射影勾配降下法を用いることで、これがある種の一次定常点へとほぼ確実に収束することを証明できます。そして、WP1とWP2は、この部分空間によって定義される多様体上を自由に動くことができます。
もう少し明確にお伝えすると、私たちが行っているのは、あるTransformerブロックにおいて、ブロック1のWP1とWP2を、UKによって定義される特定の部分空間へと制約するということです。UKは、私たちが関心を持っている部分空間を張る直交行列です。そして、位置埋め込みと、各層に保持されている固定埋め込みを差し引きます。そうすると、この部分は完全にUKによって定義される部分空間の中に存在することになります。ですから、これはロスレス(情報の損失がない)なのです。ある特定のベクトルの集合がある部分空間に存在している場合、それを同じ部分空間を張る直交行列へ射影すれば、それは完全にロスレスになります。
そして、逆伝播もまたロスレスであることを示すことができます。勾配も圧縮されますが、これもロスレスです。
5.3 グラスマン多様体上での部分空間更新と初期化バイアスの重要性
Sameera： ここで、先ほどお話しした部分空間のゆっくりとした移動について思い出していただきたいのですが、改めて3つのフェーズを振り返っておきます。まず勾配の分布があり、次に勾配がある部分空間へとスナップし、そしてこの部分空間もまた時間とともにゆっくりと移動していきます。これに対応するために、私たちはこの部分空間をグラスマン多様体(Grassmannian manifold)上で最適化します。グラスマン多様体とは、直交行列からなる多様体のことです。つまり、UKや部分空間を表す行列を最適化する際、それはこの多様体上を移動していくということです。これを行うために、私たちはリーマン勾配降下法(Riemannian gradient descent)を適用しています。
ただ、強調しておきたいのは、これは非常に低頻度で行うことができるということです。ですから、通信オーバーヘッドはほとんど発生しません。そして初期化も同様に極めて重要です。私たちが行っているのは、記法にやや不正確な部分があるのですが、これはWOであるべきで、W2はWP1とWP2であるべきです。これらはTransformerの射影行列のことを指しています。WP1とWP2を連結し、それに対してSVD(特異値分解)を行い、上位の支配的な固有ベクトルを用いて部分空間UKを定義します。これによって、モデルの収束がはるかに容易になります。つまり、この初期化バイアスは非常に重要だということです。
さらに私たちは、Adamの修正版も定義しています。詳細には立ち入りませんが、その考え方としては、WP1とWP2を反復的に同じ部分空間へ射影し直す必要がないということです。一度初期化してしまえば、この修正版のAdamが、重みをその特定の部分空間内に留めておくことを保証してくれるのです。
5.4 部分空間選択の任意性と等長写像による説明
Sameera： ここで一つ申し上げておきたいのは、WUKの構成を見ていただくと分かる通り、WP1とWP2を初期化の時点でランダムに初期化し、それを連結してSVDを取っているということです。つまり、これはほぼランダムなのです。ここには事前情報(prior)が全くありません。私たちは解について何も知らないまま、ほぼランダムな部分空間を使っているにすぎません。UKは単なるランダムな部分空間です。私たちはこの部分空間を構成する際、モデルや解について何も知らないのです。
しかし、ここで問題になるのは、ある特定の部分空間が他の部分空間よりも最適であるということがあるのかということです。それは分かりません。ただ、先ほどお話しした損失地形に関する議論に立ち返ると、かなり良い直感を持つことができます。もしモデルを十分にスケールアップすれば、劣最適な局所最小値はほとんど存在しなくなります。これは先ほどお話しした通りです。ですから、パラメータの数をおおよそ一定に保つ限り、どのような最小値に収束したとしても、少なくとも近似的には十分最適であるはずなのです。
そして、異なる部分空間があるとき、例えばこちらの枝が一つの部分空間、こちらの枝がまた別の部分空間だとします。もし2つのモデルがそれぞれ異なる部分空間上で動作していて、それらの層で起きている変換の間に等長写像(isomorphism)が存在するならば、この2つは同じ関数をモデル化することができます。つまり、モデルが十分に表現力を持っていれば、異なる部分空間を用いても同じ関数をモデル化できるのです。ですから、これは単なる座標変換にすぎません。そういう見方をすることができます。つまり、どの部分空間を選ぶかは実際にはそれほど重要ではないということです。パラメータ空間が十分に大きく、十分に過剰パラメータ化されていれば、かなり均一な解へと収束することができるのです。
6. 大規模実験による検証と今後の課題
6.1 Llama 8Bを用いた4大陸分散実験と300人参加ラン
Sameera： ここで、実際に何が起こるのかについて、もう一つグラフィカルな例をお見せしたいと思います。ある層の出力の埋め込みからPEを差し引き、それを圧縮してネットワーク経由で次の層へ送信し、受信側でUKの転置行列を掛け合わせることで復元すると、活性化のロスレスな再構成が得られます。そして、私たちはこの手法によって、従来のスパース化や圧縮、量子化といった手法を包括的に上回る性能を達成しており、スループットにおいて、ほぼ中央集権型の環境に匹敵する形で、約20倍もの向上を得ています。
これを検証するために、私たちはLlama 8Bモデルを学習させました。GPUは4つの大陸にまたがって分散配置されており、それでも私たちは、データセンターに co-located された中央集権型のモデルとほぼ同じ収束を得ることができました。このグラフの横軸は学習時間を表しています。緑色の線が非圧縮のモデルで、これがどれほど遅いかがお分かりいただけると思います。それに対して私たちの手法は非常に高速で、中央集権型モデルと同じ収束を達成しています。
そして、これはAlexも触れていたことですが、これはゼロからの実行(not zero run)ではなく、世界中から300人以上の参加者が集まった実際のランです。16GBのGPUカードさえあれば誰でも参加でき、概念実証(proof of concept)として非常にうまく機能しました。今後、私たちはより優れたランを行っていきたいと考えていますが、今回の実行から多くのことを学ぶことができました。この実行はかなり不安定(jittery)なものでした。そして強調しておきたいのですが、圧縮技術はこの取り組みのごく一部にすぎません。分散学習スタック全体を私たちは自前で書き上げました。耐障害性についても、誰でも自由に参加したり離脱したりできる仕組みが必要で、システムがどれほど不安定になりうるかもお見せできると思います。私たちは異種性(heterogeneity)への対応をはじめ、あらゆる課題に取り組まなければなりませんでした。
ですから、圧縮技術以外にも、多大なエンジニアリング上の努力と科学的な努力がこの実行には注ぎ込まれています。チーム全体に敬意を表したいと思います。そして、今回の実行から多くの教訓を得られたことを願っており、そのほとんどをすでに修正済みです。近いうちに、より良い実行を行えることを期待しています。これが意味するのは、もはやこれは不可能なことではないということです。私たちはこれが実現可能であることを示しましたが、これで問題が完全に解決したというわけでは決してありません。ですから、私はこの問題に取り組んでくれる人がもっと増えることを歓迎したいと思います。というのも、一つ申し上げておきたい点として、モデル並列における圧縮という問題は、まだ解決されていないからです。
6.2 未解決の課題と今後の展望
Sameera： 私が今お話しした圧縮アルゴリズムについても、私たちが学習させたのはChinchilla最適、つまりパラメータ数の20倍のトークン数までのモデルです。例えば10億パラメータのモデルであれば、200億トークンで学習するということです。そして、そこまではうまく機能したのですが、現代のネットワークにとってはChinchilla最適というのはかなり不十分な基準になっています。現在では、モデルはそれよりもはるかに長く学習されます。例えば、ほぼ10億パラメータのモデルが、1兆トークン規模で学習されるようなこともあります。そのように長く学習させると、ギャップが現れ始めることが分かりました。
つまり、この圧縮手法を用いて長期間学習させると、収束が少し遅くなってしまうのです。私たちは現在この点を改善するために取り組んでいます。一つのアイデアとして提案したいのは、固定された部分空間の代わりに、部分空間の混合(mixture of subspaces)という方向で考えることができるのではないかということです。つまり、固定された部分空間ではなく、変化していく、動的な部分空間の混合を考えるということです。
私は皆さんに、この問題の改善に取り組んでいただくことを推奨したいと思います。というのも、誰がこの問題を解決するかは私たちにとって重要ではなく、ただこの問題が解決されることを望んでいるからです。それは世界にとって良いことだからです。以上が私からのお話になります。何かご質問があれば、喜んでお答えします。
7. 質疑応答
7.1 動画符号化とのアナロジー、ファインチューニングへの応用、視覚と言語の違い
質問者： 皆さんがネットワーク越しにデータを送信する方法について、動画符号化(video encoding)の仕組みとの類似性を指摘させてください。動画符号化では、キーフレームがあり、その後は差分だけを送信しますが、数秒ごとに新たなキーフレームを送る必要があります。そこで、皆さんの場合、実際にデータの損失が生じて、エンコーディング全体を再送しなければならなかったような経験があったのかどうかを理解したいと思っています。
Sameera： 今回のケースでは、そういったことはありませんでした。理論上はこれはロスレスですし、実際にもかなり遠くまでうまくいきます。ただ、完全にロスレスというわけではないとも言っておきたいのですが、それは実際にはネットワークの容量を縮小しているからです。私たちは特定の部分空間の中で動作しているため、そこにはある種の容量的な損失が存在します。ですから、もし時折フルの情報を送ることができれば、それはネットワークがより良い解へと収束するための有用な情報を提供してくれるだろうという点には同意します。実は現在、私たちはファインチューニング手法に取り組んでいます。そのために、ファインチューニングの領域、つまりモデルの適応(model adaptation)の領域では、スパースな信号を送りながらも、時折フルの勾配も送ってモデルを整合させるというアプローチを採用しています。
そして、一つ違いがあると思うのは、動画や視覚情報は言語よりもはるかに冗長性が高いということです。というのも、言語は人間が構成したものだからです。私たちは言語をより効率的にしたいと考えていますが、単語は視覚トークンよりもはるかに意味的に密な情報を含んでいます。視覚トークンははるかに冗長性が高く、スパースな情報でも許容できますが、言語の場合はより難しいと私は考えています。私たちが扱っている領域は視覚情報よりも冗長性がはるかに低いため、より慎重に考える必要があるのです。
7.2 モデルアーキテクチャ、Mixture of Expertsへの適用可能性
質問者： 7.5Bモデルはどのようなアーキテクチャなのか気になっているのですが。
Sameera： これは主に、修正版のLlamaです。ただ、少し変更を加える必要がありました。というのも、大規模モデルを学習させた経験がある方ならご存知かもしれませんが、突然の損失スパイク(loss spike)が発生する問題があったからです。これは小規模では見られない現象です。だからこそ、大規模でこれらを検証することが非常に重要なのです。モデルをスケールアップすると、ほとんど相転移のようなことが起こります。小規模では観測されないことが、大規模では起こり始めるのです。私たちにはいくつかの不安定性の問題があり、これは大規模学習ではよくある課題です。そのため、元のLlamaアーキテクチャに、QKノルムやノルムの並べ替えなどを加える形で、少し修正を施す必要がありました。また、圧縮アルゴリズムと互換性を持たせる必要もありました。ですから、修正版のLlamaではありますが、基本的にはLlamaだと言えます。
質問者： なるほど、Mixture of Expertsで試したことはありますか。
Sameera： はい。社内で試しました。通常のMixture of Expertsだけでなく、異種混在のMixture of Experts(heterogeneous mixture of experts)でも機能します。つまり、異なるGPUに分散配置された、容量の異なるフィードフォワードネットワークやエキスパートを持つことができるということです。これも機能します。同じ圧縮を適用できるのは、WP1がアテンション層から出てくるものだからです。つまり、その活性化も圧縮して別のGPUへ送ることができます。ですから、Transformerブロック全体を1つのGPU、1つのノードの中に収める必要はなく、さらに細かい粒度で分割することも可能です。例えば、アテンションを1つのGPUに置き、エキスパートを別のGPUに置くといった形です。それらのGPU間の情報も圧縮することができ、それがうまく機能します。
7.3 圧縮率とデータ特性の関係、LoRAの限界に関する見解
質問者： 素晴らしい発表をありがとうございます。どれだけ圧縮できるかという話ですが、低ランク構造があるとおっしゃいましたよね。それは、実際に圧縮しようとしているデータの種類にも依存するのでしょうか。データそのものが圧縮されているということはあるのでしょうか。圧縮率と、どのようなデータを、どれだけの量送信しようとしているかという文脈の間には、何らかの関連があるべきだと感じています。
Sameera： はい、もし私の理解が正しければ、圧縮と私たちが使っているデータとの間に何らかの相互作用があるかどうかを尋ねているということですね。現時点では、そうした関連は観測していません。しかし経験的には、数学は少なくともモデル適応のフェーズ、つまりポストトレーニングのフェーズにおいて、より高いランクの更新を必要とするのではないかと考えています。これが、LoRAのような手法が数学的な推論においてあまりうまく機能しない理由だと思います。要約のような平易なタスクではうまく機能するのですが、本格的なタスクではそうはいきません。LoRAはもう時代遅れだと思っています。これは、先ほどお話しした内容とも関連していて、私たちの手法は基本的に低ランクの更新を与えるものです。
もし推論やポストトレーニングを行おうとするなら、そこは私たちがまだ検討していく必要がある部分です。現在のところ、私たちは事前学習にしか取り組んでおらず、事前学習の手法についてはそれで十分機能しています。しかし、分散型のモデル適応、分散型のファインチューニング、分散型の強化学習といった、より本格的な取り組みを行おうとするならば、おそらくより優れた圧縮アルゴリズムを考案する必要が出てくるでしょう。ですから、これはまだ解決には程遠い問題です。取り組むべき側面はたくさんあると思っています。
7.4 等長写像と損失地形の平坦性、初期化と圧縮性能の関係
質問者： 素晴らしい発表をありがとうございました。本当に見事な解決策だと思います。等長写像(isomorphism)についておっしゃっていた部分について伺いたいのですが、仮に同じような部分空間へと収束したとして、その部分空間の周辺の損失地形は似ているのでしょうか。
Sameera： すみません、最後の部分をもう一度お願いできますか。
質問者： 部分空間の周辺の損失地形は似ているのでしょうか。
Sameera： 損失地形についてですね。部分空間の周辺ということですね。分かりました、幾何学的な分析はまだ行っていません。性能面では類似しています。ただ、より幾何学的な直感は必要だと思いますが、理論的には、それらは何らかの線形モード接続性(linear mode connectivity)で繋がっているはずだということは分かっています。ただ、確実なことは分かりません。
質問者： 盆地(basin)の平坦性について気になっていたのですが。
Sameera： はい。
質問者： というのも、もし周辺の領域がより平坦であれば、モデルはより頑健である可能性が高くなると思うのですが。
Sameera： その通りです。
質問者： 非常に鋭い(sharp)場合と比べて、ということです。
Sameera： まさにその通りです。それは将来のモデル適応の部分において非常に有用になるはずです。私の直感をお伝えしますね。確実には分からないのですが、かなり良い直感的な答えはお伝えできます。私は、それはより平坦だと思っています。というのも、このような形でモデルを初期化すると、モデルがずっと滑らかになるからです。つまり、通常、もしモデルがフーリエ変換を取ったときに高ランクの情報を符号化するようなものであれば、収束する解はより鋭い(sharper)ものになる傾向があります。これは理論的に示すことができます。ある意味で、私たちはモデルをずっと滑らかにしているのです。ですから、損失盆地はより平坦になっていると思います。ただ、それは必ずしも性能が良いということを意味するわけではありません。それはただ汎化性能が良いということであって、性能面で必要なすべてのモードをカバーできているとは限りません。数学的なタスクなど、一部のタスクではかえって性能が悪くなるかもしれません。しかし、フーリエ変換を取ったときに、より低周波数に集中しているはずですから、このように学習させたモデルの損失盆地ははるかに平坦であるはずだと私は信じています。
質問者： 初期化は圧縮の性能とも密接に関係していると思いますか。
Sameera： はい。
質問者： それから、より良い圧縮のために初期化を賢く選ぶ方法はあるのでしょうか。
Sameera： 方法はあると思います。先ほどお話ししたスライドの通り、理論的にはどの部分空間を使うかは、固定された部分空間である限り重要ではありません。ただ、もちろん動的な部分空間を使うのであれば、より良い解、より良い盆地へと収束できるようになると思います。
質問者： ありがとうございました。
司会： 他にありますか。それでは、ありがとうございました、Samiraさん。
Sameera： ありがとうございました。