見出し画像

【論文瞬読】言語モデルの推論を高速化!Stepwise Internalizationが切り開く新たな可能性

こんにちは!株式会社AI Nestです。近年、自然言言語処理の分野では言語モデルの発展が目覚ましく、特にTransformerベースのモデルは様々なタスクで高い性能を示しています。しかし、複雑な推論タスクになると、明示的なChain-of-Thought (CoT) 推論が必要となり、計算コストが大きな課題となっています。そんな中、最新の研究で提案されたStepwise Internalizationという手法が注目を集めています。

タイトル:From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step
URL:https://arxiv.org/abs/2405.14838
所属:Allen Institute for Artificial Intelligence, University of Waterloo, University of Washington, Harvard University
著者:Yuntian Deng, Yejin Choi, Stuart Shieber

Chain-of-Thought (CoT) 推論とは?

CoT推論とは、言語モデルが複雑な問題を解く際に、推論の中間ステップを明示的に生成する手法です。例えば、多桁の掛け算を行う場合、途中の計算過程を出力することで、最終的な答えの精度を高めることができます。しかし、この明示的なCoT推論は計算コストが高く、特に長い推論チェーンが必要な場合は実用的ではありませんでした。

Stepwise Internalizationが切り開く新しい道

そこで登場したのが、Stepwise Internalizationです。この手法は、明示的なCoT推論から始めて、徐々に中間ステップを取り除きながらモデルをファインチューニングしていくというものです。これにより、推論ステップがモデルの隠れ状態に内在化され、最終的には完全な暗黙的なCoT推論が実現されます。

Figure1, Stepwise Internalizationの概要

Stepwise Internalizationの魅力は、その単純さと効果の高さにあります。実験では、GPT-2 Smallのようなモデルでも、9桁の乗算を99%の精度で解けるようになりました。これは、明示的なCoT推論を用いた場合と同等の精度ですが、推論速度は最大11倍にまで向上しています。

Table2, 乗算タスクにおける各手法の性能比較

さらに、Stepwise Internalizationは、速度と精度のトレードオフを可能にします。CoTトークンを内在化することで、精度を維持しつつ速度が向上するのです。

Figure2a, Stepwise Internalizationにおける速度と精度のトレードオフ

課題と今後の展望

もちろん、Stepwise Internalizationにも課題はあります。高い訓練コストや不安定性の問題、解釈可能性の低下などです。しかし、レイテンシーと精度のトレードオフを可能にする利点は大きく、実用上の価値は高いと言えるでしょう。

また、この手法は乗算や学校数学の推論タスクだけでなく、より幅広いタスクへの応用が期待されます。例えば、GSM8Kデータセットを用いた実験では、Stepwise InternalizationがGPT-4をも上回る性能を示しました。

Table3, GSM8Kデータセットにおける各手法の精度比較

今後は、訓練の効率化や安定化、様々なタスクへの適用などが研究の焦点になっていくと予想されます。

おわりに

Stepwise Internalizationは、言語モデルの推論を高速化する画期的な手法であり、自然言語処理の分野に新しい風を吹き込んでいます。明示的なCoT推論と暗黙的なCoT推論を巧みに組み合わせることで、高い精度と速度を両立することに成功しました。この研究は、言語モデルの実用化に向けた重要な一歩であり、今後のさらなる発展が楽しみです。