분포 민감도로 과개입을 억제하는 최소 침습 언어모델 스티어링

무엇인가

테스트 시점 정렬(test-time alignment)은 모델 가중치를 건드리지 않고 사용자나 상황에 따라 달라지는 보상에 맞춰 생성을 조정하려는 시도다. 이 논문이 다루는 것은 그중에서도 pre-logit steering, 즉 트랜스포머 본체를 통과한 마지막 은닉 상태에 위치별 벡터 u_t를 더해 로짓을 W(h_t+u_t)+b로 바꾸는 방식이다. 파라미터를 고정한 채 보상만 최대화하면 잔차 상태가 도달 가능한 다양체에서 멀어지는 과개입(oversteering)이 발생해 일관성과 다양성이 무너지고, 불완전한 보상 모델을 exploit하는 출력이 선택되기 쉬워진다. 저자들은 이 문제를 보상 최대화와 참조 정책으로부터의 이탈 억제 사이의 균형 문제로 정식화한다.

어떻게 동작하나

기존의 등방성(isotropic) 페널티는 ‖u_t‖²를 벌할 뿐, 같은 크기의 스티어링이라도 토큰 확률을 얼마나 바꾸는지는 전혀 반영하지 못한다. 논문은 토큰 수준 KL의 2차 전개에서 비용 행렬을 유도한다. KL(π_u^t‖π_ref^t) = ½ u_t^T F_t(0) u_t + O(‖u_t‖³)이고, 여기서 F_t(u_t) = W^T(diag(π_u^t) − π_u^t π_u^t^T)W가 u_t에 대한 피셔 정보행렬이다. 이 이차형식은 u_t^T F_t(0) u_t = Var[(Wu_t)_i], 즉 로짓 섭동의 분산과 같다. 그래서 모든 로짓을 상수만큼 평행이동하는 개입은 비용이 0이고, 확률을 실제로 흔드는 방향에만 벌점이 붙는다. 접두사 의존성을 없애기 위해 참조 정책 접두사에 대해 평균한 F̄_t를 고정 비용 행렬로 쓰며, MISVO는 Σ_t = F̄_t로 둔다.

무엇과 다른가

이론적 기여는 근사의 오차를 정량화한 데 있다. 시퀀스 수준 KL 그래디언트는 현재 토큰 분포에 대한 직접 효과 g_t^an(u) = E[F_t(u_t)]u_t와, 미래 접두사 분포 변화에서 오는 스코어 함수 항 g_t^suffix로 정확히 분해된다. 세 가지 피셔 대리행렬(반복점 의존 F̃_t^u, 스티어링된 접두사에서 0점 평가 F̄_t^u, 참조 접두사 평균 F̄_t)의 차이는 각각 C₁‖u_t‖(C₁=3σ_max(W)³)와 C₂Σ_{τ<t}‖u_τ‖(C₂=√V σ_max(W)³)로 유계이며, 고정 생성 지평선 T에서 스코어 함수 항은 O(‖u‖²)다. 결과적으로 F̄_t u_t는 전체 KL 그래디언트와 1차 항까지 일치한다. 단 이는 u→0 근방의 국소적 진술이며, 큰 개입이나 최적화 궤적을 따라 누적된 오차를 보장하지는 않는다.

어떻게 쓰나

알고리즘은 단순하다. 각 반복에서 현재 정책으로 K개의 응답을 샘플링해 보상을 평가하고, leave-one-out 베이스라인 R̄_i로 분산을 줄인 보상 그래디언트 g_t^R = (1/K)Σ(R_i−R̄_i)W^T(e_{y_t}−p_t)를 계산한다. 업데이트는 d_t = g_t^R − λΣ_t u_t, u_t ← u_t + ηd_t다. 피셔 행렬은 u=0인 첫 스텝의 롤아웃에서 추정해 최적화 내내 재사용하므로 추가 생성이 필요 없다. 행렬을 명시적으로 만들지 않고 F_t u_t = W^T C_p v, C_p v = p⊙v − p(p^T v) 형태의 행렬-벡터 곱(O(Vd))으로 처리하는 matrix-free 구현을 쓰며, 명시적 구성 시 드는 O(KTVd² + NTd²) 비용과 대비된다.

