[논문] T3D: Trajectory Self-Distillation과 Direct Discriminative Optimization을 이용한 Few-Step Diffusion Language Models

Source: arXiv - 2602.12262v1

번역할 텍스트를 제공해 주시겠어요? 텍스트를 입력해 주시면 해당 내용을 한국어로 번역해 드리겠습니다.

개요

논문은 T3D라는 새로운 학습 프레임워크를 소개한다. 이 프레임워크는 확산 기반 대형 언어 모델(DLLM)이 단 몇 단계의 디코딩만으로도 고품질 텍스트를 생성하도록 만든다. 모델이 “trajectory self‑distillation”(경로 자체 증류)을 통해 스스로를 가르치게 하고, 역 KL 손실(Direct Discriminative Optimization이라고 함)을 사용함으로써, 저자들은 속도와 생성 충실도 사이의 트레이드오프를 크게 개선하여, 몇 단계만으로도 작동하는 확산 모델을 실용적인 사용에 한층 가깝게 만든다.

주요 기여

• Trajectory Self‑Distillation: 모델 자체의 다단계 생성 궤적을 압축된 “학생” 모델에 증류하여, 훨씬 적은 단계만으로 동일한 출력을 생성하도록 하는 새로운 방법.
• Direct Discriminative Optimization (DDO): 역 KL(모드‑탐색) 목표를 사용해 학생이 교사의 고확률 모드에 집중하도록 강제함으로써, 공격적인 단계 감소 시 흔히 발생하는 품질 손실을 감소시킴.
• Few‑Step Decoding Benchmarks: Wikitext‑103, PTB, 요약 등 표준 언어 생성 작업에 대한 광범위한 실험을 통해 DDPM 기반 소수 단계 디코더 및 기존 학습 체계와 같은 강력한 베이스라인 대비 일관된 향상을 입증.
• Open‑Source Release: 전체 코드, 사전 학습 체크포인트, 학습 스크립트를 공개하여 재현성 및 후속 활용을 용이하게 함.

방법론

1. Baseline Diffusion LLM – 저자들은 표준 diffusion 언어 모델을 시작점으로 삼으며, 이 모델은 T 단계(예: 50‑100) 동안 잠재 토큰 시퀀스를 반복적으로 디노이징합니다.
2. Collect Teacher Trajectories – 학습 중에 모델은 각 학습 예제에 대해 전체 길이의 diffusion 궤적(중간 노이즈 상태)을 생성합니다.
3. Self‑Distillation Loop
   • 학생 모델은 동일한 아키텍처로 초기화되지만 단지 K ≪ T 단계만 사용하여 교사의 최종 출력을 재현하도록 학습됩니다.
   • 학생은 교사의 중간 상태를 “소프트 타깃”으로 받아들여 훨씬 더 거친 노이즈 스케줄을 동일한 최종 결과에 매핑하는 방법을 학습합니다.
4. Direct Discriminative Optimization (DDO) – 일반적인 forward KL(교사의 모든 모드를 평균) 대신, DDO는 학생 분포와 교사의 고확률 모드 사이의 reverse KL을 최소화합니다. 이는 학생이 확률 질량을 얇게 퍼뜨리는 대신 가장 가능성이 높은 토큰 시퀀스를 찾도록 유도하며, 정제 단계가 몇 개만 남았을 때 특히 중요합니다.
5. Training Objective – 전체 손실은 표준 diffusion 재구성 손실과 DDO 항을 결합한 것으로, 학생이 교사 모드에 얼마나 적극적으로 집중할지를 조절하는 하이퍼파라미터로 균형을 맞춥니다.

전체 파이프라인은 끝‑끝으로 미분 가능하며, 기존 diffusion 기반 LLM에 아키텍처 변경 없이 그대로 적용할 수 있습니다.

결과 및 발견

