[논문 리뷰] Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking

Eric Zelikman, Georges Harik외 4명|arXiv (Cornell University)|2024. 03. 14.

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.

연구 동기 및 목표

• 정제된 데이터셋이 아닌 비구조적 텍스트로부터 일반적 추론 학습을 유도한다.
• 토큰마다 중간 합리화(rationale)를 생성하고 활용하는 확장 가능한 방법을 개발한다.
• 메타 토큰과 예측에 생각을 통합하는 혼합 메커니즘을 도입하여 학습 불안정성을 초래하지 않으면서 생각을 반영한다.
• 생각이 어려운 토큰의 예측을 개선하고 추론 과제에서 제로샷 개선을 가능하게 함을 보여준다.

제안 방법

• 생각의 시작과 끝을 표시하는 토큰 마커를 사용하여 각 토큰 뒤에 다중 합리화(생각)를 병렬로 생성한다.
• 각 토큰마다 기본 LM 예측과 사후 생각 예측 사이를 보간하기 위해 mixing head를 사용한다.
• REINFORCE로 합리화 생성을 최적화하고, 토큰당 다수의 생각을 통한 예측을 평균 기반의 기준선과 비교한다.
• 즉시 다음 토큰을 넘어 다수의 미래 토큰을 감독하는 non-myopic loss(교사강요)를 적용하여 학습 안정성을 높인다.
• 다양한 데이터로부터 추론을 학습하기 위해 OpenWebText에서 파생된 OpenWebMath와 C4를 포함한 웹 텍스트 말뭉치를 7B 매개변수의 LM(Mistral 7B)으로 학습한다.
• 대각선 주의(mask)를 가진 병렬 샘플링 기법을 사용하여 모든 토큰 위치에서 생각을 효율적으로 생성한다.

실험 결과

연구 질문

• RQ1 Can a language model learn to generate useful internal rationales from broad, uncurated text data?
• RQ2 Do internal thoughts improve zero-shot reasoning performance on tasks requiring reasoning without task-specific fine-tuning?
• RQ3 How does the length of generated thoughts affect reasoning performance and token predictability?
• RQ4 What are the effects of meta-tokens and the mixing mechanism on training stability and learning signal quality?

주요 결과

• Quiet-STaR yields zero-shot improvements on CommonsenseQA (36.3% to 47.2%) and GSM8K (5.9% to 10.9%).
• Improvements scale with the number of thinking tokens used during training.
• Improvements also appear when training on C4, though to a smaller margin (GSM8K: 5.9% to 8.1%; CommonsenseQA: 36.3% to 42.6%).
• Longer thoughts generally provide greater benefits for difficult tokens, indicating better handling of complex reasoning.
• Quiet-STaR enhances direct-answering abilities without task-specific fine-tuning.
• Internal rationales can complement zero-shot chain-of-thought prompting by improving the quality and coherence of reasoning.