こんにちは、CCCMKホールディングスTECH LABの三浦です。
いつの間にか2月も中旬になりました。2月3月はあっという間に過ぎていく印象があります。きっと気が付いたら4月になっているんだろうな、と思います。
"Test-Time"
最近LLM周りで"Test-Time Scaling"というフレーズを聞く機会があって、いったい何のことだろう？？と気になって調べていました。全ての情報を追い切れていないのですが、どうやら"Test-Time"(実際のタスクを解く時、という意味だと思います。)に計算処理を加えることでLLMの回答精度を高める、というアプローチに最近注目が集まっているようです。
有名なアプローチだと"Chain-of-Thought(CoT)"が挙げられます。複雑なタスクではLLMに単にプロンプトを与えてトークンを生成させる方法では上手くいかないことが多いのですが、タスクを解くために必要なステップを考えさせ、それらを一つ一つ解かせることで複雑なタスクにも対応が出来るようになることがあります。他にも同じプロンプトに対する回答生成を複数回平行して実行し、その結果を集約して最も多い回答を採用するアプローチや、タスクを解くためのプログラムを作り、実行させるアプローチなどもあります。個人的にはOpenAIの"o1"モデルが出てから"Test-Time"における処理が注目され始めてきたように感じています。
従来の機械学習の発想で考えると、モデル推論時に学習の際に考慮していなかった複雑な追加処理を加える、というアプローチは少し違和感を感じました。しかし私たち"人間の思考"という観点で考えると、学習時に知識を取得し、テストの際にその知識を活用し、時間をかけて色々と頭を動かして問題を解く、というのはごくごく自然に行っていることです。問題が複雑になればなる程、考える時間の工夫をすることでその問題に対処出来るようになる、ということも自然なことだと思います。
特に今後LLMをエンジンにした様々なタスクに対応できる"Agent"の開発が活発化していく中で、"Test-Time"をいかに効率的に使って正確な回答を導き出せるか、という観点の研究は必要不可欠なのだと思います。
論文について
さて、私が最近読んだ論文では"Test-Time"の際に"Test-Time Training(TTT)"という学習を実行して複雑なタスクに対する精度を上げる、というアプローチを取っています。なかなか面白いアプローチで、自分が関わっている領域でも活用できる機会がありそうだな、と感じました。今回の記事ではこの論文の内容についてご紹介させて頂きたいと思います。
今回記事を書くにあたり参照した論文はこちらになります。
Title: The Surprising Effectiveness of Test-Time Training for Abstract Reasoning
Author: Ekin Akyürek, Mehul Damani, Linlu Qiu, Han Guo, Yoon Kim, Jacob Andreas
Submitted: 11 Nov 2024 arXiv: https://arxiv.org/abs/2411.07279
Test-Time Training?
"Test-Time Training"ってすごく不思議なフレーズだな、と感じました。Trainingはモデル学習時しかしないはずなのに、どういうことだろう、と思います。この論文の中で取り扱われているタスクは"Abstraction and Reasoning Corpus(ARC)"というもので、視覚的なパズルをLLMに解かせます。パズルは最大30x30の格子状に配置された複数の色で構成されるパターンを、問題の例示から推測される変換処理を使って変換すると、どんなパターンになるのかを答えさせる、という、かなり難しいタスクです。
実際にはLLMには画像情報ではなく各パターンをテキストに変換した情報が与えられます。問題の中には回転や反転、といった単純な変換だけでなく、重力を考慮した変換なども含まれているそうです。
さて、"Test-Time Training(TTT)"の手続きについてまとめてみます。ARCの各問題にはいくつかの例示サンプル(変換前と変換後のパターンのペア)とテスト用のサンプルが含まれています。通常LLMでこういった問題を扱う場合は例示サンプルをプロンプトに含めた"In-context Learning"のアプローチを取ることが多いと思いますが、TTTでは例示サンプルを増強して生成した学習用データセットを用いてLoRAチューニングを実行します。LoRAのアダプタパラメータは、全部の問題を通じて学習するより、各問題ごとに個別に用意して学習する方が精度が高くなるそうです。
TTTの学習データセット
問題に含められた例示サンプルを増強してTTT用の学習データセットを構築します。具体的な方法は、まず例示サンプルから1つのペアをピックアップし、問題用のペアとして設定します。その後、回転や左右反転といった可逆的な変換処理を施してデータを増強します。論文に掲載されている以下の図のようになります。
TTTのLoss
TTTでは次の式を最小化するようにLoRAアダプタのパラメータ を最適化します。
この式の はARCのi番目の問題に対する先述した方法で生成した学習用データセットを表していて、
は交差エントロピーを表しています。最初の項は Demonstration Lossと論文の中では呼ばれていて、例示サンプルに対するLossになっています。最後の項は例示サンプルと変換前のパターンから変換後のパターンを生成させるためのLossです。
ちなみにTTTから推論実行までを100個の問題で完了するのにNVIDIA A100GPU1基で約12時間有するそうです。
TTTの後の推論の工夫
TTTだけでなく、TTTの後の推論時にさらに工夫を加えることで、問題の正解率を向上させることが出来るそうです。"Test-Time"で使われることが多い、複数の回答を生成させて出現回数が多い結果を採用する、というアプローチを応用しています。
複数の回答を生成するための手法
ある問題に対して複数の回答を生成させるために、TTTの学習用データを作った時と同様、パターンの可逆変換を利用しています。LLMには元の問題に含まれているパターンにある可逆変換を施した問題を与え、回答を生成させます。その回答に逆の変換を施すことで、元の問題と対応した回答にすることが出来ることを利用しています。
回答の選び方
頻度が高い回答の選択は、2段階で行います。最初に変換ごとに複数の回答から出現頻度が高い回答を選択します。その後、それらの回答からさらに出現頻度が高い回答を選択する、という方法です。
TTTの後の推論については以下の論文の図のようになります。
まとめ
今回は"Test-Time Training"という興味深いアプローチについて述べられている論文を読んで、そのアプローチの内容についてまとめてみました。モデルの学習フェーズではなく、テストフェーズで新規データに対して学習させる、というのはとても面白く、自分も今度試してみたいなと思いました。 個人的に、最近LLM周りでなんとなく新しい研究の方向性が生まれてきているように感じています。最新の動向に、これからも注目していきたいです。
