엔트로피 기반 토큰 풀링으로 LLM 출력 길이 예측

초록

본 논문은 대형 언어 모델(LLM)의 출력 길이를 효율적으로 예측하기 위해 모델 내부 은닉 상태와 토큰 엔트로피를 활용한 두 가지 기법, 엔트로피‑가이드 토큰 풀링(EGTP)과 진행형 길이 예측(PLP)을 제안한다. 기존의 별도 경량 예측 모델이 갖는 높은 오버헤드와 확장성 문제를 극복하고, 특히 강화학습 샘플링과 같은 “one‑to‑many” 상황에서도 동적 예측이 가능하도록 설계하였다. 새롭게 구축한 ForeLen 벤치마크에서 EGTP는 평균 절대 오차(MAE)를 29 % 이상 감소시켰으며, 길이‑인식 스케줄러와 결합했을 때 전체 추론 처리량을 크게 향상시켰다.

상세 분석

본 연구는 LLM 서빙 환경에서 발생하는 ‘배럴 효과’를 근본적으로 완화하려는 목표를 갖는다. 배럴 효과는 배치 내 가장 긴 시퀀스에 맞춰 짧은 시퀀스를 패딩해야 하는 구조적 제약에서 비롯되며, 이는 GPU 메모리와 연산 자원을 비효율적으로 사용하게 만든다. 기존 접근법은 프롬프트만을 입력으로 하는 경량 예측 모델을 별도로 학습·배포하여 길이를 추정한다. 그러나 이러한 방식은 (1) 프롬프트 기반 정적 예측이 stochastic sampling, 특히 강화학습에서 동일 프롬프트에 대해 다수의 서로 다른 길이 결과를 생성하는 상황에 취약하고, (2) LMSYS와 같은 기존 벤치마크가 장기 시퀀스와 복합 추론을 충분히 포함하지 않아 일반화 성능이 낮으며, (3) 별도 모델 호출로 인한 추가 연산 비용과 시스템 복잡성을 초래한다는 한계를 가진다.

논문은 이러한 문제점을 해결하기 위해 LLM 자체의 은닉 상태에 내재된 ‘길이 신호’를 직접 활용한다는 핵심 아이디어를 제시한다. LLM이 토큰을 언제 출력할지를 판단하는 과정에서 내부 표현은 출력 길이에 대한 충분한 정보를 담고 있기 때문에, 이를 재활용하면 별도 모델 없이도 고정밀 길이 예측이 가능하다.

첫 번째 모듈인 Entropy‑Guided Token Pooling(EGTP)은 토큰별 엔트로피와 은닉 상태를 결합해 중요한 토큰에 가중치를 부여한다. 엔트로피는 다음 토큰 예측의 불확실성을 나타내며, 실험을 통해 엔트로피가 높을수록 길이 예측에 기여하는 gradient‑based importance가 증가함을 확인하였다(피어슨 r = 0.451). 따라서 각 토큰의 은닉 상태 h_i에 대해 엔트로피 H_i를 softmax(α·H_i) 형태의 가중치 w_i로 변환하고, 가중합 h = ∑w_i h_i 로 풀링한다. 이 과정은 기존의 평균·최대 풀링보다 중요한 정보를 보존하면서도 연산량이 거의 추가되지 않는다.

두 번째 단계는 길이 회귀를 위한 소프트 라벨 분포 학습이다. 연속적인 길이 y를 K개의 구간으로 이산화하고, 정답 구간 i에 대해 거리 기반 가우시안 형태의 소프트 라벨 p_j = exp(−|j−i|)/Z 를 만든다. 모델은 이 라벨에 대해 교차 엔트로피 손실을 최소화하고, 동시에 회귀값 ȳ = ∑p̂_i c_i (c_i는 구간 중심) 에 대해 MSE 손실을 최적화한다. λ 하이퍼파라미터로 두 손실을 균형 맞추어, 분류형 안정성과 회귀형 정확성을 동시에 달성한다.

동적 예측을 위한 Progressive Length Prediction(PLP)은 ‘one‑to‑many’ 상황을 겨냥한다. PLP는 디코딩 단계 t마다 현재까지 생성된 토큰들의 은닉 상태와 초기 프롬프트 특징 h를 결합해 z_t = Concat(h, {h′_1…h′_t}) 를 만든다. 이 z_t를 동일 회귀 헤드에 입력해 남은 토큰 수 y_rem(t)를 추정한다. 이렇게 단계별로 예측을 갱신함으로써, 샘플링 과정에서 발생하는 길이 변동성을 실시간으로 반영한다. PLP는 특히 강화학습 루프에서 여러 후보 응답을 동시에 생성해야 하는 경우, 정적 예측이 제공할 수 없는 적응성을 제공한다.

ForeLen 벤치마크는 기존 LMSYS와 달리 장기 시퀀스, 체인‑오브‑생각(Chain‑of‑Thought), 강화학습 샘플링 데이터를 포괄한다. 실험 결과 EGTP는 모든 모델(예: LLaMA‑7B, Falcon‑40B 등)에서 MAE를 평균 29.16 % 감소시켰으며, 기존 SSJF‑Reg 대비 55 % 이상 개선했다. 또한 EGTP와 PLP를 길이‑인식 스케줄러와 결합했을 때, 배치 내 패딩 비율이 크게 낮아져 전체 처리량이 1.8배~2.3배 향상되는 효과를 보였다.

이러한 결과는 (1) LLM 내부 정보를 재활용함으로써 예측 비용을 거의 제로에 가깝게 만든 점, (2) 엔트로피 기반 가중치가 중요한 토큰을 효과적으로 강조해 정적 예측 정확도를 크게 높인 점, (3) PLP가 동적 환경에서 실시간 길이 추정을 가능하게 해 강화학습과 같은 고불확실성 시나리오에 적용 가능하게 만든 점을 강조한다. 향후 연구는 EGTP와 PLP를 다양한 토큰화 방식, 멀티모달 모델, 그리고 초대형 LLM(>100B 파라미터)에도 확장하고, 스케줄러와의 통합 최적화를 통해 실제 서비스 환경에서의 비용 절감 효과를 정량화하는 방향으로 진행될 수 있다.