🤖 machine learning

Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration

본 논문은 그래디언트 상호작용에 대한 최적화 분석에서 도출된 적응형 체계인 최적 계수 보정 (OCC) 을 제안하여 다중 토큰 예측과 검증 가능한 보상에 대한 강화 학습의 효과적인 공동 훈련을 가능하게 함으로써 이전의 성능 저하를 극복하고 수학 추론 벤치마크에서 우수한 결과를 달성한다.

원저자: Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

게시일 2026-05-28

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

이 논문은 간단한 언어와 창의적인 비유를 사용하여 설명합니다.

큰 그림: 두 명의 코치, 한 명의 선수

어려운 수학 문제를 해결하도록 거대하고 매우 똑똑한 로봇 (대규모 언어 모델) 을 훈련한다고 상상해 보세요.

주 코치 (RL): 이 코치는 로봇이 이기는 법을 가르칩니다. 로봇의 답변을 살펴보고, 정답에는 '엄지척'을, 오답에는 '엄지내림'을 줍니다. 로봇은 효과가 있는 행동을 반복하고 효과가 없는 행동을 멈추도록 학습합니다. 이를 **강화 학습 (Reinforcement Learning, RL)**이라고 합니다.
보조 코치 (MTP): 이 코치는 로봇이 다음 단어가 아닌, 앞으로 몇 단어를 말할지 예측하려는 특수한 모듈입니다. 이를 **다중 토큰 예측 (Multi-Token Prediction, MTP)**이라고 합니다. 이는 로봇이 문장을 미리 계획하도록 도와주는 코치와 같습니다.

문제:
이전에는 연구자들이 두 명의 코치가 동시에 로봇을 훈련시키려 했습니다. 하지만 무언가 잘못되었습니다. 로봇은 주 코치만 있을 때보다 더 나쁜 성적을 내기 시작했습니다.

이 때문에 대부분의 팀은 '승리' 훈련 단계 동안 보조 코치를 해고하기로 결정했습니다. 그들은 보조 코치가 지켜보게 했지만, 보조 코치의 조언이 로봇의 두뇌를 바꾸는 것은 허용하지 않았습니다. 그들은 보조 코치를 '분리 (detach)'시켰습니다.

질문:
이 논문의 저자들은 이렇게 물었습니다: 왜 두 명의 코치가 함께 있으면 로봇에게 해가 될까요? 훈련을 망치지 않으면서 두 코치가 함께 일할 수 있도록 고칠 수 있을까요?

진단: '밀고 당기기' 문제

저자들은 훈련의 수학을 분석하여 답을 찾았습니다. 그들이 깨달은 바는 보조 코치가 조언을 줄 때 로봇의 두뇌에 두 가지 상반된 힘이 작용한다는 것이었습니다.

도움이 되는 밀기 (상관관계): 때로는 보조 코치의 조언이 주 코치와 같은 방향을 가리킵니다. 이는 좋습니다! 마치 두 사람이 같은 방향으로 차를 밀고 있는 것과 같습니다.
짜증 나는 흔들림 (페널티): 때로는 보조 코치의 조언이 단순한 '노이즈'일 뿐입니다. 이는 로봇의 두뇌를 승리하는 데 도움이 되지 않는 무작위 방향으로 밀어붙입니다. 이는 운전하려는 차를 누군가 흔들고 있는 것과 같습니다.

이전 방법들이 실패한 이유:

방법 A (분리): 그들은 보조 코치의 밀기를 끄었습니다. 안전하지만, '도움이 되는 밀기'를 놓쳤습니다.
방법 B (교차 엔트로피): 그들은 보조 코치가 밀게 했지만, 보조 코치는 틀린 것까지 포함해 모든 단어를 균등하게 예측하도록 로봇을 가르치려 했습니다. 이로 인해 '밀기'는 거의 없고 '흔들림'이 매우 많이 발생했습니다. 로봇은 혼란을 겪고 더 나빠졌습니다.
방법 C (정책 손실): 그들은 보조 코치에게 주 코치와 같은 '승리' 규칙을 사용하도록 지시했습니다. 처음에는 이것이 아주 잘 작동했습니다! '밀기'가 강력했습니다. 하지만 로봇이 나아질수록 두 코치는 약간씩 이견을 보이기 시작했습니다. '흔들림'은 그대로였지만 '도움이 되는 밀기'는 약해졌습니다. 결국 흔들림이 지배하게 되었고 로봇의 성능은 추락했습니다.

