기술 요약: 이중 초점 확산 언어 모델 (Bifocal Diffusion Language Models)

1. 문제 정의

이산 확산 언어 모델(Discrete Diffusion Language Models, dLLMs)은 여러 개의 마스킹된 토큰을 병렬로 복구함으로써 오토레그레시브(Autoregressive, AR) 생성의 대안으로서 유망한 가능성을 제시하며, 잠재적으로 3~8배의 실시간(wall-clock) 속도 향상을 제공할 수 있습니다. 그러나 dLLM은 문맥 접근 방식과 추론 효율성 측면에서 근본적인 구조적 딜레마에 직면해 있습니다:

• 양방향 어텐션 (Bidirectional Attention): 고품질 dLLM(예: LLaDA, Dream)은 모든 토큰이 전체 문맥에 접근할 수 있는 완전한 양방향 어텐션을 사용합니다. 이는 생성 품질을 극대화하지만, 표준 프리픽스 키-값(Prefix KV) 캐싱을 불가능하게 만듭니다. 매 디노이징 단계마다 전체 시퀀스에 대한 어텐션을 다시 계산해야 하므로, 배치 서빙(batch-serving) 시나리오에서 추론 처리량이 저하됩니다.
• 인과적 어텐션 (Causal Attention): 인과적 dLLM(예: WeDLM)은 좌측에서 우측으로 향하는 의존 구조를 유지하여 효율적인 KV 캐싱과 높은 처리량을 가능하게 합니다. 그러나 이는 문맥을 좌측으로 엄격히 제한하여 우측의 정보를 모두 버리게 되며, 결과적으로 양방향 모델에 비해 생성 품질이 크게 저하됩니다.

블록 확산(고정된 블록 내에서만 양방향 어텐션을 제한함)이나 하이브리드 AR-확산 변형과 같은 기존의 절충안들은 이러한 핵심적인 긴장 관계를 해결하지 못합니다. 즉, 우측 문맥을 작은 윈도우로 제한하거나, 누락된 우측 정보를 완전히 복구하지 못한 채 구조적 복잡성만을 높입니다.

2. 방법론: 이중 초점 dLLM 및 R2LM

본 논문은 **비대칭 양방향 문맥(asymmetric bidirectional context)**에 기반한 이중 초점(Bifocal) dLLM 패러다임을 소개합니다. 좌우 문맥을 완전한 어텐션을 통해 대칭적으로 다루는 대신, 모델은 각 방향에 대해 별도의 메커니즘을 사용합니다:

1. 좌측 문맥 (Left Context): 표준 **인과적 트랜스포머 백본(causal Transformer backbone)**에 의해 처리되어, 프리픽스 KV 캐싱과의 완전한 호환성을 보장합니다.
2. 우측 문맥 (Right Context): 병렬로 실행되는 경량화된 비-어텐션 경로를 통해 공급됩니다.

저자들은 이 패러다임을 **R2LM (Right-to-Left Mamba)**으로 구체화했습니다:

• 아키텍처: R2LM은 사전 학습된 오토레그레시브 백본에 **역방향 Mamba 상태 공간 모델(SSM) 사이드카(sidecar)**를 부착합니다.
• 메커니즘:
  • 인과적 백본은 토큰을 좌에서 우로 처리하여 정밀한 좌측 문맥을 제공하고 KV 캐시를 유지합니다.
  • R2LM 스트림은 시퀀스를 역방향(우에서 좌로)으로 처리합니다. 이 스트림은 백본으로부터 히든 스테이트를 가져와 시퀀스를 뒤집고, 선택적 스캔(selective scan, Mamba)을 적용한 뒤 다시 원래 순서로 뒤집습니다.
  • 이 역방향 스트림은 우측 문맥으로부터 얻을 수 있는 정보 이득을 근사하는 **잔차 보정 신호($\Delta^{R2L}_i$)**를 생성합니다.
• 통합: R2LM 스트림은 매 $k$번째 디코더 레이어마다 포워드 훅(forward hook)을 통해 부착됩니다. 이 기여도는 학습 가능한 스칼라 $s$(0으로 초기화됨)에 의해 게이팅되며, 이를 통해 증강된 모델이 초기화 시 인과적 베이스라인과 비트 단위로 동일(bit-identical)하도록 보장합니다.
• 추론: 병렬 디코딩 중에 프롬프트는 키(key), 값(value), 그리고 포스트 훅(post-hook) 히든 스테이트를 캐싱하기 위해 한 번만 프리필(pre-fill)됩니다. 이후의 디노이징 단계는 생성 블록만을 처리하며, KV 캐시를 읽고 생성 토큰들에 대해서만 R2L 스캔을 수행합니다. 이는 우측 문맥을 주입하면서도 인과적 서빙의 처리량 프로필을 유지합니다.

3. 주요 기여

