推論実行と学習の間で計算資源を最適に配分する

はじめに

Sam Altman は最近、OpenAI が現在1日あたり約1,000億トークン、年間でおよそ36兆トークンを生成していると述べた。現代の言語モデルは10兆トークン程度のデータで学習しており、学習中に処理するトークンは推論実行中に処理・生成するトークンより約3倍コストが高い¹。これを踏まえると、単純な分析²からは、OpenAI の年間の推論実行コストは年間のモデル学習コストと同程度の規模だと考えられる。

これは一見、奇妙な偶然に思える。なぜ一方が他方を完全に圧倒しないのだろうか。しかし、この2つの量が同程度の桁になると考えるべき十分な理由がある。学習と推論実行の計算資源のトレードオフである。本稿では、このトレードオフがどのようなものかを示し、それについての現在の実証的証拠を踏まえると、学習と推論実行に費やされる計算資源がなぜほぼ同程度になると考えるべきなのかを簡単に説明する。

トレードオフ

一般に、モデルの性能を上げる方法は2つある。モデルの学習により多くの計算資源を費やすか、モデルで推論実行を行う際により多くの計算資源を費やすかである。推論実行時に計算資源を多く費やしてモデルの性能を上げる手法はさまざまある。大規模言語モデルの例としては、より大きなデータセットでの過剰学習（overtraining）、思考の連鎖（chain-of-thought）プロンプト、ウェブ閲覧、反復サンプリングなどがある。モンテカルロ木探索（MCTS）に基づくゲームプレイ用のモデルでは、実行時により多くのプレイアウトを行うことも、性能を上げる重要な手法である。これらを組み合わせると、小さなモデルが推論実行時の計算資源をより少なく使って大きなモデルの性能に並ぶことができる。

Villalobos and Atkinson (2023) は、推論実行時に性能を上げるために使えるさまざまな手法を、次のトレードオフ表にまとめている。

手法	学習側のトレードオフの最大範囲	推論実行側のトレードオフの最大範囲	規模を大きくしたときの影響	現在のモデルの使われ方
スケーリング方針の変更	1.2 OOM	0.7 OOM	なし	学習計算量を最小化
MCTS	1.6 OOM	1.6 OOM	モデルが完璧な性能に近づくと、範囲が急に0に近づく	混在
枝刈り（pruning）	~1 OOM	~1 OOM	範囲がわずかに広がり、学習に対する見返りが逓増する	学習計算量を最小化
反復サンプリング（検証が安価な場合）	3 OOM	6 OOM	範囲が広がる	推論実行の計算資源を最小化
反復サンプリング（検証が高価な場合）	1 OOM	1.45 OOM	なし	推論実行の計算資源を最小化
思考の連鎖	NA	NA	不明	推論実行の計算資源を最小化

たとえば、スケーリング方針の変更とは、Hoffmann et al. (2022) が示した計算最適なデータセットのスケーリング則よりも学習データセットを大きくして、モデルを過剰に学習させることを指す。これにより、モデルの品質を損なわずにモデルのサイズを小さくする余地が生まれる。短いコンテキスト長では1トークンを処理する計算コストはモデルのサイズに比例するので、モデルの推論実行が安くなる。

上の表のトレードオフをすべて検討したうえで、Villalobos and Atkinson (2023) は次のように結論づけている。

手法や分野によって大きなばらつきがあるので、一般的な場合にどれだけの計算資源をトレードオフできるかについて正確な数字は示せない。しかし経験則として、それぞれの手法によって、学習か推論実行のどちらかで約1桁（OOM）の計算資源を節約でき、その代わりにもう一方を1桁よりやや多く増やすことになると予想される。

簡単のため、「やや多く」は無視して中心的な主張だけに注目する。つまり、おおよそ、学習計算量を1桁増やして推論実行を1桁減らしても（あるいはその逆でも）、できあがるモデルの性能は変わらない。このトレードオフはいずれ成り立たなくなるが、現在使われている学習と推論実行の計算資源の水準の周辺では、幅広いタスクで成り立つようである。さらに、トレードオフ表のさまざまな手法は多くの場合重ねて使えるので、学習と推論実行の計算資源のトレードオフが可能な範囲を広げられる。

投資が同程度になると予想する理由

なぜこれが、学習と推論実行への支出が同程度の水準になると予想すべきことを意味するのか。それを見るために、ある研究所が学習に1ゼタFLOP、推論実行に100ゼタFLOPの計算資源を費やしているとしよう。この場合、学習と推論実行の計算資源のトレードオフを利用して、学習計算量を10倍に増やし、推論実行の計算資源を10分の1に減らしても、ユーザーから見たモデルの品質は変わらない。この新しい推論実行の構成でもユーザーには同じ性能が提供されるが、研究所が費やす計算資源の総量は $100+1=101\text{zettaFLOP}$ ではなく $10+10=20\text{zettaFLOP}$ になり、総コストは5分の1に減る。この単純な例の論理を一般化すると、ある性能水準を達成するための計算コストを最小化する構成では、一般に学習と推論実行に同じ量を費やすことになる。

