TRIMS: Trajectory-Ranked Instruction Masked Supervision for Diffusion Language Models

1. 문제 정의 (Problem)

확산 언어 모델 (Diffusion Language Models, DLMs) 은 기존 자기회귀 (Autoregressive, AR) 모델과 달리 토큰을 병렬로 생성할 수 있어 저지연 (low-latency) 생성의 가능성을 제시합니다. 그러나 실제 적용 시 다음과 같은 한계가 존재합니다.

• 학습 - 추론 불일치 (Train-Inference Mismatch): 기존 DLM 은 균일한 무작위 마스킹 (uniform random masking) 전략으로 학습됩니다. 이는 토큰이 언제 해결되어야 하는지에 대한 명시적인 지도 신호가 부족함을 의미합니다.
• 비효율적인 디코딩 궤적: 추론 시 모델이 병렬 생성의 이점을 충분히 활용하지 못하거나, 어려운 토큰을 건너뛰거나 다음 토큰 예측 (next-token prediction) 방식으로 회귀하는 경우가 발생합니다. 이는 병렬성 (parallelism) 을 제한하고 전체적인 효율성을 떨어뜨립니다.
• 기존 방법의 비용: 기존에 제안된 궤적 (trajectory) 최적화 방법들은 대부분 DLM 자체에서 비용이 많이 드는 샘플링을 통해 궤적을 추출하고 증류 (distillation) 하는 방식을 사용합니다. 이는 학습 비용과 데이터 준비 비용을 급격히 증가시킵니다.

2. 제안 방법: TRIMS (Methodology)

저자들은 궤적 순위 기반 지시 마스킹 감독 (Trajectory-Ranked Instruction Masked Supervision, TRIMS) 을 제안합니다. 이는 DLM 학습에 궤적 정보를 주입하여 효율적인 디코딩 순서를 학습시키는 경량화된 파인튜닝 프레임워크입니다.

TRIMS 는 크게 세 단계로 구성됩니다 (그림 2 참조):

1. AR Teacher 를 통한 토큰 난이도 추정:

   • 비용이 많이 드는 DLM 샘플링 대신, 사전 학습된 자기회귀 (AR) 모델 (Teacher) 을 사용합니다.
   • 학습 데이터에 대해 Teacher-forcing 모드로 한 번의 순전파 (forward pass) 를 수행하여 각 토큰의 예측 확률을 구하고, 이를 바탕으로 토큰 난이도 점수 (Token Difficulty Score, $s_i = -\log p_T(y_i | x, y_{<i})$) 를 계산합니다.

2. 난이도 버킷화 (Difficulty Bucketing):

   • 연속적인 난이도 점수를 $K$개의 순차적 버킷 (bucket) 으로 이산화합니다.
   • 양분수 (quantile) 기반 임계값을 사용하여 각 버킷에 포함된 토큰 수를 균등하게 분배합니다.
   • 핵심 전략: 어려운 토큰 (높은 난이도 점수) 은 낮은 버킷 ID(초기 디코딩 단계) 에, 쉬운 토큰은 높은 버킷 ID(후기 단계) 에 할당합니다. 즉, "어려운 것부터 쉬운 것 (Hard-to-Easy)" 순서로 디코딩하도록 유도합니다.

3. 궤적 인식 마스킹 (Trajectory-aware Masking):

   • 학습 과정에서 무작위 마스킹 대신, 버킷 할당을 기반으로 마스킹 전략을 변경합니다.
   • 현재 샘플링된 버킷 임계값 $k$보다 낮은 ID(어려운 토큰) 를 가진 토큰은 이미 해결된 컨텍스트로 간주하고, 높은 ID(쉬운 토큰) 를 가진 토큰은 미래 토큰으로 마스킹합니다.
   • 이 과정에서 미래 토큰은 높은 확률 ($p_{future} \approx 0.95$) 로 마스킹되고, 컨텍스트 토큰은 낮은 확률 ($p_{context} \approx 0.05$) 로만 마스킹되어 모델이 어려운 토큰을 먼저 학습하도록 유도합니다.
   • 전체 학습 데이터의 10% 에만 이 전략을 적용하여 사전 학습 단계와의 불일치를 최소화하면서도 안정성을 유지합니다.