비유:
당신이 줄타기 (RL 훈련) 를 하려고 한다고 상상해 보세요.

분리는 혼자 걷는 것입니다. 안전하지만 느립니다.
구식 공동 훈련은 술에 취한 친구가 손을 잡아주는 것입니다. 그들이 당신을 줄에서 끌어냅니다.
'상승과 하락' 문제는 처음에는 도와주지만, 당신이 지쳐갈수록 잘못된 방향으로 당신을 끌어당겨 넘어뜨리는 술에 취하지 않은 친구를 둔 것입니다.

해결책: '스마트 조정기 (OCC)'

저자들은 문제가 보조 코치 자체가 나쁘다는 것이 아니라, 그들의 목소리 볼륨이 고정된 수준에 갇혀 있었다는 것을 깨달았습니다.

훈련 초기: 보조 코치는 매우 도움이 됩니다. 그들의 볼륨을 올려야 합니다.
훈련 후기: 보조 코치는 방해가 되기 시작합니다. 그들의 볼륨을 내려야 합니다.

그들은 **최적 계수 보정 (Optimal Coefficient Calibration, OCC)**이라는 새로운 시스템을 만들었습니다.

작동 원리 (창의적 은유):
OCC 를 훈련 과정의 스마트 온도 조절기라고 생각하세요.

온도를 고정된 숫자로 설정하는 대신, 이 온도 조절기는 방을 끊임없이 점검합니다.
보조 코치가 현재 도움이 되고 있는지 해를 끼치고 있는지 확인하기 위해 '프록시 (빠르고 저렴한 센서)'를 사용합니다.
보조 코치가 올바른 방향으로 밀고 있다면, 온도 조절기는 볼륨을 올립니다.
보조 코치가 흔들리며 노이즈를 만들기 시작하면, 온도 조절기는 볼륨을 내립니다.

이 과정은 모든 것을 늦출 복잡한 수학을 계산하기 위해 멈추지 않고, 단계별로 자동으로 이루어집니다.

결과: 해피 엔딩

저자들은 이 새로운 '스마트 조정기'를 AIME 와 올림피아드와 같은 여섯 가지 어려운 수학 대회에서 테스트했습니다.

이전 방법들을 능가함: '스마트 조정기 (OCC)'로 훈련된 로봇은 '분리'된 주 코치만 훈련받은 로봇보다 더 많은 문제를 해결했습니다.
추락을 수정함: 처음에는 강력하다가 실패한 '정책 손실' 방법과 달리, OCC 방법은 첫 단계부터 마지막까지 강력하게 유지되었습니다.
빠름: '스마트 조정기'는 훈련 속도를 늦추지 않았습니다. 무거운 계산을 대신하여 볼륨을 확인하는 영리한 단축키를 사용했기 때문에 이전의 '분리' 방법만큼이나 빨랐습니다.

요약

이 논문은 로봇을 주 코치와 보조 코치와 함께 훈련시킬 수 있음을 증명하지만, 단순히 고정된 설정에 맡겨둘 수는 없다는 것을 보여줍니다. 훈련 과정을 듣고 보조 코치의 영향을 실시간으로 조정하는 동적 시스템이 필요합니다. 이렇게 할 때 로봇은 그 어느 때보다 더 잘, 더 빠르게 학습합니다.

기술 요약: 최적 계수 보정을 통한 강화학습의 멀티 토큰 예측 공동 학습

문제 제기

검증 가능한 보상에 기반한 강화학습 (RLVR) 은 대규모 언어 모델 (LLM) 의 추론 능력을 향상시키기 위한 표준 패러다임이 되었습니다. 동시에, 사전 학습 중에 미래의 여러 토큰을 예측하도록 훈련된 보조 목적 함수인 멀티 토큰 예측 (MTP) 은 샘플 효율성 측면에서 상당한 이점을 보여주었으며, 추측적 디코딩을 위한 드래프트 모델로 활용되고 있습니다.

