Quiet-STaR: 풀이집 없이도 모델은 도움이 되는 중간 텍스트를 배울 수 있을까
언어 모델이 더 잘 추론하도록 만들려면 풀이가 달린 문제를 많이 준비해야 할까. Quiet-STaR는 조금 다른 출발점을 택한다. 이미 존재하는 글의 다음 부분을 더 잘 예측하는 데 도움이 되는 중간 텍스트를 모델이 생성하고, 그 유용성으로 학습한다. 정답 풀이집을 그대로 모방하는 대신 일상적인 언어 모델링 안에서 추가 계산을 학습 신호와 연결하려는 시도다.
다음 문장을 이해하기 위한 보조 경로
일반적인 언어 모델은 앞 문맥에서 다음 토큰을 예측한다. Quiet-STaR는 각 위치에서 앞 문맥을 보고 짧은 rationale을 생성한 뒤, 그 텍스트를 사용해 실제 이어지는 부분을 예측하게 한다. 논문은 이 rationale을 thought라고 부른다. 사용자가 보는 완성 답변과 별도로 계산에 쓰이는 토큰 경로다.
중요한 것은 rationale의 정답을 사람이 직접 제공하지 않는다는 점이다. 이어질 실제 글은 학습 자료에 있으므로, 보조 텍스트를 사용했을 때 그 글에 더 높은 확률을 주는지 측정할 수 있다. 유용한 경로를 더 자주 생성하도록 학습하는 것이다. Quiet-STaR 원문
설명용으로 짧은 문장을 만들어 보자. '정기 점검으로 서버를 중단했다. 따라서 요청은 잠시 ...' 뒤에 어떤 표현이 올지 예측한다고 하자. 앞의 사건과 뒤의 결과를 연결하는 보조 텍스트가 이어지는 내용을 예측하는 데 도움을 줄 수 있다. 반면 흔한 조사나 문장부호를 예측하는 위치에서는 추가 텍스트가 별로 필요하지 않을 수 있다.
이 예는 모델이 반드시 사람이 생각하는 방식으로 인과 관계를 이해한다는 뜻이 아니다. 학습에서 관찰하는 것은 미래 텍스트의 예측 개선이다. 읽기 좋은 rationale과 실제로 유용한 rationale이 같을지도 검증할 문제다. 사람에게 설득력 있어 보인다는 이유만으로 높은 보상을 주는 구조는 아니다.
모든 위치에서 시도하면 비용이 폭발한다
긴 문장 뒤에서 한 번 보조 텍스트를 만드는 것과, 문장 안의 모든 토큰 위치에서 각각 만드는 것은 다르다. 순진하게 구현하면 위치마다 별도의 모델 실행을 반복해야 한다. 원문은 어텐션 마스크와 캐시를 구성해 여러 위치의 보조 경로를 병렬로 생성한다.
각 경로는 해당 위치까지의 원문과 자신의 앞선 thought 토큰을 볼 수 있다. 다른 위치의 보조 경로를 마음대로 읽거나 아직 관찰하지 않은 원문 미래를 읽어서는 안 된다. 서로 다른 문맥 접두부에서 시작하는 분기들을 동시에 진행하는 셈이다. 보조 텍스트의 다음 한 토큰을 여러 분기에서 함께 계산하고 이를 반복한다. 원문 병렬 생성
표 하나의 각 행에서 작은 글을 따로 이어 쓰는 장면을 떠올리면 좋다. 행들은 서로 다른 앞 문맥을 갖지만 같은 모델을 사용한다. 행 사이에서 정보를 섞으면 미래를 몰래 보는 문제가 생길 수 있다. 병렬화는 단순히 배치를 늘리는 일이 아니라 어떤 정보에 접근할 수 있는지 유지하는 일이기도 하다.
생각의 시작과 끝에는 학습 가능한 특별 토큰을 사용한다. 기반 모델이 원래 이런 보조 경로를 학습하지 않았으므로, 일반 텍스트 생성과 구별할 신호가 필요하다. 이 장치와 추가 학습이 필요한 만큼, 일반 모델에 '생각한 뒤 말해'라는 프롬프트를 넣는 것만으로 Quiet-STaR를 구현했다고 볼 수 없다.
도움이 되는 경우만 반영하려면
처음 생성하는 보조 텍스트는 원래 모델의 예측을 해칠 수도 있다. 그래서 원문은 thought 없이 얻은 기본 예측과 thought 뒤에서 얻은 예측을 결합하는 mixing head를 학습한다. 두 경로의 정보를 이용해 예측 점수, 즉 로짓을 어느 정도 섞을지 정한다. 쓸모없는 보조 경로에 초기부터 모든 예측을 맡기지 않도록 하는 장치다.
여기에는 흥미로운 시간 순서가 있다. 보조 텍스트를 이미 생성한 뒤 얼마나 사용할지 판단한다. 따라서 mixing head가 있다고 추가 계산 자체를 시작 전에 건너뛰는 것은 아니다. 원문은 생각을 언제 시작하고 끝낼지 동적으로 정하는 기능이 현재 구현에서 지원되지 않는다고 한계로 적는다. 유용성 판단과 계산 예산 절약은 연결되지만 같은 문제는 아니다.
나는 이 구분이 실제 추론 시스템을 생각할 때 중요하다고 본다. 쉬운 토큰에서 추가 계산이 필요 없다는 관찰이 있다면, 다음 목표는 그 사실을 생성 전에 알아채는 것이다. 이미 긴 보조 텍스트를 쓴 다음 무시한다면 품질은 보호할 수 있어도 시간은 되돌릴 수 없다. 이는 논문이 해결한 범위에서 출발하는 실무적 질문이다.
예쁜 설명보다 이어지는 텍스트의 확률
rationale은 샘플링된 이산 토큰이므로, 그 선택을 지나 일반적인 미분만으로 학습하기 어렵다. 원문은 REINFORCE를 사용한다. 같은 위치에서 여러 rationale을 만들고, 실제 미래 텍스트에 대한 예측이 다른 후보들의 평균보다 얼마나 유용한지를 학습 신호로 삼는다. 학습 안정성을 위해 음수 보상 항을 제외하는 선택도 설명하며 그로 인한 편향 가능성을 인정한다.
또 바로 다음 토큰 하나만 평가하면 '의미를 이해한 보조 텍스트'보다 당장의 표면적 표현에 과하게 맞출 수 있다. 그래서 여러 미래 토큰을 평가에 포함하고, 앞선 실제 토큰을 조건으로 주는 teacher forcing을 사용한다. 보조 텍스트를 생성할 때 미래를 보여 주는 것과, 생성이 끝난 뒤 알려진 실제 이어짐으로 유용성을 채점하는 일은 구별해야 한다. 원문 학습 목적
직접 만든 비교를 생각해 보자. 두 후보 중 하나는 문장의 결과를 설명하고, 다른 하나는 문맥과 관계없는 일반론을 반복한다. 첫 후보가 이어지는 실제 문장을 더 잘 예측하게 한다면 그 경로를 더 자주 만들도록 학습할 수 있다. 하지만 실제 글 자체가 오류를 포함한다면, 잘 예측한 것이 현실의 참을 이해했다는 증거가 되지는 않는다. 언어 모델링 신호가 줄 수 있는 것과 줄 수 없는 것을 구분해야 한다.
개선을 볼 때 계산도 함께 보기
원문은 Mistral 7B를 출발점으로 OpenWebMath와 C4 텍스트에서 추가 학습한다. 특정 평가 문제에 직접 미세조정하지 않은 상태에서 GSM8K와 CommonsenseQA 등의 개선을 보고한다. OpenWebMath 설정의 대표 결과는 GSM8K 5.9%에서 10.9%, CommonsenseQA 36.3%에서 47.2%다. 원문 평가
이 숫자는 가능성을 보여 주지만, 범용 추론 문제가 해결됐다는 규모의 성과로 읽기는 어렵다. 데이터와 학습 길이에 따른 차이가 있고, 원문의 실험 범위는 7B 모델에 한정된다. 추가 thought 토큰을 생성하는 비용도 크다. 저자들은 계산 비용을 고려한 비교를 별도로 제시한다.
내가 후속 실험을 설계한다면 같은 총 계산 예산에서 일반적인 추가 학습과 비교하고 싶다. 추론 시간도 맞춰야 한다. 품질이 올랐지만 답변을 만드는 시간이 많이 늘었다면, 서비스의 목적에 따라 가치가 달라진다. 예측하기 어려운 토큰에서 이득이 크다는 분석을 바탕으로 선택적인 추가 계산을 연구해 볼 이유도 있다.
Quiet-STaR가 남기는 매력은 '모델이 생각한다'는 표현 자체보다 학습 신호의 선택에 있다. 사람의 풀이를 그대로 복사하지 않아도, 보조 계산이 미래 텍스트를 얼마나 잘 예측하게 하는지로 학습할 수 있다. 어디에서 추가 계산이 도움이 되는지, 그 비용을 언제 지불할지라는 질문을 언어 모델링의 한가운데로 가져온다.
읽은 논문: Eric Zelikman 외, Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking, 2024년 사전 공개 논문. 논문 · 공식 코드