【自作LLM実験③】普通のTransformerより強くなった。AIに「2つの見方」を持たせたら、得意分野まで生まれた

第一回の記事では、LLMの中身を2本にして、同じ文章を別々に処理させる実験をしました。

        ┌→ LLM A ─┐
入力 ───┤          ├→ 出力
        └→ LLM B ─┘

狙いは単純で、一人で考えるより、二人で考えた方が賢くなるんじゃないか？ というものです。

実際に学習させてみると、2本のLLMはだんだん同じ状態ではなくなり、それぞれ少し違う内部表現を作るようになりました。ところが、最終的な性能では、同じくらいの計算量を使って一本のTransformerを大きくしたモデルに負けました。

つまり、違う考え方を2つ持たせるだけでは足りなかった。

2本のモデルが別々の情報を持っていても、それをうまく利用できなければ意味がない。そこで次に気になったのが、

ということでした。

今回は、2本のLLMをそのまま並べる方法ではありません。

モデル全体は一本のままにして、各層の中で、同じ文章を「近くを見るAttention」と「文章全体を見るAttention」の2通りで見ます。そして、その2つがどれくらい一致しているか、どれくらい食い違っているかまで、次の計算を行うFFNに渡す構造を作りました。

つまり、

第一回
2本のLLMが自然に違う状態を作る
        ↓
でも、その違いをうまく使えなかった

今回
1本のLLMの中で意図的に2つの見方を作る
        ↓
その「違い」自体を計算に使う

という流れです。

この新しい構造を実際に作り、ゼロから事前学習させ、普通のTransformerとほぼ同じ計算量で比較しました。

結果は、

今回作った新しいTransformerの方が、文章予測性能で勝ちました。

しかも面白かったのは、単にNLLという数字が少し良くなっただけではありません。

ベンチマークテストでは、文章の続きを考える問題や、文章として自然かどうかを見る問題で特に強くなりました。 さらに自由に文章を書かせると、同じ言葉や文章を繰り返し続けるループも大幅に減りました。

つまり今回のモデルは、ただ全体的に少し性能が上がっただけではなく、この構造がもたらす得意分野まで現れた。

今回はLLM全体を2本にしない

今回は、モデル全体を2本にするのをやめました。普通のTransformerと同じように、メインの情報の流れは一本です。ただし、各層の中で文章を見る方法を2種類にしました。

一つは、近くを見るAttention。 もう一つは、文章全体を見るAttention。

Attentionって何？

Attentionは、今読んでいる文章の中から「どの情報を拾って、今の単語の状態をどう変えるか」を決める仕組みです。 訓練を通して、「こういう文脈では、この辺の情報を見ると役に立つ」という見方そのものを学んでいきます。

だから今回のモデルでは、

              ┌→ 近くを見るAttention
今の単語 ─────┤
              └→ 全体を見るAttention

という2つの見方を同時に作ります。

たとえば、

という文章があったとします。「冷蔵庫を開けた」のすぐ近くだけを見るのと、「昨日スーパーで牛乳を買った」という文章全体の流れを見るのでは、見えてくるものが違います。

近くだけ見れば、「冷蔵庫を開けた → 中を見た」みたいな流れが強いかもしれない。でも文章全体を見ると、「牛乳を買った → 冷蔵庫に入れる」という流れも見えてくる。

そこで思いました。この2つの見方を、すぐに一つに混ぜてしまわない方がいいんじゃないか？

「2つの答え」だけではなく、「ズレ」まで使わせる

今回作ったモデルでは、近くを見るAttentionが作った情報を F、文章全体を見るAttentionが作った情報を G とします。

大事なのは、

F + G

だけではありません。モデルに、FとGはどれくらい同じことを言っている？ どれくらい食い違っている？ どっちの主張が強い？ という情報まで渡します。

イメージとしては、

近くを見るAttention
        ↓
        F ─────┐
                ├→ 2つの違いも調べる
        G ─────┘
        ↑
全体を見るAttention

という感じです。そして、この「2つの見方の関係」を、FFN という部分に渡します。

FFNって何？