3. 주요 기여 (Key Contributions)

1. 경량화된 궤적 감독 프레임워크: DLM 의 학습 - 추론 불일치를 해결하기 위해 AR Teacher 의 토큰 난이도 신호를 활용하여 궤적 인식 감독을 도입했습니다.
2. 최소한의 변경으로 구현 가능한 마스킹 전략: 표준 Masked Diffusion Language Model (MDLM) 학습 파이프라인을 거의 변경하지 않고, "어려운 것부터 쉬운 것" 순서의 디코딩을 시뮬레이션하는 마스킹 전략을 제안했습니다.
3. 높은 효율성과 경쟁력 있는 성능: 기존 궤적 최적화 방법 (증류 기반) 에 비해 비용이 훨씬 낮습니다. DLM 샘플링이 필요 없으며, AR 모델의 단일 순전파와 1K 개의 데이터만으로도 증류 기반 방법 (93K 데이터, 수백 GPU 시간 소요) 과 경쟁력 있는 성능을 달성합니다.

4. 실험 결과 (Results)

LLaDA-Instruct 와 Dream-Instruct 두 가지 오픈소스 DLM 을 대상으로 수학 (GSM8K, MATH) 과 코딩 (HumanEval, MBPP) 벤치마크에서 평가했습니다.

• 정확도 - 병렬성 트레이드오프 개선:
  • TRIMS 는 표준 MDLM 학습 및 학습 불필요 가속화 (train-free) 기법들에 비해 정확도를 유지하면서 병렬성 (TPS, Tokens Per Step) 을 크게 향상시켰습니다.
  • 예: LLaDA 모델에서 GSM8K 기준 TPS 가 1.00 에서 6.26 으로 증가했으며, Dream 모델에서는 MBPP 에서 6.31 TPS 를 기록하여 학습 불필요 기법 대비 최대 3 배 이상의 병렬성을 달성했습니다.
• 비용 효율성:
  • 증류 기반 방법 (dParallel, d3LLM) 은 모델당 287 GPU 시간의 데이터 정제 비용과 93K 개의 샘플이 필요하지만, TRIMS 는 0.6 GPU 시간 (AR Teacher 한 번 실행) 과 1K 개의 샘플로 충분했습니다.
  • 이는 데이터 효율성과 계산 비용 측면에서 압도적인 우위를 보입니다.
• 애블레이션 연구:
  • 순서 효과: "어려운 것부터 쉬운 것 (Difficulty-descending)" 순서가 무작위나 쉬운 것부터 어려운 것 순서보다 전반적으로 더 좋은 성능을 보였습니다.
  • 버킷 크기: $K=8$이 가장 최적의 성능을 보였으나, $K$값에 따른 성능 차이는 크지 않아 하이퍼파라미터에 강건함을 보였습니다.
  • 난이도 지표: 코딩 작업에서는 NLL(음의 로그 가능도) 이 엔트로피보다 더 나은 정확도 - 병렬성 트레이드오프를 보였습니다.

5. 의의 및 결론 (Significance)

TRIMS 는 확산 언어 모델의 병렬 디코딩 잠재력을 실현하기 위한 실용적이고 확장 가능한 접근법을 제시합니다.

• 실용성: 고비용의 DLM 샘플링이나 복잡한 증류 과정 없이, 기존 AR 모델의 지식을 활용하여 DLM 의 학습 효율을 극대화합니다.
• 이론적 통찰: DLM 이 단순히 토큰을 예측하는 것을 넘어, 어떤 순서로 토큰을 해결할지 (decoding trajectory) 에 대한 명시적인 지도를 받으면 병렬성과 정확도 간의 균형을 크게 개선할 수 있음을 입증했습니다.
• 미래 전망: 이 프레임워크는 더 큰 규모의 모델로 확장 가능하며, 다양한 형태의 궤적 신호를 탐구하는 기초가 될 것으로 기대됩니다.

요약하자면, TRIMS 는 "어려운 토큰을 먼저 해결하는 순서" 를 학습하도록 유도함으로써 DLM 의 병렬 생성 능력을 극대화하고, 이를 위해 매우 낮은 비용으로 구현 가능한 혁신적인 학습 전략을 제안한 논문입니다.