[論文レビュー] Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking

Eric Zelikman, Georges Harik+4 more|arXiv (Cornell University)|Mar 14, 2024

Natural Language Processing Techniques被引用数 13

ひとこと要約

Quiet-STaR は Self-Taught Reasoner を一般化し、幅広いウェブテキストから推論根拠を学習させ、未来のテキストを予測する前に考えるようLMを訓練して、タスク特異的ファインチューニングなしでゼロショット推論の利得を得る。

ABSTRACT

When writing and talking, people sometimes pause to think. Although reasoning-focused works have often framed reasoning as a method of answering questions or completing agentic tasks, reasoning is implicit in almost all written text. For example, this applies to the steps not stated between the lines of a proof or to the theory of mind underlying a conversation. In the Self-Taught Reasoner (STaR, Zelikman et al. 2022), useful thinking is learned by inferring rationales from few-shot examples in question-answering and learning from those that lead to a correct answer. This is a highly constrained setting -- ideally, a language model could instead learn to infer unstated rationales in arbitrary text. We present Quiet-STaR, a generalization of STaR in which LMs learn to generate rationales at each token to explain future text, improving their predictions. We address key challenges, including 1) the computational cost of generating continuations, 2) the fact that the LM does not initially know how to generate or use internal thoughts, and 3) the need to predict beyond individual next tokens. To resolve these, we propose a tokenwise parallel sampling algorithm, using learnable tokens indicating a thought's start and end, and an extended teacher-forcing technique. Encouragingly, generated rationales disproportionately help model difficult-to-predict tokens and improve the LM's ability to directly answer difficult questions. In particular, after continued pretraining of an LM on a corpus of internet text with Quiet-STaR, we find zero-shot improvements on GSM8K (5.9%$ ightarrow$10.9%) and CommonsenseQA (36.3%$ ightarrow$47.2%) and observe a perplexity improvement of difficult tokens in natural text. Crucially, these improvements require no fine-tuning on these tasks. Quiet-STaR marks a step towards LMs that can learn to reason in a more general and scalable way.

研究の動機と目的

• 整理されていないテキストから一般的な推論を学ぶ動機づけ。
• 各トークンで中間的な推論を生成・使用するスケーラブルな方法を開発。
• メタトークンと混合機構を導入して、思考を予測に組み込みつつ訓練の安定性を損なわない。
• 考えることが難しいトークンの予測を改善し、推論タスクでゼロショットの改善を可能にする。

提案手法

• トークンごとに開始・思考終端のマーカーを用いて並列に複数の推論（思考）を生成する。
• 混合ヘッドを用いて、各トークンについてベースLMの予測とポスト思考予測を補間する。
• REINFORCE で推論生成を最適化し、トークンあたり複数の思考に対して平均化ベースラインと比較して思考支援予測を評価する。
• 現時点の次のトークンだけでなく、複数の将来トークンを教師強制で監督して訓練を安定化させる（非長視的損失を適用）。
• 7BパラメータLM（Mistral 7B）を用いてOpenWebText由来のOpenWebMathとC4を含む多様なデータから推論を学習する。
• 対角アテンションマスクを用いた並列サンプリング技術で、全トークン位置にわたって思考を効率的に生成する。

実験結果

リサーチクエスチョン

• RQ1 広範で未精査のテキストデータから有用な内部推論を生成する言語モデルを学習できるか？
• RQ2 内部の思考は、タスク特異的ファインチューニングなしで推論を要するゼロショット推論の性能を改善するか？
• RQ3 生成される思考の長さは推論性能とトークン予測性にどう影響するか？
• RQ4 メタトークンと混合機構は訓練の安定性と学習信号の質にどんな影響を与えるか？

主な発見

• Quiet-STaR は CommonsenseQA でゼロショットの改善を達成（36.3% から 47.2%）および GSM8K で（5.9% から 10.9%）。
• 改善は訓練時に使用する思考トークンの数とともにスケールする。
• C4 で訓練した場合も改善は現れるがマージンは小さい（GSM8K: 5.9% から 8.1%; CommonsenseQA: 36.3% から 42.6%）。
• より長い思考は難しいトークンに対して一般により大きな利益をもたらし、複雑な推論の取り扱いが改善される。
• Quiet-STaR はタスク特異的ファインチューニングなしで直接回答能力を向上させる。
• 内部の推論はゼロショットの連鎖的思考 prompting を補完し、推論の質と整合性を向上させる。