- Quiet-STaR는 정답이 있는 QA 데이터셋 대신 일반 웹 텍스트를 학습 신호로 삼아, 언어 모델이 말하기 전 내부 근거를 생성하도록 훈련하는 방법임
- 기존 STaR가 정답으로 이어진 근거만 골라 학습했다면, Quiet-STaR는 미래 텍스트 예측을 개선하는 근거를 보상해 비정형 텍스트의 암묵적 추론을 활용함
- 모든 토큰 위치에서 근거를 병렬 샘플링하고, 근거가 있는 예측과 없는 예측을 섞어 실제 다음 텍스트를 더 잘 맞히는 방향으로 학습함
- Mistral 7B를 OpenWebMath와 C4로 계속 사전학습한 결과, 태스크별 파인튜닝 없이 GSM8K 5.9%→10.9%, CommonsenseQA 36.3%→47.2%로 정확도가 상승함
- 효과는 예측하기 어려운 토큰에서 더 크게 나타났고, 내부 생각에 쓰는 토큰 수가 길수록 직접 추론 성능 개선도 커지는 경향을 보임
일반 텍스트를 추론 학습 대상으로 삼음
- 사람은 글을 쓰거나 말할 때 잠시 멈춰 생각하며, 텍스트의 의미 상당 부분은 문장 사이에 드러나지 않은 이유와 함의에 있음
- 기존 추론 중심 접근은 주로 질문에 답하거나 에이전트 작업을 완료하는 방식에 집중했지만, Quiet-STaR는 추론이 거의 모든 글에 암묵적으로 존재한다고 봄
- 암묵적 추론의 예시는 다음과 같음
- 증명에서 명시되지 않은 중간 단계
- 대화에서 상대의 마음 상태를 추정하는 이론
STaR를 언어 모델링으로 확장
- STaR(Self-Taught Reasoner) 는 질문-답변 데이터셋에서 소수 예시를 바탕으로 근거를 샘플링하고, 정답으로 이어진 근거만 학습해 더 어려운 문제를 반복적으로 풀도록 만든 방법임
- STaR는 고품질 QA 데이터셋처럼 세심하게 큐레이션된 데이터에 의존한다는 제약이 있음
- 이런 데이터셋은 특정 추론 과제만 담을 수 있음
- 근거 자체를 제공하거나 특정 추론 태스크를 제공해야 하는 경우가 있음
- Quiet-STaR는 특정 수학 QA 같은 태스크 대신, 대규모 인터넷 텍스트에서 미래 텍스트를 추론하는 데 도움이 되는 근거를 생성하도록 언어 모델을 학습시킴
- 이 접근은 “언어 모델은 비지도 멀티태스크 학습자”라는 언어 모델링 패러다임의 직관 위에 구축됨
학습 절차: think, talk, learn
- Quiet-STaR는 세 단계로 동작함
- think: 텍스트의 모든 토큰 뒤에서 미래 텍스트를 설명하는 근거를 병렬로 생성함
- talk: 근거가 있을 때와 없을 때의 다음 토큰 예측을 혼합함
- learn: REINFORCE 기반 보상으로 미래 텍스트 예측을 돕는 근거의 가능도는 높이고, 해치는 근거는 버림
- 학습 중 한 개의 thought에 대해 모델은
START와END같은 학습 가능한 메타 토큰으로 생각의 시작과 끝을 표시함 - 근거를 만든 뒤에는 mixing head가 그 근거 기반 예측을 미래 토큰 예측에 얼마나 반영할지 결정함
모든 토큰에서 생각하게 만드는 구현 과제
- 일반 텍스트에서는 모든 토큰마다 근거를 생성해야 하므로 계산 비용이 큼
- 이를 줄이기 위해 문자열 내 모든 토큰 위치에서 근거를 생성하는 토큰별 병렬 샘플링 알고리듬을 제안하고 구현함
- 언어 모델이 처음부터 내부 생각을 생성하거나 사용하는 법을 알지 못한다는 문제도 함께 다룸
- 생각의 시작과 끝을 나타내는 커스텀 메타 토큰을 도입함
- 모델이 근거를 생성해야 하는 시점과 그 근거를 바탕으로 예측해야 하는 시점을 학습하게 함
- 단일 다음 토큰만 보는 근시안적 학습을 피하기 위해 여러 토큰 앞까지 포함하는 비근시안 손실(non-myopic loss) 을 사용함
- 확장된 teacher-forcing 기법으로 개별 다음 토큰 너머의 예측까지 학습에 반영함
실험 설정과 결과
- 실험은 Mistral 7B에 Quiet-STaR를 적용해 수행됨
- 계속 사전학습에는 웹 텍스트 데이터셋인 OpenWebMath와 C4(Colossal Clean Crawled Corpus) 가 사용됨
- 태스크별 파인튜닝 없이 제로샷 직접 추론 성능이 개선됨
- GSM8K: 5.9%→10.9%
- CommonsenseQA: 36.3%→47.2%
- GSM8K와 CommonsenseQA 모두에서 Quiet-STaR 학습 중 사용한 생각 토큰 수가 길어질수록 성능 개선이 일관되게 증가함
- 자연 텍스트에서는 예측하기 어려운 토큰의 perplexity가 개선됨
- 생성된 근거는 특히 예측하기 어려운 토큰에 불균형적으로 더 큰 도움을 줌
Quiet-STaR의 기여
- Quiet-STaR는 큐레이션된 추론 태스크가 아니라 다양한 비정형 텍스트 데이터에서 추론을 학습하도록 STaR를 일반화함
- 병렬 샘플링 알고리듬으로 주어진 문자열의 모든 토큰 위치에서 근거를 생성하는 학습 절차를 확장 가능하게 만듦
- 생각의 시작과 끝을 나타내는 커스텀 메타 토큰은 모델이 근거 생성과 근거 기반 예측의 타이밍을 학습하는 데 사용됨
- mixing head는 특정 thought에서 나온 다음 토큰 예측을 현재 예측에 얼마나 반영할지 사후적으로 결정함
- 여러 토큰 앞을 포함하는 언어 모델링 손실이 생각의 효과를 개선함
- 여러 태스크에서 생각을 사용하면 같은 웹 텍스트로 학습한 모델보다 어려운 토큰을 더 잘 예측하고, 더 긴 thought에서 개선 폭이 커짐