Compute Aligned Training: Optimizing for Test Time Inference

Ousherovitch 와 Tewari 의 논문 "Compute Aligned Training: Optimizing for Test Time Inference"에 대한 상세한 기술적 요약입니다.

1. 문제 제기: 훈련 - 추론 불일치

이 논문은 현대 대규모 언어 모델 (LLM) 개발의 근본적인 단절을 다룹니다: 훈련 목표는 단일 샘플 성능을 최적화하도록 설계된 반면, 배포는 점점 더 추론 시 계산 (search) 전략에 의존합니다.

• 현황: 지도 미세 조정 (SFT) 및 강화 학습 (RL) 과 같은 표준 사후 훈련 방법은 단일 정답 출력의 확률을 최적화합니다 (예: 교차 엔트로피 최소화 또는 단일 샘플의 기대 보상 최대화).
• 충돌: 현대 추론 전략 (예: Pass@N, 다수결 투표, Best-of-N) 은 여러 후보 솔루션을 생성하여 이를 집계하고 최선의 답을 찾습니다.
  • 예시 (Pass@N): 모델이 정답일 확률이 50% 에서 90% 로 증가한다고 가정하면, 성공이 이미 거의 확실한 Pass@100 지표에는 체감 효과가 있습니다. 그러나 표준 훈련은 여전히 100% 확률을 추구하여 "쉬운" 샘플에 대한 용량을 낭비합니다.
  • 예시 (다수결 투표): 모델은 경쟁자들과 대결하기 위해 다수 (예: 35% 확률) 만 있으면 됩니다. 표준 훈련은 $p \to 1.0$을 추진하여 압도적인 승리를 장려하지만, 이는 모델이 앙상블로 작용해야 할 때 성능을 저하시킬 수 있는 견고한 합의 (robust consensus) 를 약화시킵니다.
• 결과: 이러한 불일치는 정책 붕괴 (policy collapse) (다양성 감소) 를 초래하고 모델이 증가된 추론 시 계산 예산을 효과적으로 활용하는 것을 방해합니다.

2. 방법론: 계산 정렬 훈련 (CAT)

저자들은 훈련 목표를 추론 시 사용되는 특정 추론 전략과 정렬시키는 계산 정렬 훈련 (CAT) 프레임워크를 제안합니다.

핵심 개념화

저자들은 추론 전략을 사후 휴리스틱이 아닌, 기본 정책 $\pi_\theta$를 효과적인 추론 시 분포 $\tilde{\pi}_\theta$로 변환하는 미분 가능 연산자 ($T$) 로 개념화합니다:
$$\tilde{\pi}_\theta(y|x) = T(\pi_\theta(\cdot|x), \phi)$$
여기서 $\phi$는 샘플 예산 $N$과 같은 하이퍼파라미터를 나타냅니다.

통합 손실 함수

CAT 은 기본 정책 ($\pi_\theta$) 하에서 정답의 음의 로그 가능도를 최소화하는 대신, 효과적인 시스템 수준 정책 ($\tilde{\pi}_\theta$) 하에서 손실을 최소화합니다:
$$L_{CAT}(\theta) = -\mathbb{E}_{(x, y^*)} [\log \tilde{\pi}_\theta(y^*|x)]$$

그라디언트 스케일링 메커니즘

전체 출력 공간을 샘플링하지 않고도 이를 계산적으로 실현 가능하게 만들기 위해, 저자들은 대각선 그라디언트 근사를 유도했습니다. 그들은 CAT 목적 함수의 그라디언트가 가중치 인자 $w(p, \phi)$로 표준 그라디언트를 스케일링하여 근사할 수 있음을 보였습니다:
$$\nabla_\theta L_{CAT} \approx w(p, \phi) \cdot \nabla_\theta L_{CE}$$
여기서 $p$는 정답의 기본 확률이며, $w$는 추론 시 전략에 대한 $p$ 증가의 한계 효용을 나타냅니다.

유도된 주요 스케일링 인자:

1. Pass@N (SFT 및 RL):
   • $w(p, N) = \frac{N p (1-p)^{N-1}}{1 - (1-p)^N}$
   • 행동: 성공이 이미 보장된 쉬운 샘플 (높은 $p$) 에 대한 그라디언트를 억제하여, 한계 이득이 커버리지를 향상시키는 "어려운" 샘플에 학습을 집중시킵니다.
2. 다수결 투표 (SFT 및 RL):
   • $w(p, N, k) \approx \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}$
   • 행동: 그라디언트 질량을 합의 임계값 ($p \approx k/N$) 근처에 집중시킵니다. 이미 승리하거나 절망적으로 패배한 샘플은 무시하고 오직 "전환점"에만 초점을 맞춥니다.
3. Best-of-N (RL):
   • $w \propto N (P_{<y})^{N-1}$
   • 행동: "승자 독식" 역학을 생성합니다. 그라디언트는 배치 내의 가장 높은 보상을 받는 샘플만 강화하여, 모델이 안전한 평균으로 수렴하는 대신 분포의 고분산 - 고보상 꼬리를 탐색하도록 장려합니다.

이론적 보장