직관적으로, 사후 학습 (post-training) 중에 MTP 와 RL 목적 함수를 공동으로 최적화하면 모델이 추론 능력을 습득하는 동시에 멀티 토큰 감독의 혜택을 받을 수 있어야 합니다. 그러나 경험적 증거에 따르면, 공동 MTP-RL 학습은 종종 주 모델의 성능이 크게 저하되는 결과를 초래합니다. 따라서 현재 veRL 및 slime 과 같은 프레임워크에서의 관행은 RL 학습 중에 주 모델에서 MTP 기울기를 **분리 (detach)**하는 것입니다. 이 분리는 MTP 파라미터가 독립적인 사후 미세 조정을 통해서만 적응하도록 강제하여, 동시 멀티 토큰 감독으로부터 얻을 수 있는 잠재적 이득을 포기하게 만듭니다. 본 논문이 다루는 핵심 문제는 다음과 같습니다: 왜 공동 MTP-RL 학습이 실패하며, 효과적인 공동 학습을 가능하게 하는 원칙적인 전략을 설계할 수 있는가?

방법론

이론적 분석

저자들은 최적화 관점에서 이 문제에 접근하여 RL 목적 함수 $J(\theta)$ 에 대한 MTP 의 단계별 효과를 분석합니다. $L$ -스무스 (L-smooth) 목적 함수를 가정할 때, 단계별 개선에 대한 하한을 유도합니다. 이 분해는 MTP 의 효과가 두 가지 경쟁하는 항에 의해 지배됨을 보여줍니다:

1 차 상관 항: RL 기울기 ( $g_{RL}$ ) 와 MTP 기울기 ( $g_{MTP}$ ) 사이의 방향적 일치 (내적) 에 의해 결정됩니다.
2 차 섭동 패널티: MTP 기울기의 크기가 유발하는 섭동을 반영하는 음이 아닌 항입니다.

성능 향상은 상관 항이 패널티 항을 능가할 때만 발생합니다.

기존 영역 분석

이 분해를 사용하여 본 논문은 세 가지 기존 학습 영역을 통합하고 설명합니다:

분리 (Detach): 주 모델에 대해 $g_{MTP} = 0$ 입니다. $\Delta_{MTP} = 0$ 입니다. 이는 안전하지만 잠재적 이득을 포기합니다.
교차 엔트로피 (CE) 손실: RL 은 높은 보상을 받는 샘플에 가중치를 두는 반면 CE 는 모든 샘플을 동등하게 취급하므로, 상관 항은 기대값에서 거의 0 에 가깝습니다. 결과적으로 부정적인 2 차 패널티가 우세하여 성능 저하를 초래합니다.
정책 손실 (Policy Loss): MTP 에 대해 동일한 RL 목적 함수 형식을 사용합니다. 초기에는 기울기가 정렬되어 상관 항이 우세하여 이득을 얻습니다. 그러나 학습이 진행되고 모델이 손실 지형의 평탄한 영역에 진입함에 따라 상관 관계는 감소하는 반면 기울기 노름 (따라서 패널티) 은 유지됩니다. 이는 고정된 계수가 위상 전이를 추적하지 못하게 하는 "상승 후 하락" 성능 곡선을 초래합니다.

제안된 솔루션: 최적 계수 보정 (OCC)

위상 전이를 해결하기 위해 저자들은 온라인으로 최적 MTP 계수 ( $\lambda$ ) 를 추적하는 적응형 방식인 **최적 계수 보정 (Optimal Coefficient Calibration, OCC)**을 제안합니다.

폐형 최적해: 분석은 각 단계에서 최적의 $\lambda^*$ 에 대한 폐형 표현식을 도출하며, 이는 기울기 상관 ( $c$ ) 과 제곱 기울기 노름 ( $v^2$ ) 의 비율에 비례합니다.
로그 확률 프록시: 대규모 분산 시스템에서 $c$ $c$ 와 $v^2$ $v^{2}$ 에 대한 전체 모델 기울기를 계산하는 것은 과도하게 비용이 듭니다. OCC 는 작은 단계 업데이트 하에서의 로그 확률 변화 ( $\delta \approx \log \pi_\theta - \log \pi_{old}$ $δ \approx lo g π_{θ} - lo g π_{o l d}$ ) 에서 유도된 로그 확률 프록시를 활용합니다.
- 정렬 프록시 ( $\hat{c}_t$ ): 내적을 추정합니다.
- 분산 프록시 ( $\hat{v}^2_t$ ): 제곱 노름을 추정합니다.
적응형 업데이트: 계수는 $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$ 로 업데이트됩니다. 이를 통해 OCC 는 전체 기울기를 계산하지 않고도 학습 전반에 걸쳐 양의 이득을 유지하면서 MTP 의 가중치를 동적으로 조정할 수 있습니다.

