기술 요약: 잔차 컨텍스트 확산 언어 모델 (Residual Context Diffusion Language Models)

문제 정의

확산 대규모 언어 모델(dLLM)은 병렬 토큰 디코딩을 가능하게 하여, 추론을 메모리 대역폭 제한 체제에서 연산 활용도 체제로 전환할 수 있는 잠재력을 가진, 자기회귀(Autoregressive, AR) 모델의 유망한 대안을 제시한다. 그러나 최신 블록 단위 dLLM은 AR 모델과 비교했을 때 상당한 정확도 격차를 겪고 있다. 이 격차는 추론 과정 중 사용되는 "리마스킹(remasking)" 메커니즘에 의해 더욱 심화된다. 즉, 각 디노이징 반복 단계에서 모델은 가장 확신이 높은 토큰만을 확정하고, 나머지는 정적인 마스크 토큰으로 재설정하여 폐기한다.

저자들은 이 과정이 폐기된(낮은 확신도를 가진) 토큰에 수행된 중간 연산을 사실상 낭비한다고 관찰했다. 토큰 회상(token recall) 분석을 통해 입증되었듯이, 이러한 중간 분포는 후속 디코딩 단계에 결정적인 역할을 하는 전역적 컨텍스트에 관한 구조적이고 의미론적인 정보를 포함하고 있다. 표준 dLLM은 이 정보를 버림으로써, 토큰당 더 높은 계산 비용을 지불함에도 불구하고 하위 최적의 정확도로 이어지게 한다.

방법론: 잔차 컨텍스트 확산 (Residual Context Diffusion, RCD)

본 논문은 폐기된 토큰의 잠재 표현을 "컨텍스트 잔차(contextual residuals)"로 취급하여 다음 디노이징 단계에서 모델에 다시 주입함으로써, 버려지는 연산을 재활용하는 프레임워크인 **잔차 컨텍스트 확산(RCD)**을 제안한다.

1. 엔트로피 가중치 잔차 주입 (Entropy-Weighted Residual Injection)

낮은 확신도의 토큰을 단순히 버리는 대신, RCD는 이들의 예측 확률 분포를 연속적인 임베딩 벡터(소프트 토큰)로 변환하여 다음 단계의 입력에 주입한다.

• 잔차 구성: 위치 $i$와 단계 $t$에서의 토큰에 대해, 잔차 정보 $\Delta_i^{(t)}$는 예측 확률 분포 $p_i^{(t)}$를 기반으로 한 어휘 임베딩의 가중 합으로 계산된다:
  $$\Delta_i^{(t)} = \sum_{j=1}^{V} p_{i,j}^{(t)} E_{j,:}$$
• 하이웨이 연결 (Highway Connection): 원래의 입력과 잔차 사이의 크기 불일치를 방지하기 위해, RCD는 현재 마스크된 토큰 임베딩과 이전 단계의 잔차를 보간하는 하이웨이 연결을 사용한다:
  $$\tilde{e}_i^{(t)} = (1 - \alpha_i^{(t-1)}) E(x_i^{(t)}) + \alpha_i^{(t-1)} \Delta_i^{(t-1)}$$
  (여기서 $x_i^{(t)} = [M]$은 마스크된 토큰을 의미한다).
• 엔트로피 기반 가중치 부여: 기여 가중치 $\alpha_i^{(t)}$는 예측된 토큰 분포의 정규화된 샤논 엔트로피(Shannon entropy)에 의해 동적으로 결정된다. 고엔트로피 토큰(낮은 확신도)은 시퀀스 정교화에 더 중요한 구조적 정보를 운반한다고 가설을 세웠기에 더 높은 가중치가 할당된다.

2. 분리된 2단계 학습 파이프라인 (Decoupled Two-Stage Training Pipeline)

RCD를 엔드 투 엔드로 학습시키는 것은 단계 $t$의 입력이 단계 $t-1$의 출력에 의존하는 재귀적 의존성 때문에 계산적으로 매우 부담스럽다(이는 RNN과 유사한 긴 언롤링 계산 그래프를 생성한다). 이러한 메모리 병목 현상을 우회하기 위해, 저자들은 두 단계의 학습 전략을 제안한다.

• 1단계 (참조 모델): 경량화된 사전 학습된 dLLM을 고정된 "참조 모델($M_{ref}$)"로 미세 조정한다. 이 모델은 마스크된 입력에 대해 고품질의 확률 분포와 엔트로피 가중치를 생성하는 역할을 한다.
• 2단계 (대상 모델 학습): "대상 모델($M_{target}$)"은 고정된 $M_{ref}$로부터 전달받은 신호를 사용하여 학습된다. 학습 과정에서 $M_{ref}$는 확률 분포와 엔트로피 가중치를 제공하여 잔차 벡터를 구성하며, 이는 $M_{target}$의 입력 임베딩에 주입된다. 이를 통해 잔차의 생성과 활용을 분리함으로써, 시간 역전파(backpropagation through time) 없이 표준적인 단일 단계 지도 학습이 가능해진다.