이 논문은 경쟁적 전략(경쟁자의 확률을 높이는 것이 성공에 해가 되는 경우) 에 대해 대각선 근사가 실제 그라디언트의 보수적인 과소평가임을 증명합니다. 이는 최적화기가 절대 잘못된 방향으로 이동하지 않도록 보장하여, 정책 붕괴에 대한 자연스러운 안정제 역할을 합니다.

3. 주요 기여

1. 통합 연산자 프레임워크: 추론 시 전략을 미분 가능 연산자로 공식화하여 SFT 및 RL 과 같은 임의의 추론 전략과 훈련을 정렬시키는 일반 메커니즘을 유도합니다.
2. Pass@N 을 넘어선 확장: 이전 작업 (예: Direct Coverage Optimization) 을 다수결 투표 및 Best-of-N으로 일반화하여 SFT 및 RL 패러다임 모두를 포괄합니다.
3. 모달리티 중립성: LLM 을 넘어 단백질 언어 모델 (PLMs) 에도 프레임워크가 적용 가능함을 입증하여, 복잡한 다중 모달 보상 풍경에서 성능 향상을 보였습니다.
4. 분산 분석: 그라디언트 분산 팽창에 대한 이론적 및 실증적 분석을 제공하며, 훈련을 안정화하기 위한 정규화 기법 (예: 로그 가중 추정량) 을 제안합니다.

4. 실험 결과

저자들은 MATH 벤치마크(Mistral-7B 사용) 와 단백질 생성(ProtGPT2 사용) 에서 CAT 를 평가했습니다.

A. 언어 모델 (MATH 벤치마크)

• Pass@N (SFT): $N=64$에 대해 CAT 로 훈련된 모델은 단일 샘플 정확도 (Pass@1 약 2.3% 감소) 를 희생했지만, 표준 SFT 대비 Pass@64 커버리지에서 7.8% 절대적 개선을 달성했습니다.
• 다수결 투표 (SFT): CAT 모델은 모든 추론 예산에서 베이스라인을 일관되게 능가했습니다. $N=16$에서 편향과 분산을 최적으로 균형 잡는 "적정점"이 관찰되었습니다.
• 강화 학습 (GRPO): 표준 RL 은 정책 붕괴로 인해 빠르게 포화되었습니다 (Pass@32 $\approx$ 35.8%). CAT 훈련 모델은 검색에 필요한 분산을 유지하여 Pass@32 에서 **40.0%**를 달성했습니다.
• 워밍업 효과: SFT 워밍업 단계에서 CAT 를 사용하면 이후 RL 성능이 크게 향상되었으며, 모델의 추론 경로를 풍부하게 함으로써 단일 시도 (Pass@1) 정확도까지 높였습니다.

B. 단백질 생성 (단백질 언어 모델)

• 무조건부 (죽음의 계곡): 다중 모달 보상 풍경 (국소 최적점 $h \approx 0.35$, 전역 최적점 $h \approx 0.75$, 그 사이에 "계곡"으로 분리됨) 에서 표준 RL 은 국소 최적점에 갇혔습니다. Best-of-N CAT는 높은 분포 분산을 유지함으로써 "계곡"을 성공적으로 "뛰어넘어" 전역 최대값을 찾았습니다.
• 조건부 생성: CAT 모델은 입력과 역소수성 (inverse hydrophobicity) 을 가진 시퀀스를 생성하도록 학습하여, 표준 RL 에 비해 관계형 작업과 비교할 때 우수한 확장성 및 정렬을 보여주었습니다.

5. 중요성 및 함의

• 패러다임 전환: 이 논문은 훈련 목표가 하류 추론 절차를 염두에 두고 설계되어야 한다고 주장합니다. 추론 시 계산이 모델 능력 확장의 주요 레버가 됨에 따라, 훈련은 단일 샘플이 아닌 집계된 결과를 최적화해야 합니다.
• 효율성: CAT 는 훈련 계산 비용을 증가시키지 않고 이러한 이점을 달성합니다. 단순히 추론 시 전략에 대한 샘플의 한계 효용에 따라 기존 그라디언트를 재가중할 뿐입니다.
• 안정성 vs 탐색: 이 프레임워크는 단일 샘플 정밀도와 분포 다양성 사이의 트레이드오프를 원칙적으로 처리할 수 있는 방법을 제공하여, 모델이 모드 탐색 행동으로 수렴하지 않고 Best-of-N 및 다수결 투표와 같은 검색 절차를 효과적으로 활용할 수 있게 합니다.
• 일반화 가능성: LLM 과 단백질 LM 모두에서 유효성을 입증함으로써, 이 논문은 "계산 정렬 훈련"이 추론 시 검색을 활용하는 모든 생성 모델에 대한 보편적 원칙임을 시사합니다.

결론적으로, 계산 정렬 훈련은 학습과 검색 사이의 간극을 메워, 모델이 단순히 하나의 답을 생성하는 데 "똑똑한" 것을 넘어, 검색을 위한 계산 예산이 주어졌을 때 최선의 답을 찾는 데 "똑똑하도록" 최적화되도록 보장합니다.