기술 요약: 숨을 고르세요: 자기 주도적 시퀀스 생성을 위한 적응형 계산

문제 제기

기반 모델의 추론 시간 확장 (inference-time scaling) 은 주로 모델 깊이 증가 (예: 레이어 재귀) 나 체인 오브 씽킹 (Chain-of-Thought) 과 같은 추론 전략 활용에 초점을 맞춰 왔으나, 폭 기반 확장 (width-based scaling) 이 유망한 대안으로 부상하고 있습니다. 이 접근법은 입력 스트림에 특수한 <PAUSE> 토큰을 삽입하여 모델이 응답을 생성하기 전에 컨텍스트를 처리할 수 있도록 추가적인 순전파 (forward passes) 를 부여합니다.

그러나 "말하기 전에 생각하라 (Think Before You Speak, TBYS)" 프레임워크와 같은 기존 휴즈 토큰 활용 방법들은 데이터 로더에 의해 외부적으로 계산 단계 수가 결정된다는 치명적인 한계를 겪고 있습니다. TBYS 에서는 고정된 수의 <PAUSE> 토큰이 삽입되며, 최종 휴즈 단계가 될 때까지 모델의 출력은 무시됩니다. 이는 모델이 모든 중간 단계에서 매번 탐욕적으로 행동하도록 강요하며, 향후 휴즈를 예측하거나 자체 처리를 조절할 수 없게 만듭니다. 결과적으로 모델은 응답 준비 상태를 알리거나 현재 토큰의 복잡도에 따라 계산 자원을 적응적으로 할당할 수 없으며, 추가 단계를 동적 자원이 아닌 정적 장벽으로 취급하게 됩니다.

방법론: 숨을 고르세요 (Catch Your Breath, CYB)

저자들은 각 입력 토큰에 필요한 계산 단계 수를 모델이 자율적으로 결정할 수 있도록 하는 순차적 의사결정 문제로 구성된 지도 학습 손실 함수인 숨을 고르세요 (CYB) 를 제안합니다.

핵심 메커니즘

고정된 휴즈 시퀀스 대신, CYB 는 모델이 각 단계 $i$ 에서 다음 중 하나를 결정하도록 합니다:

1. 응답: 시퀀스에서 다음 토큰을 생성합니다.
2. 거부: 추가적인 <PAUSE> 단계를 요청하여 처리를 계속하기 위해 특수한 <DON'T KNOW> 토큰을 방출합니다.

이는 모델이 내부적 불확실성에 따라 토큰당 계산 깊이를 확장할 수 있도록 하여, 모델이 연속적으로 여러 번 거부를 선택할 수 있는 가변 지연 시스템을 생성합니다.

손실 함수

CYB 손실은 모델의 의사결정 과정에 대한 기대 할인 정확도에서 유도됩니다. $d_i$ 를 단계 $i$ 에서 모델이 거부할 확률로, $t_i$ 를 단계 $i$ 에서 올바른 목표 토큰을 선택할 확률로 정의합니다. 모델은 단계 $S$ 에서 출력을 생성합니다 (거부에 의한 자기 유도 또는 강제 정지에 의한 세계 유도).

손실은 기대 할인 정확도의 음의 로그 가능도로 정의됩니다:
$$\ell_{CYB}(t, d; \omega, \gamma) = -\log \mathbb{E}_{S \sim \text{Pr}(S|d, \omega)} [\gamma_S t_S]$$
여기서:

• $\omega$ 는 강제 종료를 허용하는 세계 유도 정지 시간의 분포를 나타냅니다.
• $\gamma$ 는 단계별 할인 인자로, 정확도가 향상되지 않는 한 응답을 지연하는 것에 대해 모델을 패널티하도록 합니다.
• 기대값은 해당 단계에서 정지할 확률로 가중치를 둔 모든 가능한 정지 단계에 대해 합산하여 계산됩니다.

학습 구현

효율성을 보장하기 위해 저자들은 병렬 학습 절차를 제안합니다. 학습 중에는 모든 실제 입력 토큰 뒤에 $W_{max} - 1$ 개의 휴즈 토큰이 삽입됩니다. 트랜스포머는 인과적 마스킹과 함께 이러한 단계를 병렬로 처리합니다. 손실은 모델이 일찍 응답하도록 선택하거나 (낮은 $d_i$) 세계가 정지를 강제하는 경우, 해당 특정 토큰에 대해 후속 휴즈 단계가 손실에 기여하지 않도록 계산됩니다. 이 접근법은 고정된 휴즈를 사용하는 표준 교차 엔트로피 학습과 동일한 계산 및 메모리 발자국을 유지합니다.

