brunch

다중 토큰 예측 잠재력의 발견?

당신의 LLM은 미래를 알고 있다구

by 미미니

Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential​ 논문은 언어 모델(LLM)의 추론 속도를 획기적으로 높이는 새로운 방법을 제안하며, 기존의 순차적(autoregressive) 생성 방식의 한계를 극복하는 흥미로운 접근법을 제시합니다. 이 논문은 기술적으로 복잡해 보일 수 있지만, 결국 더 빠르고 똑똑한 언어 모델을 만드는 데 기여하는 재미있는 아이디어로 가득 차 있습니다. AI가 사람처럼 “한 번에 문장 단위로 생각”할 수 있게 만드는 첫걸음이죠! 이 글에서는 논문의 핵심 아이디어를 재미있고 이해하기 쉽게 설명하고, 왜 이 연구가 중요한지, 어떤 잠재력을 가지고 있는지를 풀어보겠습니다.


배경: 언어 모델의 속도 병목현상


우선, 언어 모델이 어떻게 작동하는지 간단히 짚고 넘어가겠습니다. 현재 대부분의 언어 모델은 자연어 생성(autoregressive generation) 방식으로 작동합니다. 즉, 한 번에 하나의 토큰(단어 또는 단어의 일부)을 예측하고, 이를 순차적으로 반복해서 문장을 완성합니다. 예를 들어, “What is two plus two?“라는 질문에 대해 모델은 “two”, “plus”, “two”, “equals”, “four”를 하나씩 차례로 생성합니다.

하지만 이 방식에는 큰 단점이 있습니다:

속도가 느리다: 한 토큰을 생성할 때마다 모델 전체를 실행해야 하므로 시간이 오래 걸립니다.

병렬 처리가 어렵다: 다음 토큰을 예측하려면 이전 토큰이 필요하니 병렬적으로 여러 토큰을 생성할 수 없습니다.

이 논문은 이런 문제를 해결하기 위해, 한 번에 여러 토큰을 예측(multi-token prediction, MTP)할 수 있는 방법을 제안합니다. 마치 사람이 문장을 먼저 머릿속으로 구성한 뒤 빠르게 말하는 것처럼, 모델도 여러 토큰을 동시에 뱉어낼 수 있게 하자는 거죠!


핵심 아이디어: 모델은 이미 미래를 알고 있다!


이 논문의 가장 흥미로운 출발점은, 기존의 순차적 언어 모델이 이미 미래의 토큰에 대한 정보를 어느 정도 가지고 있다는 관찰입니다. 예를 들어, “What is two plus two?“라는 프롬프트가 주어졌을 때, 모델은 “equals four”라는 답을 순차적으로 생성하지만, 사실 내부적으로는 이 토큰들이 나올 가능성을 이미 어느 정도 “예감”하고 있습니다.

이를 확인하기 위해 연구진은 실험을 했습니다:

1. 프롬프트 뒤에 의미 없는 플레이스홀더 토큰(예: <->)을 붙여서 모델이 어떤 출력을 내놓는지 확인했습니다.

2. 놀랍게도, 모델의 출력 logits(토큰 확률 분포)에서 정확한 미래 토큰이 상위 200개 안에 포함되어 있었습니다!

이 발견은 모델이 이미 미래를 “예측”할 잠재력을 가지고 있다는 뜻입니다. 문제는 이 잠재력을 어떻게 끌어내느냐였죠. 그래서 연구진은 모델을 약간 튜닝하여 이 능력을 극대화하는 방법을 제안했습니다.


어떻게 여러 토큰을 동시에 예측할까?


논문은 이 문제를 해결하기 위해 다섯 가지 핵심 기술을 제안합니다. 하나씩 쉽게 풀어보겠습니다:


(1) 마스크 토큰(Mask Tokens)

모델이 미래 토큰을 예측하도록 유도하기 위해, 입력 시퀀스 끝에 마스크 토큰이라는 특별한 토큰을 추가합니다. 예를 들어, “What is two plus two?” 뒤에 [m1, m2, m3] 같은 마스크 토큰을 붙여서 모델이 이 마스크를 채우도록 훈련시킵니다. 이 마스크 토큰은 “미래에 올 토큰들”을 나타내며, 모델이 한 번에 여러 토큰을 예측하도록 돕습니다.


(2) 게이트드 LoRA(Gated LoRA)

