DTCA는 디퓨전 언어모델 강화학습에서 토큰별 기여 배분을 개선한다

무엇인가

이 논문은 디퓨전 언어모델(DLM)을 강화학습으로 튜닝할 때, 최종 정답 보상만으로는 어느 토큰 커밋이 이후 디노이징에 도움을 줬는지 알기 어렵다는 문제를 다룬다. DLM은 여러 디노이징 스텝에서 토큰을 조금씩 확정하고, 한 토큰의 확정이 남은 마스크 위치의 예측에 영향을 준다. 따라서 시퀀스 단위 보상을 그대로 쓰는 RL보다 더 세밀한 토큰 수준 크레딧 배분이 필요하다는 것이 논문의 출발점이다.

어떻게 동작하나

DTCA는 토큰 기여를 남은 마스크 위치들의 예측 신뢰도 개선으로 추정한다. 특정 토큰을 확정하기 전과 후에 대해 아직 마스크인 위치들의 최대 예측 확률 합을 비교하고, 증가분이 양수일 때만 ΔP로 기록한다. 부정 변화는 0으로 두고 남은 마스크 수로 나누지 않는다. 이 값을 RL 도중 모든 토큰에 대해 계산하면 추가 forward pass가 필요하고 롤아웃별 변동이 커서, RL 전에 오프라인 롤아웃으로 추정한다. 최대 생성 길이의 절반을 디노이징 예산으로 두고 각 스텝에서 두 토큰을 확정하며 각 토큰을 별도 forward로 평가한다. 이후 토큰의 품사(POS) 범주별로 ΔP를 평균해 m_c를 만들고, 범주별 가중치 w_c를 0.8~1.2 범위로 정규화한다. 이 범위 제한은 크레딧 스케일이 과해지지 않게 하기 위한 것이다. 논문은 GSM8K 롤아웃 분석에서 NOUN은 확정 비율이 높고 VERB는 초반 높다가 감소하며, VERB가 신뢰도 이득이 가장 크고 DET와 ADP는 상대적으로 작다고 보고한다.

무엇과 다른가

학습에서는 이 POS 기여 가중치를 두 곳에 쓴다. 먼저 각 디노이징 스텝의 중요도 q_t를 확정 토큰들의 w_c(1-p_i) 합으로 계산한다. p_i는 확정 당시 토큰 신뢰도다. 낮은 신뢰도 토큰과 기여가 큰 POS에 높은 점수를 주고, 점수 누적이 전체의 μ에 도달하는 최소 스텝 집합을 고른다. 실험에서는 전체의 70%를 덮을 때까지 고르고 선택 스텝 수는 2~32로 제한한다. 다음으로 선택된 스텝의 최적화 대상 토큰들에 대해 토큰의 POS 기여 가중치를 평균 1이 되도록 정규화해 a_t,i를 만들고 원래 스텝 어드밴티지 A_t에 곱해 A_t,i=a_t,i A_t를 만든다. 대상에는 해당 스텝 확정 토큰과 RLDF의 임계값 필터로 남긴 마스크 위치가 포함되고, POS는 완성된 롤아웃의 최종 예측 토큰 기준으로 붙인다. 기여가 큰 POS 토큰은 강한 학습 신호를, 작은 POS 토큰은 약한 신호를 받는다.

어떻게 쓰나

DTCA는 여기에 OPEFO에서 착안한 엔트로피 기반 크레딧 보정을 붙인다. 각 응답 b, 선택된 디노이징 상태 t, 대상 위치 i에서 정책 업데이트가 만들 엔트로피 변화 ΔH를 1차 근사로 추정한다. ΔH가 양수인 업데이트와 음수인 업데이트를 나눠 합을 비교해 λ_b를 계산하고 -0.95~0.95로 클리핑한다. λ_b가 양수이면 엔트로피 증가 업데이트를 키우고 감소 업데이트를 줄이며, 음수이면 반대로 조정한다. 이는 자동회귀 디코딩 스텝 대신 디노이징 상태와 남은 대상 위치에 맞게 변형한 것이다.

