- LLMの標準目標である 次トークン予測 を、複数の将来トークンを同時に予測する形に置き換えることで、同じデータ量と計算予算でもコード・自然言語生成性能をさらに引き上げられる
- 構造は共有 Transformer本体 の上に複数の出力ヘッドを載せる方式で、基本推論では次トークン用ヘッドだけを使うため、従来の自己回帰生成と同様に動作する
- コードモデルでは 13Bパラメータ モデルが、比較可能な次トークンモデルより HumanEval を12%、MBPP を17%多く解けており、この利得は大規模モデルほど明確だった
- 追加ヘッドは 自己推測デコーディング に活用でき、4-token prediction モデルは最大3×、8-byte prediction モデルは 6.4× の推論速度向上を示した
- 合成課題では induction heads やアルゴリズム推論に有利であり、学習時の teacher forcing と生成時の自己回帰分布の差を縮める効果がある可能性がある
マルチトークン予測方式
- 従来の言語モデリングは各位置で次の1トークンに対する 交差エントロピー損失 を最小化する
- マルチトークン予測は、各位置で次の n 個のトークンをまとめて予測するよう学習目標を拡張する
- モデル構造は3つの部分に分かれる
- 共有 Transformer本体 が観測されたコンテキストの潜在表現を作る
- n 個の独立した出力ヘッドが各将来トークンを並列に予測する
- 共有 unembedding matrix が最終的なトークン確率を計算する
- 最も単純な推論方式は、次トークン予測ヘッドだけを使う通常の 自己回帰予測 で、残りのヘッドは捨ててもよい
- 追加の出力ヘッドは blockwise parallel decoding や Medusa-like tree attention のような 自己推測デコーディング(self-speculative decoding) に利用できる
メモリ効率のよい実装
- 単純な実装では各ヘッドの logit と gradient をすべてメモリに載せる必要があり、GPUメモリ使用量が大きくなる
- 現在のLLMでは vocabulary サイズ V が潜在表現次元 d よりはるかに大きく、logit vector がGPUメモリのボトルネックになる
- 提案手法では共有本体の forward pass 後に、各出力ヘッドの forward/backward を順次実行する
- 1つのヘッドの logit と gradient は次のヘッドへ進む前に解放される
- 本体側には累積 gradient だけを保持する
- この方式は peak GPUメモリ使用量を O(nV + d) から O(V + d) に削減し、ランタイムコストは増やさない
コードモデルの実験結果
- 実データ実験では、次トークン予測モデルと n-token prediction モデルを同じパラメータ数で比較した
- 将来予測ヘッドに n−1 層を追加する場合、共有本体から n−1 層を削除する
- 300M から 13B まで6つのサイズのモデルを、少なくとも 91B code tokens でゼロから学習した
- MBPP と HumanEval の評価では、小規模モデルがベースラインより劣る場合もあったが、スケールが大きくなるほどマルチトークン予測が優位になった
- 13B モデルは比較可能な次トークンモデルより多くの問題を解いた
- HumanEval で 12% 多くの問題を解決
- MBPP で 17% 多くの問題を解決
- 7B モデルを 200B code tokens で学習したアブレーションでは n=1, 2, 4, 6, 8 を比較した
- n=4 が HumanEval と MBPP の pass@1、pass@10、pass@100 で一貫して最良だった
- APPS/Intro では n=6 が優勢だった
- 最適な window size は入力データ分布に応じて変わりうる
推論速度と byte-level モデル
- 7B 4-token prediction モデルに greedy self-speculative decoding を適用し、学習に使っていないコード・自然言語テストプロンプトでデコーディング速度を測定した
- 結果はコードで 3.0×、テキストで 2.7× の速度向上を示した
- コードでは3つの提案のうち平均2.5トークンが受理されたトークンだった
- 8-byte prediction モデルは推論速度で 6.4× の向上を記録した
- byte-level tokenization 実験では、7B byte-level transformer を 314B bytes、約 116B tokens に相当するデータで学習した
- 8-byte prediction モデルは next-byte prediction と比べて、より多くの問題を解いた
- MBPP pass@1 で 67% 多くの問題を解決
- HumanEval pass@1 で 20% 多くの問題を解決
- multi-byte prediction は byte-level モデルをより効率的に学習する道筋になりうる
複数 epoch、微調整、自然言語での結果
- 同じデータで複数 epoch 学習しても、マルチトークン予測は次トークン予測に対して一部の優位性を維持した
- MBPP pass@1 は +2.4%
- HumanEval pass@100 は +3.2%
- それ以外の指標は類似していた
- CodeContests の微調整では、4-token prediction で事前学習した 7B モデルが次トークン基準モデルより pass@k 全般で優れていた
- 4-token prediction モデルをそのまま n′=4 loss で微調整した場合も基準モデルより良かった
- 追加ヘッドを削除して next-token target で微調整した場合が全体として最も良かった
- 自然言語では 7B モデルを 200B tokens で学習し、6つの標準 NLP benchmark を評価した
- 2-token prediction モデルは次トークン基準モデルとほぼ同等だった
- 4-token prediction モデルは性能がやや低下した
- より大きなモデルサイズが必要な可能性がある
- 生成型自然言語評価は要約と数学課題に分けて行われた
- 8つの summarization benchmark では、n=2 と n=4 モデルは 200B・500B tokens 学習の両方で ROUGE-L F1 基準において次トークン基準モデルを上回った
- GSM8K 8-shot 評価では、200B tokens 時点では n=2 が基準モデルを上回ったが、500B tokens 以降では傾向が逆転し、n=4 は全体的により悪かった
合成課題で見た induction とアルゴリズム推論
- induction とは、文中で “AB” が現れた後、後で “A” が再び現れたときに続けて “B” を予測するパターンを指す
- children stories データセットで 1M〜1B nonembedding parameters のモデルを学習し、ランダムな2トークン名を入れたテストセットで induction capability を測定した
- 30M 以下の小規模モデルでは 2-token prediction loss が induction capability の形成を大きく改善した
- 多項式算術課題では、F7[X]/(X5) における unary negation、addition、multiplication、composition を含む式を学習・評価した
- マルチトークン予測は課題難易度全般で精度を高め、out-of-domain generalization も絶対値は低いものの大きく改善した
- 30M から 100M にモデルを大きくするより、next-token prediction をマルチトークン予測に切り替える効果のほうが大きかった
なぜ機能するのか
- マルチトークン予測は teacher forcing 学習と、inference-time の自己回帰生成の間にある分布不一致を緩和できる可能性がある
- 次トークン予測は短い範囲の予測に集中する一方で、長距離依存を無視しうる
- マルチトークン予測は、後続トークンと強く関連するトークンにより大きな暗黙的重みを与える
- これは choice point の強化として解釈できる
- 有用なテキスト生成は choice point で正しい判断を選べるかどうかに左右されると考えられる
- 情報理論的な展開では、2-token prediction は X と Y の間の mutual information 項の重要度を、next-token prediction よりも強める形で現れる
限界とコスト
- 今後の課題として、マルチトークン予測で n を自動選択する方法、loss scale や loss balancing の活用、vocabulary size の調整、embedding space で動作する補助 prediction loss の開発が残されている
- すべての実験モデル学習には合計で約 500K GPU hours が使われた
- ハードウェアは A100-80GB と H100
- 推定総排出量は約 50 tCO2eq で、Meta の sustainability program により 100% オフセットされた
- 目的は言語モデルの compute と data efficiency を高めることだが、rebound effects には注意が必要であり、LLM の社会的利点とリスクの両方を考慮すべきである
1件のコメント
Hacker Newsの意見
この分野ではあまりにも多くのことが起きている
データ、事前学習、学習、推論、専門家混合、RAG といった用語が実際の流れの中でいつ出てくるのかを、工場見学のように時系列で説明してくれる資料があるとよい
たいていは、大きな全体像の中で用語がどこに入るのか分からないし、最初に事前学習を見たときは、学習の前にデータを処理する工程だと思っていたが、実際には別の学習だった
この分野は信号対雑音比が非常に低く、Langchainのような業界の先導者のドキュメントでさえ、すでに古くなっていたり、互いに矛盾していたりすることがある
ブロックチェーン過熱期にも似たようなことがあったので、過熱列車の特徴なのだと思う
特に RAG のように最近の進歩が速すぎて、しばらくは最新性を保てる参考書が出にくいため、どこから始めればよいのか分かりにくい
それでもLlamaIndexのような上位レベルのツールのドキュメントは、概念そのものを深く説明するというより、全体像の中でどこに入るのかを示してくれるので、出発点としては悪くない
YouTubeはいつものように、最新トレンドでクリックを集めようとする非専門家が海のように多く、出発点としてはあまり良くないと思う
もちろん何事もそうではあるが、それを応急処置で直そうとするより、近い将来に GPT-4よりはるかに優れたもの が出てくると仮定し、その前提の上で差別化されたプロダクトを設計するほうがよい
投機的デコーディング を知っている人にとっては、これは基本的に自己投機的デコーディング
予測したラベル列を依然として自己回帰的にネットワークへ再投入し、一致する地点までだけ予測を維持する
そのため性能は悪化せず、速くなるだけで、ここでは最大3倍だが、投機的デコーディングでは普通の水準
マルチタスク学習のおかげでさらに良くなる可能性もある。数ステップ先のターゲットを補助損失として予測するというアイデアは、すでにかなり古くからあり、良い取り組みだ
投機的デコーディングが計画能力まで改善するなら、より採用されやすくなるだろう
投機的デコーディングは、出力の正確さや品質という意味でモデル性能を下げない
数学的には、通常の自己回帰デコーディングをした場合とサンプリングされる変更後の分布が同一で、差が出る理由は単なるランダム性によるもの
もし性能を速度という意味で使っているなら、投機的デコーディングが速度を落とすこともあり得るが、ほとんどの入力と適切なドラフトモデルの選択では、そうならないはず
LLMはシーケンス予測で、特定の出力長までの可能なすべてのトークン組み合わせの 確率分布 を考慮していないのか？ すでにそうしていると思っていた
そうでないなら、今のようにうまく動いているのは驚きだ
たとえば2ビットシーケンスの可能性と確率が 00: p=0.36、01: p=0.04、10: p=0.30、11: p=0.30 なら、最も可能性の高い2ビットシーケンスは00
しかし次のトークンを1つだけ予測すると、0: p=0.40、1: p=0.60 なので、次のビットは1のように見え、その次のビットを予測するときに最適でない出発点につながる
長いシーケンスでは、結合確率分布が周辺分布へうまく分解できないほど、誤りはより目立つはず
さらに考えてみると、テキスト学習データでn番目の未来トークンだけを考慮するようにクロスエントロピー損失関数を変え、LLM性能とnの関係をプロットする簡単な研究ができそうだ
現在のLLMがすべてn=1だと仮定すれば、次の1個からn個のトークンまでの結合確率分布を予測するのに必要なリソースの爆発を、大部分回避できるという仮説だ
n番目のトークンを直接予測するには、より良いデータモデルが暗黙的に必要になるからだ。少なくとも人間が作ったテキストではそうで、すべてのデータ型に当てはまるわけではない
LLMは学習分布に従うテキストを サンプリング するように設計されていて、後続の「最も可能性の高い」テキストを教えるように設計されているわけではなく、実際に私たちもそれを望んでいない
そうすると出力の多様性がなくなる
例では、チャットアプリケーションなら40%は0、60%は1をサンプリングするのが理にかなっている
最も可能性の高い文が重要な質問応答のような用途では、他の人が言っているようにビーム探索が役立つ
また、モデルが「先を見越して」未来のトークンを事前に計算し、それを現在のトークン予測に使える点も考慮すべき
実際に [1] のような研究がこれを扱っている
最後に、一度に1トークンを予測するのは人間が話すときにしている方法なので、間違ったアプローチではない。私たちは話す前に頭の中でこうした「先読み」をしている
[1] https://arxiv.org/abs/2404.00859
記憶では、LLMの出力に奇妙なパターンが見られ、「an」が「a」より可能性が低いことが多いため、予想より 母音で始まる名詞 が少なく出る、といったものだ
つまり、ある分布で言語モデルを学習させ、温度1 でサンプリングすれば、まったく同じ分布が得られる
低い温度や貪欲にサンプリングすれば、当然ながら別の分布になる
ビーム探索のような複数の局所最適化サンプリング戦略があり、投機的デコーディングのような、より大域的なサンプリング作業も多くあった
現在のLLMは、出力トークンごとに最初から始めている、という理解で合っているのか？
「バナナを黄色くするものは？」と尋ねて、「Bananas are yellow due to a pigment called bromelain.」と答えるなら、「a」を出力する時点で、すでにニューラルネットワーク内では pigment と bromelain の概念がある程度活性化していそうに思える。
その時点でもう気が変わって「an optical illusion...」のような答えに続けることはできないので、bromelain という色素について話すと、すでに先を計画していたように見える。
LLMが「a」を出力するときにすでに行った作業を、次の回答に活用できるのだろうか？ ニューラルネットワークの状態を次の回答のために保存できるのだろうか？
「Bananas are yellow due to a」と「Bananas are yellow due to an」
前者なら「Bananas are yellow due to a pigment called bromelain.」と答えられ、後者なら「Bananas are yellow due to an organic compound called bromelain, which is a yellow pigment.」と答えられる。
どちらにしても、GPTが「a」や「an」を選んでも、応答の意味には影響しない場合がある。
極端に言えば、LLMは「due to」の次のトークンは「a」が55%、 「an」が45%という間抜けなヒューリスティックで動いている、と見ることもできる。
現実はもちろんもっと精緻だが、そのようなヒューリスティックだけでもこの挙動は説明できる。
事前学習データに bromelain に関する事実を入れていなければ、LLMは本当に「an optical illusion」のような内容で自動補完するかもしれない。
GPT-3はその種の事実誤りをかなり頻繁に起こしていたが、「a」と「an」の文法規則は見つけ出していたと記憶している。
述べたように、概念が実際に前もって活性化している必要はないと思うが、暗黙的・創発的な意味では、そのような 事前活性化 が起きている可能性はある。
KVキャッシュは、ここまでのテキストとそれに対するモデルの考えを表現したものと見ることができる。
言語モデルはシーケンス全体を最後まで学習するので、こうしたことが起きる可能性は高いと思う。
マルチトークン予測はこの挙動を明示的に促すが、定義された小さな n トークンの窓の中だけでそうする。
一方で、Transformer言語モデルの計算活用度を高めようとする取り組みとして、早期終了、深さ混合、SSMのような新しい構造も数多くある。
おそらく「a」や他の候補の後に何が来うるかについての概念的な計画はあるだろうが、そのような高水準の予測は「a」が生成されると、また最初から見直される。
モデルは各単語が生成された後に気を変えられるだけでなく、そうしなければならない。
そのため、この種の「計画」は非常にはかなく、深く考えて回答や表現を選ぶ人というより、即興で作り上げる フリースタイルラッパー に近い。
中核となるLLMはトークンを受け取り、次に来る順位付けされたトークン集合と「確信度」を出力する。
その後、通常はフィルタリングと探索の段階があり、その順位付きトークンを再びLLMに入れて、さらに多くの順位付きトークンを得て、短い確率木を作る。
たとえば上位N個のトークンを再入力すると、それぞれが新しい上位N個のトークン集合を作る。
その木を見て、合算した確信度が最も高い枝、反復トークンが最も少ない枝、入力トークンと一致するトークンが最も少ない枝を選ぶ、といった基本的なフィルタリングを行ったり、たいていはこれらの基準を組み合わせ、合算した確信度で重み付けしたランダム選択まで加えたりする。
そのため、重みが完全に固定されたLLMに同じ入力を何度与えても、異なる出力が出ることがある。
つまり具体的な質問に答えるなら、モデルは「気を変える」ことができる。生成されたすべてのトークンが、確率的な出力フィルターに、可能な出力経路の中から新しい経路を選ぶ機会を作る。
とても素朴な質問かもしれないが、文全体の意味をエンコードするベクトルを作れると仮定した場合、なぜLLMに1語ではなくその文ベクトルを予測するよう学習させられないのだろうか？
言語用のオートエンコーダを学習させること自体は、テキストに含まれる情報量が視覚/映像に比べて少ないため、実際にはかなり簡単です
難しいのは、すべての信号がトークン空間での厳密一致から来るときに、モデルを意味の部分へ集中させることです
そこでYann LeCunの共同埋め込み予測アーキテクチャというアイデアが出てきます
また、補助タスクはより多くの信号を与える一方で、焦点をずらしてしまうというトレードオフが常にあります
私たちの場合、予測するトークン数が多すぎると性能低下が見られました
したがって、潜在予測の手法では、何が有用なのかを整理する必要があります
問題は、答えを表すベクトルを得たあと、その答えを再び単語表現に戻す別のモデルのようなものが必要になる点だ
テキスト用の拡散モデルのようなものかもしれない
さらに、この拡散モデルが近似する関数は単射ではなく、うまくいっても全射であり、悪ければ数学的な意味での関数ですらない可能性がある
1つの埋め込みに対して可能なテキスト表現は多数あり、その大半は文法的にも意味的にも有効ではない可能性があるからだ
最後に、埋め込みは何らかのデータの損失のある表現なので、逆関数は多くのニュアンスや文脈を失う
LLMは、クエリと直前のn個のトークンとの自己一貫性を保つ形で次のトークン、今では次のn個のトークンを予測することで上記の問題を避けており、LLMが近似する関数はおおむね全射に近いはずだ
今は大ざっぱに言えば単語を扱っているので、AIは知っている単語しか使えないが、単語から新しい文を合成できる
AIが文単位で動くなら、すでに見た文を繰り返すだけではないのか？ そうなると新しい段落は合成できても、新しい文は作れない気がする
英語では、文がAIにとって有用な抽象化なのか確信がない。人間にとってもかろうじて有用という程度だ
普通のチャット、メール、YouTubeコメントを見ると、実際には文ではなかったり、句読点すら使っていなかったりすることが非常に多い
文が意味の単位に対応しているとは思わない
文は書き手によって2語の場合もあれば、英語論文の半分に及ぶ場合もあり、6つのアイデアをまたぐことも、1つだけを含むこともある
文がどこで終わるかは、概して意味よりも書き手の傾向に左右される
文をトークンに分割すると、各トークンがベクトル表現を得る
文レベルにすると、全トークンの語彙は無限になるはずだ
それをより高速なモデルでやることはできるかもしれないが、出力品質は落ちそうだ
論文をまだ非常に詳しく読んだわけではないが、小さな編集上のコメントがある
付録L.2は納得できたが、5.2の圧縮された論証は少し粗いと感じた
特に H(X) + H(Y) = H(X | Y) + 2I(X ; Y) + H(Y | X) において、H(Y | X)を「捨てる」としている部分が明確ではなかった
3番目のトークンZを予測するなら、H(Y | X)は暗黙の文脈Cに含まれているはずで、したがって自由に捨てられないのではないか？
付録でもこの論証が明確になったようには思えない
ただし、主張の要旨を疑うほど混乱したわけではなく、概ね表現の仕方の問題だ
ならば2トークンターゲット H(X) + H(Y) のうち、どの部分が学習を助ける補助情報で、どの部分が無駄なのでしょうか？
H(X | Y) と I(X; Y) は次トークン生成に有用ですが、H(Y | X) は定義上、次トークンXとは関係のない情報量です
そのため「マルチトークン予測は、H(Y) の有用な情報 I(X; Y) を、H(Y | X) に対する無駄な計算と引き換えにする」と言えます
ただし H(Y | X) は、接頭辞 (C, X) からYを予測する次トークンエントロピーである点に注意が必要です
アテンション機構が Y|X の予測のためにすでに行った計算を次のステップへ持ち越せるなら、その計算は実際には無駄ではなく事前計算だった可能性があります
LLMは文字どおり世界に向けた1次元の窓しか持たない、という文章を読んだ
すべてはトークン列でしかない
このようなマルチ予測が、その視野を1.1次元くらいに広げてくれるのかもしれない
いずれにせよ、その窓を何らかの形で2次元以上に拡張する必要があるという実際の論拠はある
たとえばFAIR級のリソースがあり、本当に優れたJavaコーディングモデルを学習させたいなら、トークンではなくASTを予測するよう学習させるのが理にかなっている
コメントや識別子名などを予測するには、通常のLLMと組み合わせた形が依然として必要だろうが、プログラム自体をトークンストリームとしてモデル化することはしないだろう
代わりに「ifブロックを追加」「引数4つのメソッド呼び出しブロックを追加」といったものを予測させられる
また、現在のカーソル位置における型メンバーのような情報のために、コンテキストウィンドウ内の特定位置を予約するようモデルを学習させ、推論ループをIDE/LSP風の静的解析と統合することもできる
こうすれば、モデルは実際のテキストに含まれる以上の情報を見ることができる
今こうしたモデルがあまり見られない理由は、この種の研究コストが大きく、AI側の人々がみなPython中心で、PythonはIDEの恩恵をあまり受けないからだと思う
機械学習ではheadという言葉が一貫して使われておらず、混乱しやすい
この論文にはmultihead attentionとmultiple output headsという概念が両方ある
multihead attentionはtransformer構造において入力の異なる領域に注目するもので、ここでの生物学的な比喩は中央処理装置としての頭に近い
output headはニューラルネットワークの最後の層を意味し、同じ前段の層を基に異なる出力を出すものを複数置くことができる
これもゆるい生物学的比喩だが、CPUとしての頭というより、身体の片端にある頭に近い
どちらもデータを読み取るテープヘッドの比喩ではない
LLMは、自分の出力を再び入力として与えることで、かなりの部分「考えている」ように見える。そのため、モデルに声に出して考えることを強制すると推論品質が高まることが継続的に観察されている。
つまり思考連鎖推論は、質問にすぐ答えさせるのとは異なり、何を求められているのかを言い直し、答えるためにどんな情報が必要かについての上位戦略を表現し、知っている情報を述べ、その情報が初期推論にどう影響するかを説明させる。
しかし、各時点で次の複数トークンを予測させる方式は、本質的に逆の効果をもたらすのではないかと懸念される。
思考連鎖プロンプティングは、モデルが入力として n 個のトークンだけを持つ場合より、n + m 個のトークンを持つ場合のほうがより「賢い」ことを示しているように見える。
したがって、与えられた n から次の5トークンを得ることは、n から次の1トークンを得て、n+1 から次の1トークンを得る、といった方式より結果が悪くなる可能性がある。
この特定の方法がより多くのトークンを要求するという事実は重要ではない。
安価なモデルがなければ、LLMが実際の答えではなく推定値で答えようとするバイアスに常に左右されることになる。
また、ほとんどの投機的デコーディング戦略は、モデルを逐次実行した場合と同じ出力を出す。
予測が間違っていれば、そのトークンは捨てられ、速度向上だけが失われる。
トークン/単語の +1 と +2 を独立に予測するなら、文法的にどうやってなお意味の通る結果になるのだろう？ しょっちゅう壊れそうだが？
アブストラクトでは明確ではないが、図1の説明には「推論中は次トークン出力ヘッドのみを使用する。任意で、他の3つのヘッドは推論時間を短縮するために使える」とある。
上位予測をすべて採用するなら3つのヘッドすべてを使うこともできるだろうが、そうすると一般的なサンプリング戦略は使えなくなる。
ベンチマーク以外で温度0でLLMを実際に動かしている人がどれほどいるのかは分からないし、温度を適用するより優れた何かをするならまた別かもしれない。