전제와 한계

실험은 약 1B~14B 파라미터의 네 모델(LFM2.5-1.2B-Instruct, Gemma3-4B-IT, Llama-3-8B-Instruct, Phi-4 14B)에서 수행한다. 과제는 SHP(테스트셋에서 무작위로 뽑은 500개 프롬프트, Skywork-Reward-V2-Qwen3-0.6B 보상 모델)와 MBPP+(처음 120개 문제, 홀드아웃 단위 테스트 통과율을 [0,1] 보상으로 사용)다. 베이스라인은 동일 생성 예산의 Best-of-N(top-p)과 AISP이며, RE-Control은 선행 결과가 열등하고 계산 자원 한계로 제외했다. 하이퍼파라미터는 LFM2.5-1.2B·Gemma3-4B에 K=16, N=16, Llama-3-8B·Phi-4에 K=32, N=32, 학습률 0.1, λ=1이고, 3개 시드 평균±표준편차로 보고한다. 결과적으로 7개 모델-과제 설정 중 6개에서 최고 평균 보상을 기록했다. SHP에서는 4개 모델 중 3개에서 BoN과 AISP를 모두 앞섰고, Llama-3-8B만 AISP에 역전당했다(그래도 BoN은 상회). MBPP+에서는 세 모델 모두 1위이며 가장 작은 모델에서 이득이 가장 컸다. 다양성(∏(1−rep_n))과 일관성(SimCSE 프롬프트-응답 코사인 유사도)은 BoN에 근접했다. 다만 제공된 본문에는 표의 구체적 보상 수치가 실려 있지 않아 개선폭의 절대값은 확인할 수 없다.

별도로 5.3절에서는 참조 정책 접두사 위에서 토큰 수준 KL을 직접 측정한다(100개 프롬프트). MISVO는 이 비교에서 가장 높은 보상과 함께 AISP 및 정규화 없는 스티어링보다 낮은 참조 접두사 KL을 보였고, 직접 KL을 정규화 항으로 쓴 변형은 MISVO보다 보상이 낮고 KL이 높았다. 저자들은 이 지표가 참조 접두사에서의 이탈만 측정할 뿐, 스티어링이 유발한 접두사 방문 분포 변화 전체를 포착하지는 않는다고 명시한다.

개발자 관점에서 이 논문은 가중치를 다시 학습하지 않고 추론 시점에 정책을 보상 쪽으로 기울이되 분포를 망가뜨리지 않으려면 무엇을 벌해야 하는가에 대한 실용적 답을 준다. 등방성 L2 페널티 대신 로짓 공간의 피셔 이차형식을 쓰면, 같은 노름의 개입이라도 토큰 확률을 거의 건드리지 않는 방향(예: 전 로짓 상수 이동)은 자유롭게 두고 민감한 방향만 억제한다. 구현 부담도 크지 않다. LM head W만 있으면 되고 트랜스포머 본체 역전파가 필요 없으며, 피셔 추정은 첫 스텝 롤아웃을 재활용하므로 추가 생성 예산이 들지 않는다. 다만 보상 모델이 블랙박스여도 되는 대신 보상 신호의 품질이 결과를 좌우하고, λ·K·N 같은 하이퍼파라미터를 과제별로 맞춰야 한다는 점은 그대로 남는다.

한계는 저자들이 분명히 밝힌다. 근사가 국소적이므로 큰 개입 후의 시퀀스 수준 KL이 작다는 보장이 없고, 편차 측정 지표가 참조 정책 접두사를 쓰기 때문에 시퀀스 수준 KL을 직접 재지 못한다. SHP에서는 최적화와 주 평가에 같은 보상 모델을 사용해 독립적인 선호 품질에 대한 결론이 제한된다. 또한 반복 최적화 오버헤드가 있고, 저자들 구현 기준으로 베이스라인보다 피크 메모리가 높다. 고처리량 추론 시스템으로의 통합과 더 메모리 효율적인 피셔 곱은 향후 과제로 남겨졌다.