ものすごく簡単に言うと、FFNは、Attentionが集めてきた情報を受け取って、「今の状態を次にどう変えるか」を決める部分です。

Attentionが「どこを見るか」「何を拾うか」を担当するなら、FFNは「その情報をどう解釈して、次の状態へどう変換するか」を担当している、と考えると分かりやすいです。

普通のTransformerでは、Attentionの情報が一つに混ざったあと、その一つの状態をFFNが読みます。今回のCPS Transformerでは、それに加えて、「近くを見るAttentionと、全体を見るAttentionが、どれくらい一致しているか、どれくらい食い違っているか」 という情報もFFNに渡します。

人間っぽく言えば、

という仕組みです。

このモデルを、Complementary Pair-Synthesis Transformer、略して CPS Transformer と名付けました。

本当に普通のTransformerより強いのか？

ここが一番大事です。新しい仕組みを作って「なんとなく賢そう！」では意味がありません。

そこで比較用に、普通のTransformerを横に太くしたモデル も作りました。CPSだけ余計に大量の計算をして勝っても面白くないので、両方の計算量がほぼ同じになるように調整しています。

モデルの大きさは約2億3400万パラメータ。そして両方とも、12億token までゼロから事前学習させました。

まずは大量の文章で、次の単語をどれだけ予測できるか

最初に見るのは、いちばん基本的な性能です。学習に使っていない文章を大量に用意して、次の単語を、どちらのモデルがより正確に予測できるか？ を比べました。

全部で約1678万token。これを1024tokenずつ、1万6384個の文章区間 に分けて比較しています。

結果は、

普通のTransformer     NLL 3.2235
CPS Transformer      NLL 3.1841

でした。

NLLというのは、正しい次の単語に、どれくらい高い確率を付けられたか を見る数字です。小さいほど強い。なので、文章予測性能ではCPS Transformerが勝ちました。

さらに面白いのが、平均値だけではありません。1万6384個の文章区間を一つずつ比べると、

CPS Transformerが勝った区間     15,432
普通のTransformerが勝った区間      952

でした。

つまり、94.2%の文章区間で、CPS Transformerの方が次の単語をうまく予測できていました。 一部の文章だけで大勝して平均値を押し下げたのではなく、ほとんどの文章で、少しずつCPSの方が強かった。 これはかなりきれいな結果でした。

じゃあ、ベンチマークテストではどうなる？

次に、LLMの性能を測るためによく使われるベンチマークテストもいくつか試しました。

ここで面白かったのが、全部のジャンルで均等に強くなったわけではなかったことです。CPS Transformerには、得意そうな分野 が出てきました。

LAMBADA：長い文章の続きを考えるベンチマーク

LAMBADAは、長めの文章を読ませて、その文脈から最後の単語を予測させるベンチマークテストです。かなり前の文章まで使わないと、正しく予測しにくい問題が多く含まれています。

結果は、

普通のTransformer     PPL 224.24
CPS Transformer      PPL 168.79

でした。

PPLも、小さいほど、正しい文章の続きを高い確率で予測できている という数字です。ここでは、かなり大きな差が出ました。

面白いのは、完全に正解の単語を1位にできた回数自体は、そこまで大きく変わっていないことです。でも、正しい続きを「それっぽい」と高く評価する力はかなり上がっていた。

つまりCPS Transformerは、長い文章の流れを見ながら、「この続きの方が自然そうだ」と判断する内部の確率そのものを、かなり良い方向へ動かしていたように見えます。

HellaSwag：「この続きが一番自然なのはどれ？」

HellaSwagは、「この文章の続きとして、一番自然なのはどれ？」 という選択式のベンチマークテストです。

ここでもCPS Transformerが勝ちました。

普通のTransformer     28.63%
CPS Transformer      29.79%

差だけ見ると1%ちょっとですが、問題数が1万問以上あるので、かなり安定して出た差です。

BLiMP：英文として自然なのはどっち？

BLiMPは、「英語として自然なのはどちら？」 を見る文法系のベンチマークテストです。

今回使った固定670問では、

