Papers·5일 전
Texas A&M, ARLM을 DLM으로 변환하는 On-Policy Distillation — 15x~7,000x 적은 학습 토큰
Texas A&M DIVE Lab이 autoregressive 언어모델(ARLM)을 diffusion 언어모델(DLM)로 변환하는 On-Policy Distillation(OPD) 기법을 제안했습니다. 기존 방식은 causal attention을 bidirectional로 바꾸고 DLM 목적함수로 재학습했지만, 목적함수 차이로 인한 지식 손실과 학습-추론 불일치 문제가 있었습니다. OPDLM은 ARLM을 student로, 원본 ARLM을 teacher로 삼아 student가 생성한 trajectory에서 teacher의 logit을 증류함으로써 두 문제를 동시에 해결합니다. 결과적으로 15x에서 7,000x까지 적은 학습 토큰으로도 다양한 태스크에서 강력한 성능을 보였습니다. 단, bidirectional attention으로 전환해야 하므로 추론 시 전체 시퀀스를 한 번에 처리해야 한다는 점은 유의할 만합니다.
Texas A&M 팀이 ARLM을 DLM으로 변환할 때 발생하는 두 가지 분포 이동을 동시에 해결하는 On-Policy Distillation 방법을 제안했습니다.
핵심 결론
- 태스크 — ARLM을 DLM으로 변환 (ARLM-to-DLM transformation).
- 성능 — 15x~7,000x 적은 학습 토큰으로 원본 ARLM 대비 경쟁력 있는 성능 유지.
- 벤치 — 다양한 언어 이해 및 생성 태스크에서 검증.
방법
- 핵심 아이디어 — Student는 bidirectional attention을 가진 ARLM, teacher는 원본 frozen ARLM. Student가 생성한 on-policy trajectory에서 teacher의 logit을 distillation target으로 사용.
- 차별점 — 기존 방식은 random masking 기반 DLM loss를 사용해 학습-추론 불일치가 있었지만, OPDLM은 student의 실제 생성 궤적을 활용해 이 문제를 제거.
- 또한 teacher distillation을 통해 원본 ARLM의 지식을 보존하므로 목적함수 변경에 따른 지식 손실도 방지합니다.
한계·조건
- 구조 제약 — Bidirectional attention을 사용하므로 autoregressive 디코딩이 불가능하며, 전체 시퀀스를 한 번에 생성해야 함.
- 리소스 — Teacher 모델을 고정한 채 student를 학습하므로 추가 메모리와 연산이 필요.
- 코드 — 현재 공개 여부 불명 — 논문 내 실험 설정만 확인 가능.
편집자 한 줄
DLM을 ARLM의 post-training 단계로 위치시킨 관점이 흥미롭습니다. 다만 bidirectional attention으로의 전환이 실제 배포 환경에서 얼마나 실용적일지는 더 봐야 할 것 같네요.
- #diffusion-language-model
- #knowledge-distillation
- #on-policy
- #texas-a&m
DIVE Lab, Texas A&M University