PARTREP: Decoder-only LLM을 위한 반복 대상 학습

원문 발행 2026. 07. 02. 16:07원문 언어 영어AI 한국어 번역arXiv cs.CL (NLP) 원문 보기

요약

Decoder-only LLM의 비대칭적 정보 흐름 문제를 해결하기 위해 프롬프트의 일부 토큰만 선택적으로 반복하는 PartRep 기법을 제안합니다. NLL 기반의 경량 게이트를 통해 정보량이 많은 토큰만 증강함으로써, 성능은 유지하면서 KV 캐시와 연산 비용을 획기적으로 절감합니다.

핵심 포인트

인과적 어텐션의 정보 비대칭성 문제를 프롬프트 반복으로 해결
전체 반복 대신 정보량이 높은 토큰만 선택하는 PartRep 방식 제안
NLL 기반의 경량 게이트를 활용한 효율적인 토큰 선택 메커니즘
KV 캐시 사용량 59.4%, 프리필 FLOPs 79.0% 수준으로 절감
Qwen2.5, Llama3.2, Gemma4 등 다양한 모델에서 성능 검증

Decoder-only LLM은 방대한 양의 자연어 작업에서 탁월한 성능을 보이지만, 인과적 어텐션 (Causal Attention)으로 인해 발생하는 비대칭적 정보 흐름 문제로 어려움을 겪습니다. 즉, 뒤쪽 토큰이 앞쪽 토큰보다 문맥적 근거 (Contextual Grounding)가 더 풍부합니다. 간단하고 효과적인 해결책은 프롬프트 반복 (Prompt Repetition)입니다. 생성 전에 프롬프트의 두 번째 복사본을 단순히 추가하는 것만으로도 위치 전반에 걸쳐 근거를 재분배하고 추론 성능을 향상할 수 있습니다. 그러나 원본 프롬프트를 전체 반복하는 것은 KV 캐시 (KV Cache) 점유율을 두 배로 늘리고, 프리필 (Prefill) 단계에서의 어텐션 비용을 네 배로 증가시켜 긴 문맥 (Long-context) 설정에서는 비실용적입니다.

우리는 전체 프롬프트 대신 가장 정보량이 많은 토큰만을 추가하는 선택적 증강 방법인 PartRep을 제안합니다. 우리는 토큰 단위의 음의 로그 가능도 (Negative Log-Likelihood, NLL)를 선택 신호로 사용합니다. 이는 예측 가능성이 낮은 토큰일수록 주변 문맥으로부터 복구하기 어렵기 때문에, 뒤쪽 위치의 반복으로부터 더 많은 이득을 얻을 것이라는 가설에 기반합니다. 점수 산출을 위한 전체 순전파 (Forward Pass)의 막대한 비용을 피하기 위해, 우리는 초기 레이어의 은닉 상태 (Hidden States)로부터 높은 NLL 토큰을 예측하는 경량 게이트 (Gate)를 학습시켜, 미드-프리필 (Mid-prefill) 단계에서 조기 종료 (Early Exit)를 통한 토큰 선택을 가능하게 합니다. 8개의 벤치마크 (MMLU, GSM8K, RULER 포함)와 3개의 모델 제품군 (Qwen2.5, Llama3.2, Gemma4)에 걸쳐 테스트한 결과, PartRep은 전체 반복 방식의 이점을 대부분 유지하면서도 KV 캐시의 59.4%와 프리필 FLOPs의 79.0%만을 사용합니다.

AI 자동 생성 콘텐츠

원문 바로가기

PARTREP: Decoder-only LLM을 위한 반복 대상 학습

요약

핵심 포인트

댓글