• 품질 격차 감소: 단 5개의 diffusion 단계만 사용해도 T3D는 전체 단계 모델 대비 perplexity 격차의 약 60 %와 BLEU 격차의 약 70 %를 메워줍니다.
• 다양한 작업에 대한 견고성: 요약 (ROUGE‑L) 및 대화 생성 (Distinct‑n)에서도 유사한 개선이 관찰되어, 이 방법이 단순 언어 모델링을 넘어 일반화됨을 보여줍니다.
• 소거 실험: DDO를 제거하고 일반 forward KL을 사용할 경우 성능이 약 10‑15 % 상대적으로 저하되어, mode‑seeking distillation의 중요성을 확인할 수 있습니다.

전반적으로 T3D는 서브선형 속도 향상(3‑5배 빠름)을 제공하면서도 많은 다운스트림 애플리케이션에 충분히 허용 가능한 수준의 생성 품질을 유지합니다.

실용적 함의

• 실시간 챗봇 및 어시스턴트 – 몇 단계 확산 디코딩은 이전에 자동회귀 모델에서만 달성할 수 있었던 지연 제한(응답당 100 ms 미만)을 충족할 수 있습니다.
• 엣지 및 모바일 배포 – 노이즈 제거 단계 수가 감소함에 따라 연산량과 에너지 소비가 낮아져, 자원 제한 디바이스에서도 확산 LLM을 활용할 수 있게 됩니다.
• 병렬 토큰 생성 – 확산 단계가 전체 시퀀스를 동시에 처리하기 때문에 T3D는 DLLM 고유의 병렬성을 유지하며, 최신 GPU/TPU 배치 처리 파이프라인을 보다 효율적으로 활용할 수 있습니다.
• 파인튜닝 및 도메인 적응 – 자체 증류 프레임워크를 사전 학습된 확산 LLM 위에 적용하면, 개발자가 몇 단계 효율성을 유지하면서 모델을 특정 도메인에 빠르게 맞출 수 있습니다.

요약하면, T3D는 확산 언어 모델을 연구 단계의 호기심에서 생산 준비가 된 대안으로 전환시켜, 속도와 병렬성이 중요한 상황에서 활용할 수 있게 합니다.

제한 사항 및 향후 작업

• Full‑Step 우수성: T3D를 사용하더라도 최고의 품질은 여전히 원래 50‑step 디코더에서 나오므로, 절대적인 높은 충실도를 요구하는 미션‑크리티컬 작업은 여전히 풀‑스텝 또는 자동회귀 모델을 선호할 수 있습니다.
• 하이퍼파라미터 민감도: 재구성 손실과 DDO 항 사이의 균형은 신중한 튜닝이 필요합니다; 최적이 아닌 설정은 모드 붕괴나 다양성 저하를 초래할 수 있습니다.
• 초대형 모델에 대한 확장성: 실험은 약 1.3 B 파라미터 모델까지 진행되었습니다; 이 방식을 수십억 파라미터 규모의 LLM에 적용하면 새로운 안정성 문제가 나타날 수 있습니다.
• 향후 방향: 저자들은 입력마다 K를 변동시키는 적응형 스텝 스케줄 탐색, 제어 가능한 생성을 위한 classifier‑free guidance와 T3D 결합, 그리고 검색 기반 파이프라인과의 통합을 제안합니다.

If you’re interested in trying T3D yourself, the authors have released the code and pretrained checkpoints on GitHub (https://github.com/Tyrion58/T3D). Feel free to experiment, benchmark on your own workloads, and contribute back to the community!

저자

• Tunyu Zhang
• Xinxi Zhang
• Ligong Han
• Haizhou Shi
• Xiaoxiao He
• Zhuowei Li
• Hao Wang
• Kai Xu
• Akash Srivastava
• Hao Wang
• Vladimir Pavlovic
• Dimitris N. Metaxas

논문 정보

• arXiv ID: 2602.12262v1
• 분류: cs.CL, cs.LG
• 출판일: 2026년 2월 12일
• PDF: PDF 다운로드