모델을 새로 훈련시키는 건 비용이 많이 들죠. 그래서 연구진은 LoRA(Low-Rank Adaptation)라는 경량화된 미세 조정 기법을 사용했습니다. 하지만 기존 LoRA는 모델의 기존 성능(즉, 다음 토큰 예측)을 망칠 위험이 있습니다. 이를 해결하기 위해 게이트드 LoRA를 제안했는데, 이는:

NTP(Next-Token Prediction): 기존의 순차적 예측은 그대로 유지.

MTP(Multi-Token Prediction): 새로운 경로를 통해 미래 토큰 예측을 추가.

• 이를 위해 바이너리 마스크를 사용해 NTP와 MTP를 분리합니다. 결과적으로 기존 모델의 성능은 손상되지 않으면서 새로운 기능을 추가할 수 있죠!


(3) 샘플러 헤드(Sampler Head)

여러 토큰을 한 번에 예측하면, 이 토큰들이 서로 일관성(coherence) 있게 연결되지 않을 수 있습니다. 이를 해결하기 위해 연구진은 샘플러 헤드라는 경량의 2층 MLP(다층 퍼셉트론)를 추가했습니다. 이 샘플러는:

• 모델이 예측한 토큰 분포를 받아서, 이전에 생성된 토큰을 고려해 더 자연스러운 토큰 시퀀스를 선택합니다.

• 예를 들어, “two” 다음에 “plus”가 올 확률을 높여서 문장이 매끄럽게 이어지도록 합니다.


실험 결과: 속도와 품질의 완벽한 조화


논문은 Tulu3-8B 모델을 사용해 실험을 진행했습니다. 이 모델은 다양한 도메인(수학, 코딩, 대화 등)에서 테스트되었으며, 최대 9개의 토큰을 한 번에 예측하도록 훈련되었습니다. 주요 결과는 다음과 같습니다:

• 속도 향상:코딩과 수학 작업에서 약 5배 빠른 생성 속도를 달성했습니다. 일반 대화 및 지식 관련 작업에서는 약 2.5배 빠른 속도를 보였습니다.

• 품질 유지: 게이트드 LoRA 덕분에 기존 모델의 품질 저하 없이 이러한 속도 향상을 달성했습니다.

특히, 쿼드라틱 디코딩(Quadratic Decoding)이라는 기술을 통해 더 안정적인 토큰 검증을 가능하게 했고, LCM(Latent Consistency Matching) 손실 함수를 통해 예측의 일관성을 높였습니다. 결과적으로, 모델은 더 빠르게, 더 정확하게 문장을 생성할 수 있었습니다.


왜 이 연구는 중요한가?


이 연구는 언어 모델의 추론 속도라는 핵심 병목현상을 해결하는 데 큰 기여를 합니다. 특히,

효율성: 추론 속도가 2~5배 빨라지면, 대규모 언어 모델을 실시간 애플리케이션(예: 챗봇, 코드 생성기)에 더 쉽게 활용할 수 있습니다.

비용 절감: 더 적은 GPU 리소스로 동일한 작업을 수행할 수 있어, 클라우드 컴퓨팅 비용이 줄어듭니다.

확장성: 이 기술은 사전 훈련(pretraining)이나 다운스트림 작업에 적용될 가능성을 열어, 더 큰 모델에도 활용될 수 있습니다.


마무리:미래 예측의 미래 전망


이 논문은 언어 모델이 이미 가지고 있는 미래 예측 잠재력을 끌어내어, 속도와 품질이라는 두 마리 토끼를 모두 잡는 방법을 제시했습니다. 특히, 마스크 토큰, 게이트드 LoRA, 샘플러 헤드, 쿼드라틱 디코딩, LCM 손실이라는 다섯 가지 혁신은 모델의 효율성을 극대화하면서도 기존 성능을 유지합니다.

미래에는 이 기술이 사전 훈련 단계나 다양한 작업 적응 단계에 적용될 가능성이 큽니다. 또한, 확산 기반 생성(diffusion-based generation)과 결합하면 더 혁신적인 결과를 낼 수 있을 거예요. 언어 모델의 “미래 예측” 능력을 활용한 이 연구는 AI의 실시간 응용 가능성을 한층 더 높이는 중요한 발걸음입니다!


keyword
매거진의 이전글AI의 성격을 조종하다: Persona Vectors