주요 기여

이론적 통합: 본 논문은 1 차 상관과 2 차 섭동 간의 상호 작용에 기반하여 분리, 교차 엔트로피, 정책 손실 영역의 성공 또는 실패를 설명하는 이론적 분석을 제공합니다.
OCC 알고리즘: 로그 확률 프록시를 사용하여 온라인으로 MTP 계수를 보정하는 적응형 절차인 최적 계수 보정을 제안합니다. 이 방법은 무시할 수 있는 계산 오버헤드를 요구하며 전체 모델 기울기 계산을 피합니다.
경험적 검증: MiMo-7B, GLM-4.5-Air 등 여러 모델과 DAPO, GSPO 등 여러 알고리즘에 대한 경쟁 수준의 수학 추론 벤치마크를 통한 광범위한 실험은 OCC 가 분리 (Detach) 기준선을 일관되게 능가하거나 일치함을 보여줍니다.

실험 결과

저자들은 AIME24, AIME25, AMC, MATH, Minerva, OlympiadBench 등 여섯 가지 벤치마크에서 OCC 를 평가했습니다.

성능: 모든 벤치마크와 모델 규모 (7B 밀집 및 106B MoE) 에서 OCC 는 분리 기준선과 고정 계수 정책 손실보다 일관되게 우수한 성과를 거두었습니다.
- DAPO 를 사용한 MiMo-7B 에서 OCC 는 평균 정확도 **61.7%**를 달성하여 분리 (58.9%) 와 정책 손실 (57.4%) 을 능가했습니다.
- GLM-4.5-Air 에서 OCC 는 분리 (65.9%) 대비 **67.6%**의 평균 정확도를 달성했습니다.
학습 동역학: 정책 손실은 상승 후 하락 경향 (후기 단계에서 저하) 을 보인 반면, OCC 는 학습 전반에 걸쳐 안정적인 개선을 유지했습니다. CE 손실은 분리보다 일관되게 낮은 성과를 보였습니다.
효율성: OCC 는 무시할 수 있는 오버헤드를 도입했습니다. OCC 의 단계별 월시계 (wall-clock) 학습 시간은 8.07 초로 분리 (8.11 초) 와 유사한 반면, 전체 모델 역전파는 45.06 초가 소요되어 5.6 배의 속도 저하가 발생했습니다.
절대 실험 (Ablation): 민감도 분석은 사전 정의된 스케일링 인자 $\lambda_+$ 가 다양한 값 범위에서 견고함을 보여주었으며, $\lambda_+ = 1.0$ 이 가장 좋은 평균 성능을 제공함을 확인했습니다.

중요성과 주장

본 논문은 공동 MTP-RL 학습의 실패가 본질적인 것이 아니라, 학습 중 기울기 정렬의 동적 드리프트와 정적 계수 간의 불일치에서 비롯된다고 주장합니다. 최적화 지형의 이론적 분해를 제공하고 OCC 를 제안함으로써, 저자들은 MTP 를 RL 사후 학습에 안전하게 재도입할 수 있음을 입증합니다.

이 연구의 중요성은 모델이 RL 을 통해 추론 능력을 동시에 습득하면서도 성능을 희생하거나 과도한 계산 비용을 치르지 않고 멀티 토큰 표현을 활용할 수 있게 한다는 점에 있습니다. 이 작업은 비볼록 딥러닝 지형에서의 공동 학습을 위해 적응형 보정이 필수적임을 시사하며, 현재 산업 표준인 기울기 분리 대안에 대한 원칙적인 대안을 제시합니다.

한계점

저자들은 이론적 분석이 $L$ -스무스 가정에 의존하며, 이는 비볼록 지형에서 정확히 성립하지 않을 수 있어 정량적 지침보다는 정성적 지침을 제공한다고 지적합니다. 로그 확률 프록시는 1 차 근사이며 큰 학습률 하에서는 충실도가 손실될 수 있습니다. 또한, 실험은 검증 가능한 보상을 가진 수학 추론에 초점을 맞추었으며, 학습된 보상 모델을 가진 RLHF 나 개방형 생성과 같은 설정으로의 일반화는 아직 검증되지 않았습니다.