전제와 한계

실험은 LLaDA-8B-Instruct에서 추가 SFT 없이 진행됐다. RL은 MATH 학습 분할에서 RLDF와 같은 데이터 처리로 수행했고, POS 기여 테이블은 MATH와 GSM8K 학습 데이터의 오프라인 디노이징 롤아웃으로 미리 만들어 모델 파라미터는 업데이트하지 않았다. 평가는 MATH500과 GSM8K 최종 정답 정확도, 생성 예산 128·256·512에서 이뤄졌다. 비교 대상은 LLaDA-8B-Instruct, Uniform SFT, LLaDA-1.5, d1, ESPO, TraceRL, Coupled-GRPO, RLDF다. DTCA는 베이스 모델 대비 MATH500에서 예산 128/256/512에 각각 9.0/6.8/4.8포인트, GSM8K에서 5.3/2.1/2.5포인트 향상했다. 가장 작은 예산 128에서 이득이 가장 컸고 256과 512에서도 양수였다.

Ablation에서 두 구성 요소가 모두 기여했다. 256 예산에서 Contribution Credit을 빼면 MATH500 3.8포인트, GSM8K 2.4포인트 하락했고, Entropy Calibration을 빼면 각각 1.4포인트 하락했다. 512 예산에서는 Entropy Calibration 제거가 MATH500 4.2포인트, GSM8K 3.5포인트 하락을 냈다. 학습 동역학 분석에서 기여 기반 크레딧을 빼면 보상 궤적은 비슷하지만 gradient norm이 일관되게 더 컸고, 엔트로피 보정을 빼면 후반 KL divergence가 더 빨리 커지고 보상이 낮아졌다. 토큰 커밋 분석에서는 DTCA가 NUM 확정 비율을 줄이고 PUNCT를 늘리는 등 범주별 확정 분포를 바꿨고, VERB는 확정 비율 변화는 작지만 확정 시 신뢰도 상승과 엔트로피 감소가 컸다. 엔트로피 보정 분석에서는 POS 범주별 크레딧 변화와 엔트로피 변화가 항상 같은 방향은 아니었다.

개발자 관점에서는 디퓨전 LLM을 RL로 튜닝하면서 결과 보상만 쓰는 GRPO류에서 토큰별 학습 신호를 어떻게 더할지 고민할 때 참고할 만하다. 특히 여러 토큰이 한 스텝에 확정되는 마스크 디퓨전에서 스텝 선택과 토큰 가중치를 POS라는 저비용 그룹으로 안정화하는 설계, 그리고 오프라인 기여 테이블을 RL 전에 고정하는 운영 방식이 실무적으로 중요하다. 다만 재현하려면 POS 태거, 오프라인 롤아웃 수집, 추가 forward pass 비용, RLDF식 대상 토큰 필터와 엔트로피 보정 구현을 확인해야 한다.

원문은 별도의 한계 절을 두지 않았지만, 방법의 전제와 분석 범위는 명시되어 있다. 토큰 기여 추정은 RL 도중이 아니라 MATH와 GSM8K 학습 데이터의 오프라인 롤아웃으로 만든 공유 POS 테이블에 의존하고, 개별 토큰 효과를 별도 forward pass로 평가하므로 오프라인 비용과 데이터 분포 의존성이 남는다. POS 가중치는 0.8~1.2로 제한되며, 스텝 선택의 μ와 2~32 범위 같은 하이퍼파라미터가 있다. 엔트로피 보정은 1차 근사이고, 분석은 같은 토큰의 업데이트 직전·직후를 측정한 것이 아니라 훈련 전반의 추세를 비교한 것이다. 실험도 LLaDA-8B-Instruct와 수학 추론 벤치마크, 128~512 예산에 한정되어 있다.