MARS: 자기회귀 모델의 멀티 토큰 생성을 가능하게 하는 기법

기존의 멀티 토큰 생성 방식은 별도의 드래프트 모델을 운영하거나 아키텍처를 수정해야 하는 부담이 있었다. MARS는 가벼운 Fine-tuning만으로 기존 모델의 성능을 유지하면서도 한 번의 연산으로 여러 토큰을 생성해 추론 효율성을 극대화한다.

기존의 Autoregressive(AR) 모델은 Attention Mechanism을 통해 이전 토큰들을 참조하여 바로 다음 하나의 토큰만을 예측한다. 이 과정에서 'The answer is'와 같이 매우 예측 가능한 문구조차 매번 독립적인 연산을 거쳐야 하므로 계산 자원이 낭비되는 한계가 있다.

MARS는 이 문제를 해결하기 위해 모델이 미래의 여러 위치를 동시에 바라보게 한다. 핵심은 [MASK] 토큰을 활용하는 것이다. 학습 시 입력 시퀀스의 일부를 [MASK]로 치환하고, 모델이 이 마스크된 위치의 정답 토큰들을 한 번에 맞히도록 훈련시킨다. 이때 기존 AR 모델의 Causal Attention 구조를 그대로 유지하면서 마스크된 위치들 사이에도 인과 관계를 부여하여 모델의 기본 동작 원리를 해치지 않는다.

결과적으로 모델은 확신이 있는 구간에서는 여러 토큰을 한 번에 쏟아내고, 불확실한 구간에서는 기존처럼 하나씩 생성하는 유연함을 갖게 된다. 이는 마치 숙련된 타자가 뻔한 공은 미리 예측해서 휘두르고 까다로운 공은 신중하게 보고 치는 것과 유사한 원리다.

MARS는 Mask AutoRegreSsion의 약자로, 기존 SFT 체크포인트에서 시작하여 두 가지 스트림을 병렬로 학습시킨다. 첫 번째는 Clean Stream으로, 원본 토큰들을 그대로 입력하여 표준적인 Next-token Prediction을 수행한다. 두 번째는 Noisy Stream으로, 특정 블록 크기 B만큼의 토큰을 [MASK]로 교체하고 모델이 이를 복구하도록 유도한다.

학습 과정에서 특수한 Attention Mask를 사용한다. [x; x̃] 형태로 연결된 입력에서 Clean Stream은 표준 Causal Attention을 사용하고, Noisy Stream의 각 마스크 토큰은 이전 블록의 Clean 토큰들과 현재 블록 내의 이전 마스크 토큰들만 참조할 수 있도록 제한한다. [입력 토큰 및 마스크 위치 → 인과적 마스크 연산 → 각 위치별 Logit 출력] 과정을 거쳐 모델은 불완전한 문맥에서도 미래 토큰을 예측하는 법을 배운다.

손실 함수는 L = Lmask + LAR 구조를 취한다. 마스크된 위치의 예측 오차(Lmask)와 Clean Stream의 예측 오차(LAR)를 동일한 비중으로 합산한다. [정답 토큰과 예측 확률 분포 → Cross-Entropy 계산 → 오차 합산] 과정을 통해 모델은 멀티 토큰 예측 능력을 습득하면서도 기존의 단일 토큰 생성 능력을 잃지 않도록 정렬된다.

Qwen2.5-0.5B 및 7B 모델을 대상으로 6개 표준 벤치마크에서 평가를 진행했다. 단일 토큰 생성 모드(τ=1.0)에서 MARS-7B는 평균 58.1점을 기록하여 기본 AR SFT 모델의 56.6점보다 높은 성능을 보였다. 특히 GSM8K(+4.5)와 HumanEval(+3.0)에서 유의미한 향상이 관찰되었다.

멀티 토큰 생성 모드에서는 성능 손실을 최소화하면서도 높은 효율성을 증명했다. Confidence Threshold τ=0.95 설정 시, Qwen2.5-7B 모델은 1.5~1.7배의 처리량(Throughput) 향상을 보였으며, 정확도 하락은 평균 1.3점에 불과했다. 이는 기존의 Speculative Decoding이나 Medusa와 같은 기법들이 추가적인 메모리나 파라미터를 요구하는 것과 대조적인 성과다.

실제 벽시계 시간(Wall-clock time) 분석 결과, 새롭게 제안된 블록 단위 KV Caching을 적용했을 때 Qwen2.5-7B 배치 추론 환경에서 표준 AR 대비 최대 1.71배 빠른 속도를 기록했다. 이는 모델의 구조적 변경 없이 순수하게 학습 기법과 추론 전략의 개선만으로 달성한 수치다.

MARS의 핵심 차별점은 기존 Block Diffusion 방식들이 범했던 세 가지 오류(비인과적 어텐션, 로짓 정렬 불일치, 비순차적 생성)를 해결했다는 점이다. MARS는 모든 과정에서 엄격한 Left-to-right 인과 관계를 유지하며, Logit 또한 표준 AR 모델과 동일하게 Right-shifted 방식을 고수하여 사전 학습된 지식과의 충돌을 방지한다.

학습 시 Clean Stream을 유지하는 것은 'Signal Decay' 현상을 막기 위한 필수적인 장치다. 블록 크기 B가 커질수록 마스크된 위치에서 받는 AR 신호의 비율은 1/B로 급감하는데, Clean Stream의 LAR 손실 함수를 병렬로 계산함으로써 모델이 항상 50% 이상의 AR 신호를 유지하며 학습할 수 있도록 설계했다.

추론 시에는 Sliding Window 방식을 채택한다. 매 단계에서 B개의 [MASK]를 붙여 연산한 뒤, 설정된 임계값 τ를 넘는 토큰들만 확정(Accept)한다. 최소 한 개의 토큰은 반드시 확정되도록 보장함으로써 모델이 확신이 없을 때는 자연스럽게 표준 AR 모드로 퇴행(Graceful Degradation)하게 만든다.

MARS 학습 시 Clean 시퀀스와 Noisy 시퀀스를 이어 붙여 처리하므로, 샘플당 시퀀스 길이가 두 배가 되어 학습 연산량이 표준 SFT 대비 약 2배 정도 증가한다. 또한 블록 단위 KV Caching은 배치 내 샘플 간의 동기화가 필요하므로, 매우 큰 배치 사이즈에서는 효율성이 다소 저하될 수 있는 여지가 있다.