기술 요약: 전향적 이산 확산 모델을 위한 학습된 릴레이 표현

1. 문제 정의: 마스킹 확산 모델의 "하드 리셋"

마스킹 확산 모델 (MDM) 은 반복적인 탈노이즈 과정을 통해 이산 시퀀스를 생성합니다. 완전히 마스킹된 시퀀스에서 시작하여 모델은 일정 비율의 위치를 반복적으로 마스킹 해제합니다. 그러나 표준 MDM 추론에는 중요한 비효율성이 존재합니다. 각 순전파 (forward pass) 후, 모델은 마스킹된 위치를 위해 생성된 풍부한 내부 연속 계산 (은닉 상태) 을 폐기합니다. 다음 단계는 새로 확정된 이산 토큰들로부터만 시작됩니다.

저자들은 이를 하드 리셋 문제라고 명명했습니다. MDM 은 생성 과정에서 여러 번의 순전파를 수행하지만, 중간 연속 계산을 재사용할 수 없기 때문에 모델은 단계 간 추론을 효과적으로 분산 (amortize) 할 수 없습니다. 이는 순환 계산과 대조적입니다. 순환 계산에서는 고정된 파라미터 모델을 단계에 걸쳐 풀어서 (unrolling) 적용했을 때 모델이 근사할 수 있는 함수 클래스가 확장되어 복잡한 추론 작업에서 성능이 향상되는 것으로 입증되었습니다. 여기서 다루는 핵심 질문은, 중간 잠재 상태를 폐기하지 않고 순차적 마스킹 해제 단계 간에 더 풍부한 정보를 전달할 수 있는 순환 계산을 MDM 이 지원하도록 하는 방법입니다.

2. 방법론: 학습된 릴레이 표현 (Relay)

제안된 솔루션인 Relay는 디코딩 단계 간에 전달되는 연속적이고 미분 가능한 잠재 상태를 도입하여, 모델이 "전향적 (forward-thinking)"이 되도록 합니다.

핵심 메커니즘

• 증강된 상태 궤적: 각 추론 단계 $k$에서 모델은 증강 상태 $s_k = (x_{t_k}, h_k)$를 유지합니다. 여기서 $x_{t_k}$는 현재 부분적으로 마스킹된 시퀀스이고, $h_k$는 학습된 릴레이 상태 (연속 잠재 벡터) 입니다.
• 아키텍처: 모델은 백본 $f_\theta$, 토큰 임베딩 $Emb_\theta$, 언임베딩 헤드 $UnEmb_\theta$, 그리고 특정 릴레이 모듈 $R_\theta$로 구성됩니다.
  • 다음 은닉 상태는 다음과 같이 계산됩니다: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$.
  • 릴레이 상태 $h_k$는 임베딩된 토큰에 투영되어 추가된 후 백본에 입력됩니다.
  • 출력 로짓 $\ell_k$는 $UnEmb_\theta$를 통해 $h_{k+1}$에서 유도됩니다.
• Truncated BPTT 를 통한 학습: 독립적인 단계로 학습되는 표준 MDM 과 달리, Relay 는 **Truncated Backpropagation Through Time (BPTT)**을 사용하여 학습됩니다.
  • 모델은 완전히 마스킹된 시퀀스에서 시작하여 $K$단계 (예: $K=2$) 의 롤아웃을 수행합니다.
  • 롤아웃 동안 위치들은 정책 (예: 신뢰도 임계값) 에 따라 마스킹 해제되며, 확정된 토큰들은 롤아웃 중 오류 누적을 방지하기 위해 '교사 강제 (teacher-forced)'되어 정답 값으로 고정됩니다.
  • 손실은 $K$단계에 걸친 교차 엔트로피 손실의 합입니다.
  • 기울기는 릴레이 경로 ($h_k \to h_{k+1}$) 를 통해 역전파되어, 모델이 현재 단계 예측뿐만 아니라 미래 단계에 가장 유용한 정보를 $h_k$에 인코딩하도록 학습하게 합니다.

추론

추론 시에는 마스킹 해제 일정이나 샘플링 정책 측면에서 프로세스가 크게 변하지 않습니다. 유일한 추가 사항은 모델이 확정된 토큰과 함께 릴레이 상태 $h_k$를 앞으로 전달한다는 점입니다. 릴레이 상태는 아직 이산 토큰에 확정되지 않은 정보를 위한 미분 가능한 메모리 채널 역할을 합니다.

3. 주요 기여

