기술 요약: 빠른 확산 언어 모델 디코딩을 위한 지원적 토큰 공개 (AXON)

1. 문제 정의

이산 확산 대규모 언어 모델(dLLM)은 병렬 텍스트 생성을 가능하게 함으로써 자기회귀(autoregressive) 모델의 유망한 대안을 제공합니다. 모든 토큰을 순차적으로 생성하는 자기회귀 모델과 달리, dLLM은 완전히 마스킹된 시퀀스에서 시작하여 이를 반복적으로 노이즈를 제거(denoise)하며, 각 단계에서 모든 마스킹된 위치를 동시에 예측합니다. 이러한 양방향 정제 과정은 병렬 업데이트를 가능하게 하지만, 결정적인 **품질-지연 시간 트레이드오프(quality-latency trade-off)**를 초래합니다.

기존의 훈련 불필요(training-free) 병렬 디코딩 전략들은 단일 단계에서 어떤 마스킹된 토큰을 "안전하게" 공개(commit)할지 결정함으로써 이 문제를 해결하려고 시도합니다. 이러한 방법들은 주로 다음을 기반으로 합니다:

• 신뢰도 기반 기준: 높은 한계 확률(marginal probabilities)을 가진 토큰을 공개하여, 해당 토큰들이 신뢰할 수 있고 독립적이라고 가정합니다.
• 의존성 인식 기준: 결합도가 강한 토큰들을 동시에 확정하는 것을 피하기 위해 어텐션(attention)이나 분포 유사성을 사용합니다.

그러나 본 논문은 이러한 접근 방식에서 근본적인 병목 현상을 식별했습니다: 안전하지 않은 확정을 피하는 것이 남은 마스킹된 시퀀스의 디코딩을 용이하게 보장하지는 않는다는 점입니다. 안전하지 않은 토큰을 억제하더라도, 남은 불확실한 토큰들이 이를 신뢰성 있게 해결하는 데 필요한 충분한 가시적 문맥(visible context)을 갖추지 못할 수 있습니다. 이는 디코딩이 점점 더 보수적으로 변하거나 정체되는 "병목"으로 이어져, 더 많은 추가 디노이징 단계를 요구하게 만듭니다. 핵심 문제는 현재의 방법들이 무엇이 즉시 공개하기에 안전한가에만 집중할 뿐, 어떤 토큰을 공개했을 때 남은 불확실한 위치들에 가장 유용한 문맥을 제공할 것인가를 고려하지 않는다는 것입니다.

2. 방법론: AXON

저자들은 기존의 병렬 디코딩 전략을 보완하기 위해 설계된 훈련 불필요 모듈인 AXON을 제안합니다. AXON은 기본 디코더를 대체하는 것이 아니라, 현재 상태가 불충분한 문맥적 지원을 시사할 때 개입합니다. AXON은 디코딩 문제를 "안전한 확정 선택"에서 "지원적 앵커(supportive anchor) 선택"으로 재정의합니다.

핵심 메커니즘

AXON은 두 단계로 작동합니다:

1. 게이팅 메커니즘 (개입 시점 결정):
   AXON은 가벼운 게이트를 사용하여 기본 디코더가 어려움을 겪고 있는지 감지합니다. 이는 포워드 패스 중에 이미 계산된 신호(어텐션 맵 및 신뢰도 점수)를 바탕으로 두 가지 결핍 기준을 평가합니다:

   • 속도 결핍(Pace Deficit): 남은 디노이징 단계에 비해 토큰 공개 속도가 너무 느린지 감지합니다.
   • 문맥 결핍(Context Deficit): 남은 불확실한 토큰들이 현재 공개된 토큰들로부터 충분한 어텐션 지원을 받지 못하고 있는지 감지합니다.
     게이트는 이러한 결핍이 기본 디코더의 제안이 남은 마스킹된 시퀀스를 지원하기에 불충분하다고 나타낼 때만 활성화됩니다.

2. 앵커 선택 (무엇을 공개할 것인가):
   게이트가 활성화되면, AXON은 불확실한 위치들에 의해 강하게 어텐션되는 확신 있는 마스킹된 토큰들인 "앵커"의 작은 집합을 선택합니다. 선택 과정은 다음과 같습니다:

   • 스코어링: 후보 토큰 $j$는 불확실한 토큰 $i$를 지원하는 능력에 따라 점수가 매겨집니다. 점수 $w_{ij}$는 어텐션 가중치 $A_{ij}$, 대상의 불확실성 $(1 - c_i)$, 그리고 후보의 신뢰도 $c_j$를 결합합니다:
     $$w_{ij} = A_{ij} (1 - c_i) c_j$$
   • 커버리지 목적 함수: 중복되지 않는 앵커 서브셋을 선택하기 위해, AXON은 시설 위치(facility-location) 커버리지 목적 함수를 최대화합니다. 이 목적 함수는 각 불확실한 위치가 오직 선택된 최선의 앵커로부터만 크레딧을 받도록 보장하여, 선택된 집합의 다양성을 장려합니다:
     $$C(S) = \sum_{i \in U} \max_{j \in S} w_{ij}$$
     여기서 $S$는 선택된 앵커의 집합이고, $U$는 불확실한 위치의 집합입니다.
   • 최적화: 커버리지 함수가 단조 하부 가법성(monotone submodular)을 가지므로, AXON은 탐욕 알고리즘(greedy algorithm)을 사용하여 앵커를 선택하며, 이는 최적 집합에 대한 $(1 - 1/e)$ 근사치를 보장합니다.

