1. 문제 정의 (Problem)

기존의 자기회귀 (Autoregressive, AR) 언어 모델은 문맥이 매우 예측 가능하더라도 매번 한 번의 포워드 패스 (forward pass) 에 단 하나의 토큰만 생성합니다. 이는 불필요한 계산 낭비를 초래합니다.

• 기존 접근법의 한계:
  • Speculative Decoding: 타겟 모델 외에 별도의 '드래프트 (draft)' 모델을 유지해야 하므로 메모리 사용량이 두 배가 되고 오케스트레이션이 복잡해집니다.
  • Multi-head 접근법 (Medusa, EAGLE 등): 모델 구조에 추가 예측 헤드를 붙여야 하므로 파라미터가 증가하고 별도의 학습이 필요합니다.
  • Block Masked Diffusion: 기존 AR 모델을 블록 단위로 마스킹된 확산 모델로 변환할 때, 주의 메커니즘 (attention) 이나 생성 순서 등의 설계 차이로 인해 추론 능력 (특히 논리 추론 및 코딩) 이 크게 저하되는 문제가 있었습니다.

2. 방법론 (Methodology)

저자들은 **MARS (Mask AutoRegreSsion)**를 제안합니다. 이는 아키텍처 수정이나 추가 파라미터 없이, 기존 지시 학습 (Instruction-tuned) 된 AR 모델을 가벼운 파인튜닝을 통해 한 번의 포워드 패스로 여러 토큰을 예측할 수 있도록 만드는 방법입니다.

핵심 아이디어 및 설계 원칙

MARS 는 AR 모델의 본질적인 특성을 해치지 않으면서 다중 토큰 생성 능력을 추가하는 것을 목표로 합니다. 이를 위해 기존 Block Diffusion 과 AR 모델 간의 4 가지 격차 (Gap) 를 분석하고, 3 가지는 제거 가능한 설계 선택임을 규명했습니다.

1. 격차 분석 및 해결:

   • 격차 1 (토큰 마스킹): 다중 토큰 예측을 위해 필수적인 요소이므로 수용.
   • 격차 2 (주의 패턴): 기존 Diffusion 은 블록 내에서 양방향 (Bidirectional) 주의를 사용했으나, MARS 는 인과적 (Causal) 주의만 유지하여 AR 모델의 특성을 보존합니다.
   • 격차 3 (Logits 정렬): AR 모델은 $t$ 위치에서 $t+1$ 토큰을 예측하는 'Right-shifted' 방식을 사용합니다. MARS 는 이 방식을 그대로 유지합니다.
   • 격차 4 (생성 순서): 기존 방식은 블록 내에서 토큰을 순서 없이 예측했으나, MARS 는 왼쪽에서 오른쪽 (Left-to-Right) 으로 엄격하게 토큰을 받아들이도록 설계합니다.

2. 학습 전략 (Training):

   • 이중 스트림 입력: 입력 시퀀스를 두 가지로 복사합니다.
     • Clean Stream: 원본 토큰을 유지하여 일반적인 AR 다음 토큰 예측 (Next-token prediction) 을 학습.
     • Noisy Stream: 블록 단위의 토큰을 [MASK] 로 대체하여, 이전 블록의 Clean 토큰을 문맥으로 사용하여 마스킹된 토큰을 예측하도록 학습.
   • SFT Loss 의 중요성: 블록 크기가 커질수록 AR 와 유사한 학습 신호가 희석되는 문제를 해결하기 위해, Clean Stream 에 대한 **표준 AR 손실 함수 (SFT loss)**를 병렬로 적용합니다. 이는 모델이 마스킹 예측을 배우는 동안에도 원래의 AR 능력을 유지하도록 보장합니다.

