기술 요약: Thoughtbubbles

문제 정의

현재의 트랜스포머(Transformer) 아키텍처는 고정된 계산 예산과 작업 기억(working memory)을 가진다는 근본적인 한계에 직면해 있습니다. 연쇄 사고(Chain-of-Thought, CoT)와 같은 기법들은 모델이 복잡한 문제를 단계별로 분해할 수 있게 해주지만, 이는 직렬로 생성되는 자연어 토큰에 의존하므로 비효율적이며 사전 학습(pretraining) 단계에 적용할 수 없습니다. "일시 정지 토큰(pause tokens)"이나 추가적인 잔차 스트림(residual streams)을 삽입하는 대안적 접근 방식들은 대개 수동적인 설계가 필요하고, 모든 레이어에 균일하게 적용되며(적응성 제한), 표준 언어 모델링 이상의 명시적인 지도 신호에 의존합니다. 따라서 모델이 명시적인 단계별 지시 없이도 어려운 토큰에 대해 더 많이 "생각"할 수 있도록, 잠재 공간(latent space)에서 병렬 계산을 동적으로 할당할 수 있는 비지도 방식의 아키텍처가 필요합니다.

방법론: Thoughtbubbles

저자들은 Thoughtbubbles를 제안합니다. 이는 잠재 공간 내에서 병렬 적응형 계산을 네이티브하게 수행하는 트랜스포머 변형 모델입니다. 핵심 메커니즘은 순전파(forward pass) 중에 잔차 스트림을 동적으로 "포킹(forking, 갈래 나누기)"하거나 삭제하는 법을 학습하는 것입니다.

핵심 메커니즘

1. 포킹 및 잔차 스트림 (Forking and Residual Streams):

   • 모델은 트랜스포머 레이어 사이에서 "포킹" 연산을 도입합니다. 각 잔차 스트림 $x^{(k)}_i$에 대해, 모델은 "유지 점수(keep score)"와 "포크 점수(fork score)"를 계산합니다.
   • 이 점수들은 이전 레이어로부터 전파된 "누적 점수(cumulative score)"와 곱해져 해당 스트림의 중요도를 결정합니다.
   • 이러한 점수에 Top-k 선택을 적용하여 어떤 스트림을 유지하고 어떤 스트림을 포크(복제)할지 결정합니다. 만약 스트림이 포크되면, 새로운 잔차 스트림이 생성되어 해당 토큰을 위한 잠재 계산의 "버블(bubble)"을 형성합니다.
   • 예측을 위한 소스를 유지하기 위해, 가장 오른쪽(rightmost) 토큰은 항상 유지되도록 강제됩니다.

2. 점수 감쇠 (Score Attenuation):

   • 모델이 중요한 스트림에 높은 점수를 부여하도록 훈련하기 위해, 저자들은 누적 점수를 사용하여 어텐션 메커니즘과 잔차 업데이트를 감쇠(attenuate)시킵니다.
   • 구체적으로, 어텐션 로짓(attention logits)과 밸류 벡터(value vectors)가 점수에 의해 조절됩니다. 이는 모델이 중요하다고 판단한 스트림(높은 점수)에는 더 크게 의존하고, 제거하려는 스트림(낮은 점수)에는 덜 의존하도록 강제하며, 이를 통해 어려운 토큰이 자연스럽게 더 많은 계산 자원을 축적하게 되는 자기 강화 루프를 생성합니다.

3. 위치 임베딩 (Positional Embeddings):

   • 가변적인 포크 수를 처리하기 위해, 저자들은 회전 위치 임베딩(Rotary Position Embeddings, RoPE)을 수정합니다. 포크된 토큰에 대해 포크 수에 비례하는 "부분 회전(partial rotation)"을 적용하여, 동일한 토큰의 포크들이 위치 공간에서 서로 "가깝게" 유지되도록 합니다.

4. 출력 집계 (Output Aggregation):

   • 최종 레이어에서 모델은 각 잔차 스트림을 개별적으로 디코딩합니다.
   • 최종 출력 분포는 누적 점수에 의해 결정된 가중치를 가진 분포들의 가중 평균입니다.
   • 1.9B 파라미터 규모의 효율성을 위해, 최종 소프트맥스 투영(softmax projection)을 적용하기 전에 잔차 벡터 자체를 평균내는 더 저렴한 근사 방식을 사용합니다.