3. 추론 전략

추론 단계에서 모델은 스스로 자신의 잔차를 생성하는 자기 참조 루프로 전환된다. 학습(프록시) 단계와 추론(자기 생성) 단계 사이의 분포 차이를 메우기 위해 저자들은 다음을 도입한다:

• 온도 조절 엔트로피 (Temperature-Scaled Entropy): 스칼라 $T_{res}$가 엔트로피 계산에 사용되는 확률 분포의 "부드러움(softness)"을 조절하여, 모델의 확신도를 학습 분포에 맞게 교정한다.
• 초기화: 프로세스는 참조 모델의 초기 예측을 사용하여 "웜 스타트(warm-start)"하거나, 제로 벡터를 사용하여 "콜드 스타트(cold-start)"할 수 있다.

주요 기여

1. 잔차 컨텍스트 메커니즘: 버려지는 토큰의 표현을 재활용하여, 낭비되는 연산을 가이드 역할을 하는 컨텍스트 신호로 변환하는 새로운 모듈을 제안한다.
2. 엔트로피 기반 집계: 정규화된 샤논 엔트로피를 기반으로 잔차 정보의 가중치를 결정하는 원칙적인 방법을 제시하여, 불확실한 토큰이 컨텍스트 정교화에 더 많이 기여하도록 보장한다.
3. 효율적인 학습 파이프라인: 시간 역전파의 메모리 병목 현상을 피하는 분리된 2단계 학습 방식을 통해, 최소한의 데이터(약 3억 토큰)만으로 표준 dLLM을 RCD 패러다임으로 전환할 수 있게 한다.
4. 새로운 파레토 트레이드오프 (Pareto Trade-off): 디노이징 단계와 잔차 전달 사이의 새로운 조절 노브(knob)를 도입하여, 동일한 지연 시간에서 더 높은 정확도를 얻거나 동일한 정확도에서 훨씬 적은 단계를 사용할 수 있게 한다.

실험 결과

저자들은 LLaDA(양방향 글로벌 컨텍스트)와 SDAR(준-자기회귀 블록 단위)라는 두 가지 서로 다른 dLLM 제품군에 대해 RCD를 검증했다.

• 정확도 향상: RCD는 GSM8K, MATH500, MinervaMath를 포함한 벤치마크에서 프런티어 dLLM의 성능을 일관되게 4~11 퍼센트 포인트 개선했다.
• 복잡한 추론: 까다로운 AIME24/25 벤치마크에서 RCD는 베이스라인 정확도를 거의 두 배로 높였다 (예: SDAR-8B-b64는 AIME24에서 7.08%에서 18.75%로 향상).
• 효율성: RCD는 최소한의 추가 연산 오버헤드로 이러한 이득을 달성한다. 베이스라인의 최고 정확도에 도달하는 데 4~5배 적은 디노이징 단계만 필요하다.
• 데이터 효율성: RCD는 약 3억 토큰의 학습 데이터만으로 높은 추론 정확도에 도달하는 반면, 경쟁 관계에 있는 잠재 기반 방법들(예: Loopholing)은 동일한 예산 하에서 일관된 시퀀스를 생성하는 데 실패한다.
• 확장성: 이 방법은 4B에서 8B 파라미터 모델 및 다양한 블록 크기(32 및 64 토큰)에 걸쳐 견고하게 확장된다.

의의 및 주장

본 논문은 현재 dLLM의 주요 병목 현상이 모델 용량이나 학습 단계의 부족이 아니라, 리마스킹 전략으로 인한 정보 손실에 있다고 주장한다. 버려지는 토큰으로부터 "낭비되는" 연산을 재활용함으로써, RCD는 병렬 디코딩의 이점을 유지하면서도 자기회귀 모델의 추론 능력을 따라잡거나 능가할 수 있는 확산 모델의 잠재력을 끌어올린다.

저자들은 RCD를 고충실도 병렬 텍스트 생성을 위한 실용적이고 확장 가능한 솔루션으로 포지셔닝한다. 또한 이 방법이 백본 모델의 구조적 변경을 요구하지 않고 디코딩 및 학습 패러다임의 변화를 요구한다는 점을 강조하며, 이는 차세대 확산 언어 모델의 정확도와 효율성을 개선할 수 있는 실행 가능한 경로임을 시사한다.