주요 기여

1. 원칙적인 손실 함수: TBYS 의 정적이고 외부적으로 부과된 지연과 대조적으로, 계산 단계를 스스로 결정하도록 트랜스포머를 학습시키는 CYB 손실의 도입.
2. 효율적인 병렬 학습: 적응형 계산을 가능하게 하면서도 표준 교차 엔트로피의 계산 및 메모리 효율성과 일치하는 학습 절차.
3. 가설 검증: 이 논문은 네 가지 가설을 검증합니다:
   • H1: CYB 는 읽기 타이밍을 스스로 결정하기 때문에 TBYS 보다 성능이 우수합니다.
   • H2: CYB 는 <DON'T KNOW> 확률이 지연을 통한 정확도 향상 가능성과 상관관계가 있는 잘 교정된 모델을 생성합니다.
   • H3: 강제 조기 정지를 통해 모델의 정지 시간 제어 권한을 제거하면 성능이 저하되어 자율성의 가치를 확인합니다.
   • H4: 속도 - 정확도 트레이드오프는 할인 인자 ($\gamma_0$) 를 통해 조작할 수 있어, 모델이 일부 정확도 손실을 감수하고 더 빠르게 응답하도록 할 수 있습니다.
4. 경험적 우위: TBYS 대비 추가적인 계산 비용 없이 사전 학습 및 미세 조정 시나리오 모두에서 CYB 가 퍼플렉시티를 감소시키고 하위 작업 정확도를 향상시킨다는 증명.

결과

Gemma1-2B 및 Gemma3-4B 모델을 C4 데이터셋 및 다양한 하위 작업 벤치마크 (HellaSwag, MMLU, GSM8K) 에서 사용하여 실험을 수행했습니다.

• 퍼플렉시티 감소: CYB 는 TBYS 를 일관되게 능가했습니다. 특히, 세 개의 휴즈 토큰을 가진 TBYS 모델보다 하나의 휴즈 토큰만 가진 CYB 모델이 더 낮은 퍼플렉시티를 달성했습니다. 사전 학습에서 3 개의 휴즈를 가진 CYB 는 TBYS 대비 3.6% 의 상대적 퍼플렉시티 개선을 보였으며, 미세 조정에서는 2.7% 를 보였습니다.
• 데이터 효율성: CYB 는 기준선 성능을 달성하는 데 베이스라인의 33% 만의 데이터가 필요했으며, 3 휴즈 TBYS 모델을 능가하는 데 66% 미만의 데이터만 필요했습니다. 이는 데이터 포화 아티팩트가 아닌 학습 효율성의 진정한 개선을 나타냅니다.
• 하위 작업 정확도: CYB 로 학습된 모델은 베이스라인 및 TBYS 모두보다 HellaSwag, MMLU, GSM8K 에서 더 높은 Top-1 정확도를 달성했습니다.
• 교정 및 행동: 분석 결과, <DON'T KNOW> 확률과 추가 단계로 인한 잠재적 정확도 향상 사이에 통계적으로 유의미한 양의 상관관계가 있음을 보여주었습니다. 정성적 분석에 따르면 휴즈 지속 시간은 컨텍스트에 민감합니다. 예를 들어, 모델은 복잡한 토큰 (예: "traditional") 에서는 더 오래 멈추지만 예측 가능한 시퀀스 (예: 숫자) 에서는 멈추지 않습니다.

중요성과 주장

이 논문은 계산 할당에 대한 자율성이 CYB 와 기존 폭 기반 확장 방법 사이의 핵심 차별점이라고 주장합니다. 모델이 언제 응답할 준비가 되었는지 결정하도록 허용함으로써, CYB 는 휴즈 토큰을 정적 장벽에서 정보를 통합하기 위한 동적 자원으로 변환합니다.

저자들은 CYB 가 주로 추론 시간 계산 비용 절감을 목적으로 하는 것이 아니라, 계산에 대한 유연하고 컨텍스트 의존적인 제어를 아키텍처에 제공하기 위해 고안되었다고 강조합니다. 이는 처리 요구 사항에 따라 고정 기간이 변하는 인간의 독서 행동을 반영합니다. 이 논문은 병렬화 가능성, 단순성, 형식적 표현력 덕분에 루프형 트랜스포머와 같은 깊이 기반 적응형 방법보다 구조적으로 우월한 대안으로 CYB 를 위치시키며, 계산 오버헤드를 증가시키지 않고 모델 성능을 향상시키는 추론 시간 확장의 새로운 패러다임을 제시합니다.