LLMに潜在空間で思考させる"Chain of Continuous Thought"についての論文を読みました。

こんにちは、CCCMKホールディングス AIエンジニアの三浦です。

先日海に行きました。景色もそうですが、海の近くでは普段の生活であまり感じない匂いなんかも感じられて新鮮な気持ちになりました。

はじめに

LLMの推論精度を向上させる方法に"CoT(Chain of Thought)"があります。これはLLMにすぐに回答を生成させるのではなく回答に必要になる思考のプロセスを生成し回答させることで、論理的な思考が必要になる問題にも対処出来るようにするテクニックです。

たとえば

AならばB, BならばC, DならばE, CならばF
が成り立つとき、AならばFは成立する？

といった問題を考えてみます。CoTを使わない場合はLLMに"成立する", "成立しない"だけを回答させるのですが、CoTを使う場合は"AならばB,BならばC,CならばF,よってAならばF"といったような中間思考を生成させた上で"成立する"または"成立しない"という問題に対する最終回答を生成させます。

CoTの特徴は中間思考を自然言語で表現しているという点です。言い換えると思考を言語(トークン)に変換して行っています。LLMはテキストを受け取ると、それらをトークンに分割し、埋め込み表現化した後に複数の層を通過しながら計算し、最後に得られる潜在表現をソフトマックス関数に入力し、スコアが高いトークンを1つだけ出力します。最後のトークンを出力する箇所で、潜在表現に含まれる大部分の情報を破棄している、とも言えます。モデルの出力を言語化するためにはトークン化が必要ですが、中間思考においては必ずしも言語化は必要ありません。あえてトークン化を行わず、途中で得られた潜在表現をすべて利用したほうが良いという考え方も出来そうです。

最近読んだMeta社の論文"Training Large Language Models to Reason in a Continuous Latent Space"では、"COCONUT(Chain of Continuous Thought)"という手法が提案されています。これは中間思考の段階で生成される潜在表現をトークン化せずにそのままモデルに入力する、というアイデアです。

このアイデアがとても面白いな、と思ったので今回は論文の内容をまとめてみたいと思います。

参照論文

Title: Training Large Language Models to Reason in a Continuous Latent Space
Authors: Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, Yuandong Tian
Submitted: 9 Dec 2024 (v1), last revised 11 Dec 2024
arXiv: https://arxiv.org/abs/2412.06769

CoTにとって難しい問題

CoTで解くのが難しい問題は、回答にたどり着くまでに色々な可能性を探る必要があるような問題です。たとえば論文の中で扱われている"ProsQA"のようなタスクです。

ProsQAはこの論文の中で提案されているタスクで、下の図のようにランダムなコンセプトをノードで表現し、それらを"すべてのAはBである"といった依存関係で接続した有向グラフに基づいて、ある二つのノード間に依存関係が成立するかどうかを判断させる問題設定です。1つのノードから複数のパスが存在するため、それぞれを探索する必要がありますが、CoTのように1つのトークンを決定してしまう方法だと複数の可能性を探索させることが難しくなってしまいます。

Training Large Language Models to Reason in a Continuous Latent Space, Figure 6

一方トークン化する前の潜在表現の中には次のトークン(パス)の可能性が連続値として含まれています。たとえば下の図のような計算の文章題をLLMに解かせる場合、答え 「540 (=3×3×60)」 に至るまでの計算過程の数値として潜在表現の中には"180"と"9"の可能性も含まれます。これらの情報を破棄せずにそのまま思考に利用することで、COCONUTでは複数の可能性を考慮した思考が実現出来ます。

Training Large Language Models to Reason in a Continuous Latent Space, Figure 4

COCONUTの仕組み

COCONUTでの推論処理では、"latent"モードと"language"モードという二つの動作モードを切り替えて実行されます。LLMに質問文を入力すると、"<bot>"という特殊なトークンを追加した後にすぐにlatentモードに移行します。latentモードではLLMが出力する最終層の潜在表現をトークン化せずにそのまま次の入力に加えられます。latentモード中はこれを繰り返し、latentモード終了時は"<eot>"という特殊トークンを追加した後にlanguageモードに移行します。languageモードでは従来通りトークンが生成されます。

latentモードからlanguageモードにいつ切り替えるのかについては分類モデルを用いて判断する方法、latentモードの実行回数をあらかじめ決めておく方法が挙げられています。論文の検証はあらかじめ回数を決めておく方法が採用されています。

COCONUTの学習

COCONUTではLLMの学習にCoTのデータを使用し、複数のステージに分けて学習が行われます。以下の図のように、ステージを進めるにつれCoTのステップを取り除き、その分latentモードによる思考処理に置き換えて学習を行っていきます。

Training Large Language Models to Reason in a Continuous Latent Space, Figure 2

試行回数を増やすと精度が向上する

COCONUTのlatentモードにおける試行回数を増やすことで回答精度が向上する、という検証が行われています。下の図のように、回数を増やすことでAccuracyが向上しています。ただ回数が3, 4回くらいで精度向上の頭打ちになるようです。回答の内訳(右の図)を見ると、回数を増やすことでハルシネーションの発生を抑えることが出来、それが精度向上につながっていることがうかがえます。

Training Large Language Models to Reason in a Continuous Latent Space, Figure 5

潜在空間を利用することの利点

COCONUTで潜在空間を利用して思考を行うことの利点は、可能性のあるパスを打ち切らずに幅広く考慮して問題を解くための計画に利用できる点です。そのためProsQAのような複数のパスを同時に考慮し計画を立てなければならない問題で精度の向上が見込まれます。下の図ではlatentモードで思考を行った後に遷移先として考えられる各ノードに対し、その単語を生成する生成確率を計算しスコアとして評価をしています。1回の思考が終わった後(左)はいくつかのノードに対して高いスコアが付与されていますが、2回の思考を実行すると1つのノードに高いスコアを出力しており、可能性のあるノードを絞ることが出来ていることが分かります。

Training Large Language Models to Reason in a Continuous Latent Space, Figure 7

まとめ

今回はCOCONUTという、LLMに潜在空間を利用して思考させる、というアプローチに関する論文を読み、その内容をまとめてみました。思考の内容は言語化しないほうが上手くいく、というのは発想がとても面白いですし、よくよく考えてみると人間の思考プロセスと近しいのかも、と思いました。論文の中では比較的サイズが小さいGPT-2が検証に使われていますが、今後思考能力に長けたLLMの学習において、このような手法が採用されていくのではないか、と感じました。