임베딩 공간 프로빙을 통한 효율적인 학습 불필요 멀티 토큰 예측

기존의 멀티 토큰 예측(MTP)은 별도의 보조 모델을 학습시키거나 가중치를 수정해야 하는 번거로움이 있었다. 이 논문은 얼어붙은(frozen) 모델의 임베딩 공간에 특수한 '마스크 토큰'을 주입하는 것만으로도 미래 토큰을 정확히 예측할 수 있음을 증명하여, 추가 비용 없이 추론 효율을 극대화한다.

기존 LLM은 한 번에 하나의 토큰만 생성하는 Autoregressive 방식을 사용하며, 이는 GPU 자원을 충분히 활용하지 못해 추론 속도가 느리다는 한계가 있다. 이를 해결하기 위해 여러 토큰을 동시에 예측하는 Multi-Token Prediction(MTP) 연구가 진행되어 왔으나, 대부분 별도의 예측 헤드를 학습시키거나 가중치를 변경해야 하는 부담이 있었다. 이 논문은 모델이 이미 학습 과정에서 미래 토큰에 대한 잠재적인 정보를 내부 표현(Hidden State)에 담고 있다는 점에 주목한다. 연구진은 임베딩 공간에서 추출한 '마스크 토큰'을 프롬프트 뒤에 붙여 모델을 '프로빙(Probing)'하면, 모델이 별도의 학습 없이도 자연스럽게 미래 토큰의 분포를 출력한다는 사실을 발견했다. 특히 Transformer의 Decoder 레이어를 거칠수록 마스크 토큰의 표현이 실제 정답 토큰의 상태와 코사인 유사도 측면에서 점점 일치해가는 현상을 이론적, 실험적으로 입증했다. 이를 통해 추가 모델 없이도 높은 정확도의 미래 예측이 가능해지며, 결과적으로 한 번의 연산으로 여러 단어를 동시에 확정 짓는 가속화가 이루어진다.

Mask Token Injection. 입력 프롬프트 x_{1:t} 뒤에 임베딩 공간에서 동적으로 생성된 마스크 토큰 m_1, ..., m_k를 주입한다. 초기 마스크 토큰은 프롬프트 임베딩의 평균값(Mean-of-prompt)으로 설정하며, 생성 단계마다 마지막으로 생성된 토큰의 정보를 반영하여 업데이트한다. [m_i[s+1] = m_i[s] + lambda(e_{t+s} - m_i[s]) 연산을 통해 이전 상태에 새로운 토큰 임베딩의 차분값을 더해 문맥 정보를 주입한다.] Dynamic Tree Construction. 마스크 토큰의 출력 로짓에서 Top-K 후보를 샘플링하여 투기적 토큰 트리를 구성한다. 누적 확률을 기반으로 유망한 경로만 확장하며, 부모 토큰과 동일한 토큰이 자식 노드에서 반복되는 경우를 제거하는 경량 프루닝을 적용한다. [각 노드의 확률 P(c)에 자식 토큰의 조건부 확률 P(t_j|l_n)을 곱하여 전체 경로의 신뢰도를 계산하고, 이를 기준으로 트리의 너비를 조절한다.] Parallel Verification. 생성된 토큰 트리 후보들을 원본 모델을 통해 한 번에 병렬 검증한다. 검증 과정에서 원본 모델의 Next-token 분포와 일치하는 토큰들만 수용하며, 이는 수학적으로 손실 없는 생성을 보장한다. GPU 효율을 높이기 위해 트리 구조에 최적화된 어텐션 마스크와 위치 인덱스 오프셋 기법을 사용하여 연산 오버헤드를 최소화한다.

SpecBench 벤치마크 평가 결과, LLaMA3 모델에서 기존 학습 불필요 기법 대비 수용 길이(Acceptance Length)가 약 12% 향상되었으며, Qwen3 모델에서는 8-12%의 개선을 보였다. 이는 추가 학습 없이도 모델 내부의 잠재력을 끌어내어 더 긴 문장을 한 번에 예측할 수 있음을 의미한다. 처리량(Throughput) 측면에서는 LLaMA3.1-8B-Instruct 모델 기준 표준 디코딩 대비 최대 26% 높은 토큰 생성 속도를 기록했다. 특히 코딩(Coding) 작업에서 가장 높은 성능 향상을 보였으며, 이는 정형화된 패턴이 많은 작업에서 마스크 토큰 프로빙이 더 효과적임을 시사한다. Ablation Study를 통해 마스크 토큰의 초기화 전략 중 '프롬프트 임베딩 평균값'을 사용하는 방식이 가장 안정적인 성능을 보임을 확인했다. 또한, 제안된 트리 프루닝 기법은 평균 토큰 수용률을 최대 4%까지 추가로 향상시키는 효과를 입증했다.

본 연구는 별도의 MLP 헤드나 드래프트 모델을 사용하는 대신, 얼어붙은 모델의 임베딩 레이어와 LM 헤드를 그대로 재사용한다. 핵심은 마스크 토큰이 레이어를 통과하며 실제 미래 토큰의 Hidden State와 정렬된다는 점이다. 이론적 분석(Lemma 3.1)을 통해, 마스크 토큰과 실제 토큰 간의 코사인 유사도가 특정 임계값 이상일 경우, 실제 정답 토큰이 마스크 토큰의 Top-K 예측 집합에 포함될 확률이 보장됨을 증명했다. 실험적으로 레이어 15 이후부터 유사도가 급격히 상승하여 약 0.45 수준에 도달함을 확인했다. 블록 복잡도(Block Complexity)를 도입하여 병렬 처리되는 토큰 수와 연산 비용 간의 트레이드오프를 관리한다. 동적 트리 확장 알고리즘은 모델의 불확실성에 따라 트리 구조를 유연하게 변경하여, 고정된 트리 구조를 사용하는 기존 방식보다 다양한 도메인에서 강건한 성능을 보여준다.

오픈 엔드 작업(작문, 롤플레이)에서는 단일 마스크 토큰이 유리한 반면, 폐쇄형 작업(번역, 수학)에서는 더 깊은 트리가 유리한 등 작업 성격에 따라 최적의 마스크 토큰 수가 달라지는 경향이 있다. 또한 현재 구현은 동적 어텐션 마스크 생성 시 일부 메모리 오버헤드가 발생할 수 있어 향후 vLLM 등 추론 프레임워크와의 통합 연구가 필요하다.