통합

AXON은 모듈식입니다. AXON은 디코딩 상태를 모니터링하다가 게이트가 활성화되면, 선택된 앵커를 다음 디노이징 단계의 가시적 문맥으로 주입합니다. 이는 해결되지 않은 의존성을 해소하고, 기본 dLLM 백본을 수정하지 않고도 남은 토큰들의 수렴을 가속화합니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 설명합니다:

1. 빠른 디코딩에 대한 보완적 관점: 저자들은 단순히 "안전한" 토큰을 선택하는 것을 넘어, 불확실한 위치에 필요한 문맥을 제공하는 "정보력이 있는" 마스킹된 토큰을 식별하는 것으로 초점을 전환할 것을 제안합니다.
2. AXON 모듈: 기존 병렬 dLLM 디코더를 위한 훈련 불필요, 플러그 앤 플레이(plug-and-play) 모듈입니다. 이는 언더서포트(under-supported) 상태를 감지하기 위해 게이팅 메커니즘을 사용하고, 영향력이 큰 비중복 앵커를 선택하기 위해 하부 가법적 커버리지 목적 함수를 활용합니다.
3. 실증적 검증: 다양한 dLLM 백본(LLaDA, Dream), 추론 및 코드 생성 벤치마크(GSM8K, Minerva-Math, HumanEval, MBPP), 그리고 다양한 병렬 디코딩 전략(신뢰도 기반, 로컬리티 인식, 의존성 인식)에 걸친 광범한 평가를 수행했습니다.

4. 실험 결과

저자들은 세 가지 모델 제품군에 대해 강력한 베이스라인(Confidence, LocalLeap, DAWN)을 대상으로 AXON을 평가했습니다. 결과는 품질-지연 시간 트레이드오프의 개선을 보여줍니다:

• 효율성 (NFE): AXON은 생성 완료에 필요한 함수 호출 횟수(NFE)를 일관되게 줄였으며, 종종 상당한 차이로 줄였습니다 (예: 다양한 작업에서 NFE를 256에서 ~70-90으로 감소).
• 처리량 (TPS): 많은 구성에서 AXON은 초당 토큰 수(TPS)를 증가시켜, 더 빠른 실제 실행 시간(wall-clock inference)을 나타냈습니다.
• 정확도: AXON은 작업 정확도를 유지하거나 향상시킵니다. 특히, 몇몇 설정에서는 정확도와 처리량 모두에서 동시적인 이득을 보였습니다 (예: LLaDA-8B와 LocalLeap 조합에서 HumanEval 정확도가 3.05포인트 향상되는 동시에 TPS도 증가함).
• 절제 연구(Ablation Studies):
  • 커버리지 목적 함수: 시설 위치 목적 함수(distinct 불확실 위치의 커버리지를 최대화)가 가산적 커버리지(Graph-Cut)보다 우수한 성능을 보였으며, 이는 중복 앵커를 피하는 것이 효율성에 매우 중요하다는 것을 확인시켜 줍니다.
  • 게이트 가중치 ($\alpha$): 속도 결핍과 문맥 결핍을 균형 있게 조절하는 하이브리드 게이트가 필수적임이 입증되었습니다. 극단적인 설정(속도 또는 문맥에만 의존하는 경우)은 최적이 아니었으며, 이는 균형 잡힌 개입 전략의 필요성을 검증합니다.
  • 앵커 예산(Anchor Budget): 작은 고정 예산(예: $k=1$) 또는 적응형 커버리지 매칭이 최적인 것으로 나타났습니다. 더 큰 예산은 유의미한 이득 없이 노이즈를 유발했습니다.

5. 의의 및 주장

본 논문은 AXON이 현재의 병렬 디코딩이 가진 특정 한계인 '남은 토큰을 위한 문맥적 지원 부족'을 해결한다고 주장합니다. 앵커 선택을 커버리지 문제로 다룸으로써, AXON은 공개되는 정보가 단순히 "안전한" 것뿐만 아니라 후속 디노이징 단계를 위해 "유용한" 것임을 보장합니다.

저자들은 AXON이 백본 모델을 수정하거나 별도의 훈련 없이 이러한 개선을 달성한다는 점을 강조합니다. 이는 상태 모니터링과 전략적 앵커 선택을 통해 문맥적 지원의 구축을 능동적으로 관리함으로써, 효율적인 dLLM 디코딩이 단순한 커밋 규칙을 넘어설 수 있음을 시사합니다.

논문은 AXian이 어텐션과 신뢰도 신호에 의존하며, 이들이 실제 토큰 영향력의 대리 지표(proxy)라는 점을 언급하며 한계에 대해 겸허한 태도를 유지합니다. 만약 이러한 신호들이 제대로 보정되지 않는다면 그 이점이 제한될 수 있습니다. 또한, 현재는 텍고 전용 추론 및 코드 생성에 초점을 맞추고 있으며, 긴 문맥(long-context) 및 멀티모달 설정은 향후 과제로 남겨두었습니다.