디퓨전 언어 모델의 강화학습을 위한 토큰 신용 할당
원제 DTCA: Denoising Token Credit Assignment for Reinforcement Learning in Diffusion Language Models
추천 256댓글 4
Key Point
디퓨전 모델 기반 언어 모델에 강화학습을 효과적으로 적용하는 방식이 제시되었으며, 토큰별 신용 할당 문제를 해결함으로써 수학 추론 성능에서 실질적 개선을 보였다.
핵심 요약
- 디퓨전 언어 모델(DLM)은 반복 제거(denoising)로 텍스트를 생성하며 양방향 문맥 모델링과 병렬 토큰 정제가 가능하다.
- DLM에 강화학습을 적용하면 추론 성능이 개선되지만, 반복 제거 과정에서 토큰별 신용 할당이 어렵다는 문제가 있다.
- 제안된 DTCA(Denoising Token Credit Assignment) 프레임워크는 토큰 약속 후 남은 마스크 위치의 예측 신뢰도 개선으로 토큰 기여도를 추정한다.
- 기여도 가중치를 품사(POS) 카테고리별로 그룹화하여 정보량 많은 제거 시점을 선택하고 강화학습 훈련 중 토큰 단위 학습 신호를 할당한다.
- 정책 업데이트로 인한 토큰 엔트로피 변화에 따라 학습 신호를 추가로 보정한다.
- LLaDA-8B-Instruct에서 MATH500에서 기존 대비 최대 9.0점, GSM8K에서 5.3점 성능 향상을 달성했다.