대규모 언어 모델 추론 가속화: 자기 지도 학습 조기 종료의 힘

초록

본 논문은 대규모 언어 모델(LLM)의 추론 비용을 줄이기 위해 트랜스포머 중간 레이어에 ‘조기 종료 헤드’를 추가하는 모듈식 방법을 제안합니다. 각 헤드는 메인 모델의 예측을 모방하도록 자기 지도 학습 방식으로 훈련되며, 교정된 신뢰도 임계값에 도달하면 계산을 조기에 중단합니다. 여러 신뢰도 지표 중 엔트로피가 정오답 구분에 가장 효과적임을 보였으며, Pythia 모델(70M~2.8B 파라미터) 실험을 통해 추론 비용을 크게 줄이면서도 정확도를 유지함을 입증했습니다. 또한, 이 방식을 추측 디코딩에 적용한 ‘동적 자기 추측 디코딩(DSSD)‘을 소개하며, 최소한의 하이퍼파라미터 튜닝으로 기존 방법보다 1.66배 높은 토큰 수용률을 달성했습니다.

상세 분석

이 논문의 핵심 기술적 기여는 기존 LLM의 추론 비효율성, 즉 모든 토큰 생성에 동일한 계산 비용이 소모되는 문제를 ‘조기 종료(Early Exit)’ 메커니즘으로 해결한 점에 있습니다. 저자들은 사전 학습된 모델의 가중치를 동결한 상태에서, 중간 레이어(예: 총 레이어 수 L의 1/5, 2/5, 3/5, 4/5 지점)에 간단한 MLP 헤드를 추가합니다. 이 헤드들은 메인 모델의 최종 출력을 ‘교사’로 삼아 자기 지도 학습 방식(목적 함수 L_self, 즉 KL 발산)으로 훈련됩니다. 이는 추가 데이터 없이 모델 자체의 지식을 활용한다는 점에서 실용적입니다.

가장 주목할 만한 통찰은 ‘교정(Calibration)’ 프로세스와 신뢰도 지표 선택에 있습니다. 저자들은 최대 확률(Max Probability), 엔트로피(Entropy), 상위 두 확률 차이(Top-2 Difference) 등 세 가지 지표를 비교한 결과, 모든 모델 크기(70M~2.8B)에서 엔트로피가 가장 높은 AUC(ROC 곡선 아래 면적)를 보여주며 정답과 오답 예측을 가장 잘 구분함을 실험적으로 증명했습니다. 엔트로피는 예측 분포의 불확실성을 정량화하므로, “이해하기 쉬운” 토큰은 낮은 엔트로피(높은 신뢰도)를, “어려운” 토큰은 높은 엔트로피(낮은 신뢰도)를 보일 가능성이 높습니다. 사용자가 설정한 최소 정확도 임계값 ε에 따라 교정 데이터셋에서 이 엔트로피 임계값을 동적으로 결정함으로써, 속도와 정확도 간의 trade-off를 체계적으로 제어할 수 있습니다.

이 방법론을 ‘추측 디코딩(Speculative Decoding)‘에 적용한 ‘동적 자기 추측 디코딩(DSSD)‘은 또 다른 중요한 확장입니다. 기존 LayerSkip 방식이 어떤 중간 레이어를 드래프트 모델로 사용할지 미리 고정하고 최적의 추론 길이를 수동으로 탐색해야 하는 번거로움을 해결합니다. DSSD는 각 생성 단계에서 현재 토큰의 예측 신뢰도(엔트로피)를 기반으로 동적으로 종료 레이어를 선택합니다. 즉, 쉬운 토큰은 더 얕은 레이어에서, 어려운 토큰은 더 깊은 레이어(또는 최종 레이어)에서 예측을 시도합니다. 이로 인해 단일 정확도 임계값 ε만으로 조정 가능하며, LayerSkip 대비 1.66배 높은 토큰 수용률을 달성했습니다. 이는 동일한 계산 예산으로 더 많은 토큰을 효율적으로 검증할 수 있음을 의미합니다.

한계점으로는, 조기 종료 헤드 훈련과 교정에 추가적인 계산 비용(비록 추론 비용에 비하면 미미하지만)이 필요하며, 매우 짧은 시퀀스나 모든 토큰이 높은 불확실성을 보이는 특수한 텍스트 유형에서는 효과가 제한될 수 있습니다. 또한, 실험이 Pythia 모델군에 국한되어 있으며, 최신 초대규모 모델(예: 100B 이상)이나 다른 아키텍처에서의 일반화 가능성은 추가 검증이 필요합니다. 향후 연구 방향으로는 다중 토큰에 대한 조기 종료 결정, 다양한 작업(번역, 요약 등)에 대한 평가, 그리고 하드웨어 가속기와의 통합 최적화 등이 있을 것입니다.