この結果はトレードオフの大きさに異常に敏感なのではないか、という懸念が考えられる。学習計算量を1桁（10倍）増やせば推論実行の計算資源を1桁節約できる（つまり推論実行に費やす計算資源を10分の1にできる）のではなく、学習計算量を0.5桁増やすだけで推論実行の計算資源を1桁節約できるというバランスだったらどうだろうか。その場合、学習と推論実行への計算資源の最適な配分はどうなるだろうか。

実は、学習計算量 $\beta$ 桁を推論実行の計算資源 $\alpha$ 桁とトレードオフできるとき、正しい戦略は、割合 $\alpha /(\alpha +\beta )$ を学習に、割合 $\beta /(\alpha +\beta )$ を推論実行に費やすことである。³ その形式的な証明はこちらに示すが、本文が複雑にならないよう、ここでは具体例に絞る。学習計算量1桁（10倍）を推論実行の計算資源2桁（100倍）とトレードオフでき（つまり $\beta =1,\alpha =2$）、計算資源が合計300ゼタFLOPあるとしよう。これを単純に学習と推論実行に等分して、それぞれに150ゼタFLOPを割り当てると、最適な行動にはならない。学習計算量を0.1桁増やし、推論実行の計算資源を0.2桁減らせば、性能を変えずに計算資源の総使用量を $150\cdot 10^{0.1}+150\cdot 10^{-0.2}\approx 283.4\text{zettaFLOP}$ にでき、計算資源の予算300ゼタFLOPより少なくて済むからである。つまり、性能を落とさずに計算資源の支出を減らせたことになり、最初の配分が最適でなかったことがわかる。

しかし、学習に $(2/3)\cdot 300=200\text{zettaFLOP}$、推論実行に $(1/3)\cdot 300=100\text{zettaFLOP}$ という最適な配分にしている場合、先ほどと同じトレードオフを行うと、計算資源の総量は $300$ から $200\cdot 10^{0.1}+100\cdot 10^{-0.2}\approx 314.9\text{zettaFLOP}$ に変わり、最初より悪くなる。実際、学習に200ゼタFLOP、推論実行に100ゼタFLOPという配分から出発すると、配分をどう変えても、性能を保ったまま計算資源を節約することはできない。だからこの配分が効率的なのである。

この結果から得られる主な教訓は、学習計算量と推論実行の計算資源の大きさがほぼ同程度にならないためには、$\alpha$ が1からわずかに異なるだけでは不十分だということである。$\alpha$ が1から少しずれても影響は小さい。2つの計算資源の配分を大きく違うものにするには、$\alpha$ が1から大きく離れている必要がある。具体的には、学習と推論実行の計算資源を異なる桁にしたいなら、その値は0.1未満か10より大きくなければならない。トレードオフ表にまとめた実証的証拠は、1からそれほど離れた $\alpha$ の値を支持していない。これが、推論実行と学習に費やされる計算資源がほぼ同程度になると予想する主な理由である。

結論

学習と推論実行のトレードオフが学習計算量と推論実行の計算資源の十分に広い範囲で成り立つなら、大規模なモデルの学習と、そのモデルによる推論実行には、同程度の計算資源が費やされると予想すべきである。言い換えれば、将来、これらの支出の一方が他方を1桁以上も上回ると予想すべきではない。この結果は、トレードオフの大きさ、つまり推論実行のコストを1桁減らすのに学習計算量を何桁余分に払う必要があるかについての、ありうる程度の不確実性に対しても頑健であるようだ。

この結論が誤りとなる最もありそうな道筋は、既知のトレードオフが「息切れ」することである。たとえば、知られている推論実行を節約するトレードオフから得られる利得をすべて使い果たした世界では、計算資源の大半を推論実行に費やしているかもしれない。すべてのトレードオフを推論実行に最大限有利になるよう調整してもなお推論実行が計算コストの大半を占めるなら、推論実行のコストを節約するために私たちにできることは、もうあまり残っていないのかもしれない。一方で、現在知られていない、学習計算量と推論実行の計算資源をトレードオフする別の方法があるかもしれず、そうした方法が見つかれば、トレードオフが可能な範囲は広がるかもしれない。この面で事態がどう展開するかははっきりしない。

ここで考慮していない、より重要度の低い要因もいくつかある。たとえば、推論実行は学習よりもメモリ帯域幅の制約が厳しいかもしれない。また、ムーアの法則と時間割引のために、将来費やす計算資源は先に費やす計算資源より安くなるので、推論実行に費やされる計算資源の総量のほうが多くなるかもしれない。さらに、このモデルでは、AI開発者が特定の性能水準のモデルにどれだけの需要があるかを完全に見通せると仮定したが、実際にはそうではない。これらの効果は重要だが、ここで扱っても分析が複雑になるだけで、学習と推論実行への計算資源の支出は同じくらいの規模になるはずだという基本的な定性的結果は変わらない。今後、本稿の論理を土台に、より現実的なモデルを構築するつもりである。