[박찬 기자]

<이미지를 클릭하시면 크게 보실 수 있습니다>

에이전트형 AI 워크플로우가 확산면서 긴 추론 체인에 따른 비용과 지연이 급증하는 가운데, 미국 연구진이 모델 구조를 거의 바꾸지 않으면서도 추론(reasoning) 속도를 최대 3배까지 높일 수 있는 새로운 학습 기법을 제시했다.

메릴랜드대학교, 로렌스 리버모어 국립 연구소, 컬럼비아대학교, 투게더AI 공동 연구진은 23일(현지시간) 추론 가속 메커니즘을 모델 내부 가중치에 내장하는 학습 기법 '자기 증류를 통한 멀티 토큰 예측(Multi-Token Prediction via Self-Distillation)'을 온라인 아카이브에 공개했다.

현재 대형언어모델(LLM)은 '다음 토큰 예측(NTP)' 방식으로 작동한다. 한번 계산할 때마다 토큰 하나씩만 만들어내는 구조다.

문제는 추론에 특화된 모델이 수천개에 이르는 '사고 사슬(CoT)' 문장을 생성해야 할 때다. 토큰을 하나씩 차례로 만들다 보니 속도가 느려지고, 계산 비용도 크게 늘어나는 병목 현상이 발생한다.

이 문제를 해결하기 위한 방법으로 '멀티 토큰 예측(MTP)'이 제안됐다. 이는 한번의 계산으로 여러 개의 토큰을 동시에 예측하는 방식이다. 다만 기존 MTP 학습 방식은 각 토큰을 따로 맞히는 데 초점을 두다 보니, 문장이 어색해지거나 "the the the…"처럼 같은 단어가 반복되는 오류가 생기는 한계가 있었다.

연구진은 이런 한계를 극복하기 위해 '온라인 증류(online distillation)' 방식의 교사-학생(student–teacher) 구조를 도입했다.

학생 모델은 한번에 여러 개의 토큰 묶음(블록)을 생성한다. 즉, 빠르게 여러 단어를 동시에 만들어내는 역할을 한다. 교사 모델은 기존의 성능이 검증된 단일 토큰 예측 모델이다. 학생 모델이 생성한 토큰 블록이 자연스럽고 일관성 있는지, 확률적으로 타당한지를 평가하고 지도한다.

예를 들어 학생 모델이 "lion bamboo"처럼 어색한 단어 조합을 만들어내면, 교사 모델은 이에 대해 높은 손실값(loss)을 부여한다. 그러면 학생은 부자연스러운 출력을 줄이도록 다시 학습하게 된다.

이 과정은 강화 학습(RL)의 온-폴리시(on-policy) 학습과 비슷하다. 학생 모델이 스스로 만들어낸 결과를 바탕으로, 그때그때 피드백을 받고 수정해 나가는 방식이기 때문이다.

이 방법의 핵심은 모델 구조를 거의 바꾸지 않아도 된다는 점이다. 기존 모델의 임베딩 행렬에서 비어 있던 자리에 특수 토큰 '' 하나만 추가하면, 토큰을 순서대로 하나씩 계산하던 방식을 여러 개를 동시에 처리하는 병렬 연산 구조로 바꿀 수 있다. MoE, 윈도우드 어텐션, SSM 레이어 같은 내부 구조도 그대로 유지된다.

별도의 작은 모델을 추가로 붙여야 하는 '추측 디코딩(speculative decoding)'과 달리, 복잡한 인프라 변경 없이 기존 모델을 거의 그대로 활용할 수 있다는 설명이다.

여러 토큰을 한번에 생성하는 방식은 속도를 크게 높일 수 있다는 장점이 있지만, 동시에 정확도가 떨어질 위험도 따른다. 연구진은 이러한 한계를 보완하기 위해 '콘프어댑트(ConfAdapt)'라는 적응형 디코딩 전략을 제안했다.

콘프어댑트는 단계마다 일정한 신뢰도 기준값(90%)을 설정하고, 모델이 예측한 토큰 묶음의 확신 정도를 평가한다. 신뢰도가 기준 이상으로 충분히 높다고 판단되면 해당 토큰 블록을 한번에 채택해 빠르게 다음 단계로 넘어간다.

반대로 확신이 낮은 구간에 대해서는 다시 기존의 단일 토큰 예측 방식으로 전환해 다시 세밀하게 계산한다.

이 방식은 예측이 비교적 쉬운 구조적 문장에서는 여러 토큰을 동시에 출력해 시간을 절약하고, 복잡하거나 애매한 부분에서는 연산을 집중함으로써 정확도를 유지하도록 설계됐다.

연구진은 이 기법을 오픈 웨이트 기반의 지시 튜닝 모델인 '라마-3.1-8B-매그파이'와 '큐원3-4B-인스트럭트-2507'에 적용해 성능을 검증했다.

수학 추론 데이터셋인 '메타매스QA(MetaMathQA)'에서 실험한 결과, 라마-3.1-8B 모델은 기존 대비 약 3배 빠른 속도를 기록하면서도 정확도 감소는 3% 미만에 그쳤다. 큐원3-4B도 속도는 약 3배 향상됐으며, 정확도는 약 7% 감소하는 수준이었다.

더 공격적인 설정을 적용할 경우 최대 5배까지 속도를 끌어올릴 수 있었지만, 그만큼 정확도 손실도 커졌다.

<이미지를 클릭하시면 크게 보실 수 있습니다>

흥미로운 점은 멀티 토큰 예측 학습에 직접 포함되지 않았던 영역, 예를 들어 창작 글쓰기나 요약 작업에서도 추론 가속 효과가 유지됐다는 사실이다. 이는 이 기법이 특정 과제에만 국한되지 않고, 일반적인 성능 개선 가능성을 지니고 있음을 시사한다.

기존 배치 기반 NTP는 총 처리량에는 최적화돼 있지만, 개별 사용자의 응답 속도 개선에는 한계가 있다. 이번 접근은 GPU를 단일 사용자 요청에 최대한 활용해 체감 속도를 줄이는 데 초점을 둔다.

연구진은 학습된 모델을 허깅페이스에 공개했으며, MTP 프레임워크 코드도 곧 배포할 예정이다.

박찬 기자 cpark@aitimes.com

<저작권자 Copyright ⓒ AI타임스 무단전재 및 재배포 금지>