普通のTransformer     75.82%
CPS Transformer      79.40%

でした。ここもかなり面白い。

でも、全部で勝ったわけではない

一方で、MMLU、ARC、BoolQ、WinoGrandeのような知識問題や一般的な選択問題では、大きな差は出ませんでした。

これがかなり面白いところです。ただ全体が少し賢くなったわけではない。 今回入れた仕組みに関係しそうなところで、特に差が出た。

文章の流れ。続きを読む力。文法的な自然さ。こういうところで強くなった。

まるで、アーキテクチャによって「このモデルはこういう問題が得意」という性格が生まれた ように見えます。

そして、文章のループがほぼ半分になった

個人的に一番びっくりしたのはこれです。

小さいLLMに文章を書かせると、途中から、

みたいに、同じ言葉や似た文章を延々繰り返す ことがあります。LLMを触ったことがある人なら、一度くらい見たことがあると思います。

そこで5000個の文章の書き出しを用意して、普通のTransformerとCPSに、まったく同じ条件で続きを書かせました。

結果、同じ4単語のまとまりを繰り返す割合が、

普通のTransformer     6.56%
CPS Transformer      3.37%

になりました。

ほぼ半分です。

さらに、64token以上にわたって同じような文章を繰り返す長いループは、

普通のTransformer     6.42%
CPS Transformer      1.78%

まで減りました。かなり大きな差です。

反復を抑える設定を入れても、差は残った

「いや、それなら生成設定で反復を抑えればいいじゃん」と思う人もいると思います。そこで、反復しにくくする設定を入れた状態でも比較しました。

結果は、

普通のTransformer     3.00%
CPS Transformer      1.25%

でした。

反復を抑える設定を入れても、CPSの方がさらにループしにくい。 しかも、両方とも出力する文章の長さはほぼ同じでした。CPSが早く文章を終わらせているから、ループが少なく見えたわけでもありません。

これはかなり興味深い結果です。

なぜループが減ったんだろう？

ここからは、結果を見ながら考えた話です。

文章を生成している途中で、モデルが最近の数単語だけを強く見すぎるとします。すると、

という感じで、自分で自分を反復方向へ押してしまう ことがあります。

でもCPSには、近くを見るAttention とは別に、文章全体を見るAttention があります。

近くだけ見ている方が「またこれを書こう！」と言っていても、全体を見る方は「いや、文章全体ではもう別の話に進んでるぞ」という違う情報を持っているかもしれません。そしてCPSは、この2つの意見がズレていること自体 をFFNの計算に使います。

だから、局所的なループに吸い込まれそうになったところで、全体側の情報がブレーキになっている可能性があります。これが本当の理由だと証明したわけではありません。でも、この構造と今回の結果はかなりきれいにつながっています。

モデルの中を覗いたら、本当に「2つの見方」が残っていた

さらにモデル内部も観察しました。

もし学習しているうちに、近くを見るAttentionと全体を見るAttentionが、結局まったく同じことを考えるようになっていたら、この仕組みはあまり面白くありません。

でも実際には、かなり違う状態のまま残っていました。

さらに、「この2つがどう違うか」を表す内部情報も、学習が進むほど複雑になっていきました。最初は比較的単純だったものが、学習を続けるにつれて、いろいろな方向を使うようになっていった。

しかも、その「2つの違い」がFFNへ与える影響も、学習が進むほど大きくなっていきました。モデル自身が、「この2つの違い、使えるじゃん」 と学習していったようにも見えます。

もしかして「意味空間を広くする」方法は、幅を広げるだけじゃない？

ここで前回の記事につながります。

前回は、大きな一本のTransformerと、小さいTransformerを2本にしたものを比べました。結果は、大きな一本の勝ち。

だから、

と考えました。

今回の結果を見ると、もしかすると答えの一つは、「複数の見方を持たせるだけではなく、その見方同士の関係まで計算に使わせる」 ことなのかもしれません。

ただ2人並べるだけではない。2人が何を考えていて、どこで意見が一致して、どこで意見が割れているのかまで考える。 そこに情報がある。