1. 새로운 패러다임: 표준 인과적 프리픽스 KV 캐싱을 유지하면서도 연속적이고 위치 인식이 가능한 우측 문맥을 제공하는 최초의 설계인 이중 초점 dLLM을 도입했습니다.
2. R2LM 구현: 인과적 트랜스포머와 역방향 Mamba SSM 사이드카를 결합한 특정 아키텍처를 제안합니다. 이는 공동 지속 사전 학습(joint continued pretraining)과, 백본을 동결한 채 R2L 스트림만 학습하는 "플러그인(plug-in)" 모드를 모두 지원하여 기존 AR 체크포인트의 가중치를 해치지 않고 통합할 수 있습니다.
3. 이론적 기초: 이 접근 방식은 우측 문맥 정보를 인과적 좌측 문맥 항에 대한 잔차 보정으로 취급하는 가법적 로그 사후 확률 분해(additive log-posterior decomposition)에 근거합니다.

4. 실험 결과

저자들은 Qwen3-1.7B 모델을 사용하여 60B 토큰의 지속 사전 학습을 통해 제어된 세 가지 방식(양방향, 인과적, R2LM)의 비교 평가를 수행했습니다.

• 생성 품질:

  • 장기 타겟 작업 (Long-Target Tasks): R2LM은 멀티 토큰 작업에서 평균 **47.44%**의 정확도를 달ol하여, 인과적 베이스라인(43.90%)보다 3.54%p 높았고, 양방향 베이스라인(44.78%)보다 2.66%p 높았습니다.
  • 단기 타겟 작업 (Short-Target Tasks): R2LM은 단일 토큰 작업(예: MMLU, BoolQ)에서 인과적 베이스라인보다 성능이 향상되었으며, 양방향 모델과의 격차를 크게 좁혔습니다.
  • 종합: R2LM은 **전체 평균(47.71% vs 46.74%)**에서 양방향 dLLM을 넘어섰으며, 대부분의 벤치마크에서 인과적 베이스라인을 상회했습니다.
  • 플러그인 효율성: R2LM-PI 변형(백본을 동결하고 5B 토큰으로 R2L 스트리즘만 학습)은 가장 높은 전체 평균(47.76%)을 달성했으며, 공동 학습 모델의 장기 타겟 이득 중 59%를 단 1/12의 학습 데이터와 1/10의 학습 파라미터만으로 회복했습니다.

• 추론 효율성:

  • 처리량 (Throughput): 프롬프트 길이가 4096인 배치 서빙($B=8$) 환경에서, R2LM은 양방향 dLLM보다 12.9배 높은 처리량을 기록했으며, AR 베이스라인보다 1.9배에서 2.9배 빠른 속도를 보였습니다.
  • 확장성: 양방향 처리량은 프롬프트 길이가 길어짐에 따라 (전체 어텐션을 재계산하기 때문에) 급격히 떨어졌으나, R2LM의 KV 캐싱 경로는 거의 평탄하게 유지되었습니다.
  • 메모리: R2LM은 인과적 모델의 메모리 프로필을 유지하여, 긴 시퀀스 길이($P=16384$)에서 양방향 모델 대비 약 32%의 메모리를 절감했습니다.

• 절제 연구 (Ablation Studies):

  • 실험을 통해 이러한 이득이 단순히 파라미터 수 증가 때문이 아님을 확인했습니다. 동일한 파라미터를 가진 MLP 어댑터(방향성 없음)는 퍼플렉시티(perplexity)를 개선하지 못했고, 좌측-우측 Mamba(반대 방향)는 베이스라인보다 성능이 낮았습니다. 오직 **역방향 Mamba(Reverse Mamba)**만이 양방향 성능과의 격차를 성공적으로 줄였습니다.

5. 의의 및 주장

본 논문은 이중 초점 dLLM이 생성 품질과 추론 효율성 사이의 오랜 트레이드오프를 해결한다고 주장합니다. 우측 문맥 경로를 어텐션 메커니즘에서 분리함으로써, R2LM은 다음을 달ей합니다:

• 높은 품질: 우측 문맥을 효과적으로 활용하여 평균 벤치마크에서 양방향 dLLM을 능가합니다.
• 높은 효율성: 인과적 모델의 KV 캐싱 능력을 유지하여, 배치 서빙 시나리오에서 상당한 처리량 향상과 함께 병렬 디코딩을 가능하게 합니다.
• 모듈성: 플러그인 기능은 최소한의 데이터와 파라미터 업데이트만으로 기존의 인과적 사전 학습 모델에 우측 문맥 메커니즘을 추가할 수 있게 합니다.

저자들은 모델 크기(1.7B)와 학습 예산이 dLLM이 가장 강력한 수치를 발표하는 규모보다 작다는 점과, R2L 스트림이 약 10.8%의 파라미터 오버헤드를 추가하여 단일 요청 추론 속도를 약간 늦출 수 있다는 점을 한계점으로 언급했습니다.