3. 추론 전략 (Inference):

   • 슬라이딩 윈도우: 현재 접두사 (Prefix) 뒤에 $B$개의 [MASK] 토큰을 붙이고 한 번의 포워드 패스를 실행합니다.
   • 신뢰도 기반 수용 (Confidence Thresholding): 모델이 예측한 토큰의 확률이 임계값 ($\tau$) 이상이면 해당 토큰을 즉시 수용하고 다음 토큰을 예측합니다.
   • 적응형 생성: 모델이 확신할 때는 여러 토큰을 한 번에 생성하고, 불확실할 때는 단일 토큰 생성으로 자연스럽게 회귀합니다.
   • KV Cache 최적화: 배치 추론 시 블록 단위 KV 캐싱 전략을 도입하여 재계산을 최소화하고 실제 처리 속도를 높입니다.

3. 주요 기여 (Key Contributions)

1. 격차 분석 및 해결: AR 와 블록 마스킹 예측 간의 4 가지 격차 중 3 가지는 제거 가능한 설계 선택임을 증명하고, 이를 해결함으로써 아키텍처 변경 없이 품질을 유지하는 방법을 제시했습니다.
2. MARS 제안: 추가 파라미터나 구조 변경 없이, 기존 SFT 데이터를 재사용하여 다중 토큰 생성 능력을 부여하는 경량 파인튜닝 방법론을 제안했습니다.
3. SFT Loss 의 역할 규명: 블록 크기가 커질수록 AR 신호가 약해지는 문제를 해결하기 위해 Clean Stream 에 대한 SFT Loss 가 필수적임을 증명했습니다.
4. 실용적 속도 - 품질 트레이드오프: 신뢰도 임계값 ($\tau$) 을 실시간으로 조절하여 서비스 부하에 따라 처리량 (Throughput) 을 동적으로 조정할 수 있는 유연한 시스템을 구현했습니다.

4. 실험 결과 (Results)

Qwen2.5-0.5B 및 7B 모델을 기반으로 한 실험 결과입니다.

• 단일 토큰 모드 (Quality Preservation):
  • MARS 는 단일 토큰 생성 모드 ($\tau=1.0$) 에서 기존 AR SFT 모델보다 동등하거나 더 높은 성능을 보였습니다 (예: 7B 모델에서 평균 점수 56.6 $\rightarrow$ 58.1).
  • Block Diffusion 은 BBH, MMLU-Pro 등에서 성능이 급격히 하락한 반면, MARS 는 추론 및 코딩 능력을 유지했습니다.
• 다중 토큰 모드 (Throughput Gain):
  • 신뢰도 임계값을 낮추어 다중 토큰 생성을 허용하면, 정확도 손실은 미미한 반면 처리량은 1.5~1.7 배 증가했습니다.
  • 예: 7B 모델에서 평균 정확도 1.3 포인트 감소 (58.1 $\rightarrow$ 56.8) 로 1.68 토큰/포워드 생성.
• 실제 처리 속도 (Wall-clock Speedup):
  • 블록 단위 KV 캐싱 전략을 적용한 배치 추론에서, Qwen2.5-7B 기준 최대 1.71 배의 실제 실행 시간 단축을 달성했습니다 (Batch size 4 기준).
  • 캐싱이 없을 경우 배치 크기가 커질수록 성능이 저하되었으나, 캐싱 전략으로 이를 극복했습니다.

5. 의의 및 결론 (Significance)

• 단일 모델, 이중 기능: MARS 는 하나의 모델 체크포인트로 고품질 단일 토큰 생성과 고속 다중 토큰 생성을 모두 지원합니다. 모델 교체나 재학습 없이 서비스 환경에서 실시간으로 속도 조절이 가능합니다.
• 배포 효율성: 별도의 드래프트 모델이나 추가 헤드가 필요 없어 메모리 오버헤드가 없고, 기존 AR 파이프라인과 호환됩니다.
• 실용성: 추론 비용과 지연 시간 (Latency) 을 상황에 따라 유연하게 조절할 수 있는 'Latency-Quality Knob'을 제공하여, 실제 서비스 배포에 매우 실용적인 솔루션을 제시합니다.

결론적으로 MARS 는 자기회귀 모델의 근본적인 한계 (단일 토큰 생성) 를 구조 변경 없이 학습 전략만으로 극복하여, 효율성과 성능을 동시에 잡은 혁신적인 방법론입니다.