1. 릴레이 프레임워크: Truncated BPTT 를 통해 학습된 학습된 잠재 릴레이를 디코딩 단계 간에 전달함으로써 MDM 에 순환 계산을 통합하는 일반적인 방법입니다. 이는 아키텍처에 구애받지 않으며, 블록 확산 및 KV 캐싱과 같은 기존 MDM 가속화 기술과 호환됩니다.
2. 확장성 및 적응: **Fast-dLLM v2 (1.5B)**의 전체 파라미터 적응을 통해 대규모 언어 모델 (LLM) 규모에서 Relay 의 유효성을 검증했습니다. 이 방법은 최소한의 추가 학습 FLOP 로 최첨단 이산 언어 모델 (DLM) 을 성공적으로 적응시켰습니다.
3. 성능 - 지연 시간 개선: Relay 가 성능 - 지연 시간 파레토 프론티어를 개선하여, 표준 지도 미세 조정 (SFT) 보다 코딩 작업에서 더 우수한 성능을 내면서도 추론 지연 시간을 크게 단축한다는 실증적 증거를 제시했습니다.
4. 디자인 공간 분석: 롤아웃 절차, 릴레이 채널, BPTT 학습 목표의 기여도를 분리하여 분석한 난이도 높은 스도쿠 기반 계획 작업에 대한 광범위한 절사 연구 (ablation studies) 를 수행했습니다.

4. 실험 결과

스도쿠 계획 작업

저자들은 Sudoku-Extreme 데이터셋을 사용하여 네 가지 학습 목표를 비교했습니다:

• MLM: 표준 균일 마스킹 확산 (단일 단계, 릴레이 없음).
• Rollout: 교사 강제와 함께 수행된 다단계 롤아웃이지만 릴레이 채널이 없음 ($R_\theta \equiv 0$).
• Relay (sg): 릴레이 경로를 포함하지만 그라디언트 정지 (stop-gradients) 를 사용함 (릴레이를 통한 BPTT 없음).
• Relay: 릴레이를 통한 $K=2$ BPTT 를 포함한 전체 방법.

결과:

• 정확도-NFE 프론티어: Relay 는 정확도와 순전파 평가 횟수 (NFE) 간의 가장 우수한 균형을 달성했습니다.
• 구성 요소의 영향: MLM 에서 Rollout 으로 이동하면 성능이 향상되었습니다. 릴레이 채널을 추가 (Relay sg) 하면 상당한 도약이 있었습니다. 마지막으로 BPTT 를 활성화 (Relay) 하면 최상의 결과가 나왔는데, 이는 모델이 더 공격적인 마스킹 해제 하에서도 부분 보드의 자기 일관성을 유지하도록 학습했음을 보여줍니다.
• 효율성: 고정된 신뢰도 임계값에서 Relay 는 각 순전파 시 더 많은 셀을 확정하면서도 합법성 (행/열/박스 위반 없음) 을 유지하여, 기준선보다 더 적은 총 순전파 횟수로 동일한 정확도에 도달했습니다.

사전 학습 모델 적응 (Fast-dLLM v2)

저자들은 코드 및 수학 데이터셋의 혼합물에 Relay 를 사용하여 Fast-dLLM v2 (1.5B) 를 적응시켰습니다.

• 코딩 성능: HumanEval 및 MBPP 벤치마크에서 Relay 는 표준 SFT 를 능가했습니다.
  • HumanEval: Relay 는 **42.1%**의 정확도를 달성했습니다 (SFT 는 38.4%). NFE 는 88.3회였습니다 (SFT 는 130.7 회).
  • MBPP: Relay 는 **46.6%**의 정확도를 달성했습니다 (SFT 는 43.9%). NFE 는 78.8회였습니다 (SFT 는 84.8 회).
• 지연 시간: Relay 는 정확도를 향상시키면서 표준 SFT 대비 최대 **32%**까지 추론 지연 시간을 단축했습니다.
• 메모리 오버헤드: A100 GPU 에서의 프로파일링 결과, Relay 는 두 번째 순전파 중 라이브 활성화 메모리를 증가시키지만, 릴레이 상태 저장소가 아닌 어휘 투영 역전파가 병목 현상이기 때문에 피크 메모리 사용량은 바닐라 SFT 와 거의 동일하게 유지되었습니다 (약 20.1 GiB 대 21.2 GiB).

5. 중요성 및 주장

이 논문은 Relay가 MDM 이 중간 연속 계산을 누적할 수 있게 함으로써 반복적 비자기회귀 생성의 본질을 근본적으로 변화시킨다고 주장합니다.

• 전향적 사고: 이 방법은 각 단계를 새로운 예측 문제로 취급하는 대신, 모델이 미래의 마스킹 해제 결정을 위해 중간 표현을 명시적으로 최적화할 수 있게 합니다.
• 파레토 프론티어 진전: 목표 정확도에 도달하는 데 필요한 순전파 횟수를 줄임으로써 Relay 는 성능 - 지연 시간 프론티어를 진전시켜, 복잡한 추론 및 생성 작업에 대해 이산 확산 모델을 더 효율적으로 만듭니다.
• 호환성: 이 접근법은 아키텍처 개편 없이 기존 최첨단 DLM 과 가속화 기술 (예: 블록 확산, KV 캐싱) 과 함께 작동하는 '드롭인 (drop-in)' 향상 기능으로 제시됩니다.

저자들은 Relay 가 비탐욕적이고 전향적인 생성을 향한 의미 있는 단계이며, 해석 가능성과 이미지, 분자 등 다른 모달리티로의 적용을 위해 더 탐구될 수 있는 잠재적 추론을 위한 연속적 기질을 제공한다고 결론지었습니다.