5. 추론 전략 (Inference Strategy):

   • 저자들은 자기회귀적 생성(autoregressive generation)을 위한 **동적 포킹(dynamic forking)**을 도입합니다. 추론 예산($\kappa$)은 입력 시퀀스 길이에 비례하여 스케일링되는데, 이는 짧은 시퀀스가 긴 시퀀스에 비해 과도하게 포킹되어 발생하는 분포 변화(distribution shift)를 방지하기 위함입니다.

주요 기여

1. 비지도 적응형 계산 (Unsupervised Adaptive Computation): Thoughtbubbles는 비지도 방식으로 동적 잠재 병렬 계산을 할당할 수 있는 최초의 아키텍처입니다. 이 모델은 추가적인 지도나 명시적인 추론 궤적 없이 오직 언어 모델링 손실만을 사용하여 표준 디코더 전용 언어 모델로서 훈련됩니다.
2. 성능 효율성: 이 방법은 1.9B 규모에서 단 절반의 훈련 토큰 예산만을 사용하면서도, 표준 디코더 및 비적응형 병렬 계산 베이스라인(예: 중복된 필러 토큰)보다 우수한 성능을 달성합니다.
3. 확장성 (Scalability): 이 접근 방식은 150M에서 1.9B 파라미터에 이르는 모델 크기 전반에 걸쳐 일관된 성능 향상을 보여줍니다. 특히, 319M 규모의 Thoughtbubbles 모델은 퍼플렉서티(perplexity) 측면에서 772M 규모의 베이스라인을 능가합니다.
4. 해석 가능성 (Interpretability): 모델은 불확실성(높은 사후 엔트로피)이 높은 영역에 계산을 할당하도록 자율적으로 학습하며, 이는 어려운 토큰이 더 많은 "생각"을 필요로 한다는 가설과 일치합니다.

결과

• 제로샷 평가 (Zero-Shot Evaluation): 다양한 제로샷 벤치마크(LAMBADA, HellaSwag, BLiMP, AI2-ARC, PIQA)에서 Thoughtbubbles는 파라미터 및 계산량이 매칭된 베이스라인들을 지속적으로 앞섭니다.
• GSM8K 추론: 1.9B Thoughtbubbles 모델은 베이스라인의 절반에 불과한 토큰 예산을 사용하여 GSM8K 수학 추론 데이터셋에서 경쟁력 있는 결과를 달성했습니다.
• 퍼플렉서티 (Perplexity): 이 모델은 모든 평가에서 가장 낮은 퍼플렉서티를 달치하며, 일부 경우에는 더 큰 베이스라인 모델들까지 능가합니다.
• 절제 연구 (Ablation Studies):
  • 어텐션 마스킹(감쇠)을 제거하면 성능이 크게 저하되며, 이는 포크/유지 결정을 위한 점수 기반 조절의 필요성을 확인시켜 줍니다.
  • 포크된 토큰의 로짓을 단순히 평균내는 것이 가장 오른쪽 토큰만 유지하는 것보다 성능이 좋으며, 이는 병렬 스트림의 유용성을 입증합니다.
  • 포킹을 모든 레이어로 확장하는 것은 미미한 이득만을 주며, 이는 초기 레이어에서의 포킹만으로도 대부분의 작업에 충분함을 시사합니다.

의의 및 주장

본 논문은 Thoughtbuilders가 **통합된 훈련 시 및 테스트 시 스케일링 동작(unified train-time and test-time scaling behaviors)**을 향한 중요한 진전이라고 주장합니다. CoT가 훈련 과정에서 단계별 지시에 대한 노출을 요구하는 것과 달리, Thoughtbubbles는 표준 사전 학습 단계에서 적응형 계산을 암시적으로 수행하도록 학습됩니다.

저자들은 이 방법이 트랜스포머의 **입력 적응성(input-adaptivity)**을 끌어올려, 특정 토큰에 대해 작업 기억과 계산 깊이를 동적으로 확장함으로써 복잡한 다단계 문제를 해결할 수 있게 한다고 상정합니다. 이 연구는 잠재 적응형 계산이 고정된 예산의 아키텍처와 명시적인 직렬 추론 궤적의 한계를 넘어, 언어 모델 사전 학습의 표준적이고 통합된 구성 요소가 되는 미래를 제시합니다.