そう考えると、かなりワクワクします。

じゃあ「2つ」で止める必要なくない？

そして当然、もっと先を想像したくなります。

今回は、

近くを見る
全体を見る

の2つでした。

でも、3つにしたら？

近くを見る
少し広く見る
全体を見る

さらに、5つにしたら？

近くを見る
中距離を見る
全体を見る
登場人物を追う
過去の出来事を覚える

みたいな状態を、別々に持たせる。

そして、「5人のうち4人は同じ意見なのに、1人だけ違う」とか、「近くを見る2人はAと言っているけど、長期記憶はBと言っている」とか、「全体としてはBだけど、文法的にはAしかありえない」みたいな関係まで計算に使わせる。

これ、面白くないですか？

「一つの巨大な頭」から「複数の視点を持つ頭」へ

今のTransformerは、ものすごく大雑把に言えば、巨大な一つの状態 を何層も変換していく仕組みです。

だったら将来は、一つの巨大な状態をさらに巨大化するだけではなく、いくつもの異なる状態を同時に持って、それらを比較しながら考える 方向もあるかもしれません。

たとえば、複数の内部状態がそれぞれ違う役割を持つ。そして最後に全部混ぜるのではなく、混ぜる前の違いそのものを使う。

もしこれがモデルを大きくしても効くなら、かなり面白いことになると思います。

234Mでこれなら、もっと大きくしたら？

今回のモデルは、約2億3400万パラメータです。ChatGPTやClaudeのような巨大モデルと比べれば、かなり小さい。学習量も12億tokenです。

それでも、

• 文章予測性能が上がった

• 長い文脈を使うベンチマークで差が出た

• 文法系でも強くなった

• 自由生成のループが大幅に減った

• モデル内部でも二つの見方が別々に育った

という変化が出ました。

では、10億パラメータなら？ 30億なら？ 100億なら？

さらに、2つではなく3つ、5つ、もっと多くの見方を持たせたら？ それぞれの見方をもっと専門化させたら？ 文章だけではなく、画像、音、記憶、検索結果まで、別々の「見方」として持たせたら？

ここから先は、完全に妄想です。でも、最初の2状態版を実際に作って、普通のTransformerより性能が上がった という結果があるので、ただの空想ではなく、「次に試したくなる空想」 になりました。

自作LLMはここが面白い

前回の記事では、

という話を書きました。

今回まさに、「俺の考えたTransformer、普通のTransformerより強かった」 を実際にやることができました。

もちろん、世界最大のAIを作ったわけではありません。家庭用GPUで訓練できる小さいモデルです。

でも、自分で「こういう構造にしたら面白くない？」と考えて、AIと一緒に設計して、実装して、何日もGPUを回して、結果を見たら、本当に性能が上がった。

これは普通に、めちゃくちゃ面白いです。

論文として公開しました

今回のCPS Transformerは、構造、数式、学習結果、内部観察をまとめて論文にし、Zenodoで公開しました。

Complementary Pair-Synthesis Transformer

「こういう構造を作って、実際に学習させたらこうなった」 という記録を残しました。

まだTransformerには、かなり遊べる余地がありそう

Transformerは、もう完成された仕組みに見えるかもしれません。でも実際に中身をいじってみると、まだ、

みたいな、かなり素朴な疑問を試せる場所が残っています。

そして、その素朴な疑問から作ったモデルが、今回は本当に普通のTransformerより強くなりました。

以前なら、「AIの中に2つの見方を作って、その意見の違いまで使わせたら面白くない？」なんて思いついても、そこで終わっていたと思います。

でも今は、AIと一緒に構造を考えて、コードにして、家庭用GPUで学習させて、普通のTransformerと比較して、ベンチマークテストまでできる。

「こんなのどうだろう？」という妄想が、数日後には、数字になって返ってくる。 そして数字が良かったら、また次の妄想が始まる。

今回のCPS Transformerは、その中で初めて、「新しい構造を作ったら、本当に性能が上がった」 ところまで行けたモデルになりました。

Transformerには、まだかなり面白い設計空間が残っていそうです。