RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference

이 논문은 기존 조기 종료 방법의 한계를 극복하고 추론 속도를 높이면서도 제로샷 성능을 유지하기 위해 유사 데이터의 중간 계층 정보를 검색하여 모델의 종료 시점을 안내하는 'RAEE'라는 강건한 검색 기반 조기 종료 프레임워크를 제안합니다.

핵심

이 논문은 RAEE(Robust Retrieval-Augmented Early Exit) 라는 새로운 기술을 소개합니다. 쉽게 말해, "**거대한 인공지능 **(LLM) 하는 방법입니다.

기존의 인공지능은 질문을 받으면 처음부터 끝까지 모든 과정을 거쳐서 답을 내놓습니다. 마치 아주 똑똑한 요리사가 간단한 토스트를 만들 때도 고급 스테이크를 요리하듯 모든 재료를 다 준비하고 긴 과정을 거치는 것과 비슷하죠. 이렇게 하면 정확하지만 시간과 전기가 많이 낭비됩니다.

RAEE 는 이 문제를 해결하기 위해 두 가지 핵심 아이디어를 사용합니다.

1. "이미 비슷한 문제를 해결한 친구를 찾아봐!" (검색 기반 학습)

기존의 '조기 종료 (Early Exit)' 기술들은 모델 스스로 "아, 내가 이제 답을 알겠다!"라고 판단하게 하려고 노력했습니다. 하지만 이 방법은 모델을 다시 가르치는 데 시간이 많이 걸리거나, 오히려 답을 틀릴 확률이 높아지는 문제가 있었습니다.

RAEE 는 완전히 다른 접근법을 취합니다.

• 상황: 새로운 질문 (예: "이 영화 리뷰는 긍정일까?") 이 들어옵니다.
• 행동: RAEE 는 거대한 데이터베이스에서 이 질문과 매우 비슷한 과거의 질문들을 찾아냅니다.
• 비유: 마치 시험을 볼 때, "어? 이 문제는 전에 내가 푼 문제랑 비슷하네? 그때는 3 단계에서 답이 나왔었어!"라고 기억해 내는 것과 같습니다.
• 결과: 비슷한 과거 사례들이 "3 단계에서 답이 나왔으니 여기서 멈춰도 돼"라고 알려주면, 모델은 굳이 마지막 단계까지 가지 않고 3 단계에서 멈추고 답을 냅니다.

2. "실수한 부분을 중간에 바로잡아!" (오류 수정 기능)

기존 기술들은 "빨리 끝내는 것"만 중요하게 여겨, 중간에 멈추면 정확도가 떨어질까 봐 걱정했습니다. 하지만 RAEE 는 놀라운 사실을 발견했습니다.

• 발견: 가끔은 모델이 마지막 단계까지 다 계산한 결과보다, 중간 단계에서 나온 답이 더 정확할 때가 있습니다.
• 비유: 수학 문제를 풀 때, 마지막에 계산 실수를 해서 엉뚱한 답을 낸 경우, 중간에 계산한 단계가 사실은 정답에 더 가까웠던 경우를 생각해 보세요.
• RAEE 의 역할: RAEE 는 "이 질문과 비슷한 과거 문제들을 봤을 때, 중간 단계에서 정답을 맞춘 경우가 많았어"라고 알려줍니다. 그래서 모델이 마지막까지 계산해서 틀릴 뻔한 것을, 중간 단계에서 정답을 찾아내게 **보정 **(Correct)해 줍니다.

요약: RAEE 가 가져온 변화

1. **더 빠름 **(Efficiency) 불필요한 계산을 건너뛰기 때문에 전기를 아끼고 응답 속도가 빨라집니다.
2. **더 정확함 **(Accuracy) 단순히 빨리 끝내는 게 아니라, 과거의 성공 사례를 참고해서 틀릴 뻔한 답을 중간에 고쳐주기 때문에, 오히려 원래 모델보다 더 잘 맞출 때도 있습니다.
3. **학습 불필요 **(No Training) 모델을 다시 가르치는 (Fine-tuning) 과정이 필요 없습니다. 이미 있는 데이터베이스를 활용하기만 하면 됩니다.

한 줄 요약:
RAEE 는 거대한 인공지능이 매번 처음부터 끝까지 고민할 필요 없이, "비슷한 과거 사례를 검색해서 가장 효율적인 시점에 멈추고, 틀릴 뻔한 답을 중간에 고쳐주는 똑똑한 비서"입니다.

기술 요약

1. 문제 정의 (Problem Statement)

대규모 언어 모델 (LLM) 의 추론 (Inference) 은 높은 계산 비용과 메모리 요구 사항으로 인해 배포가 어렵습니다. 이를 해결하기 위해 Early Exit(조기 종료) 기법이 제안되었으나, 기존 방법론들은 다음과 같은 한계를 가집니다.

• 학습 기반 (Training-based): 내부 분류기 (Internal Classifier) 와 백본 모델을 함께 최적화하여 막대한 학습 오버헤드가 발생합니다.
• 반-학습 기반 (Semi-training): 백본은 고정하고 가벼운 분류기만 학습하지만, 수동 특징 공학에 의존하며 일반화 성능이 낮을 수 있습니다.
• 무학습 (Training-free): 휴리스틱 기준 (예: 엔트로피 임계값) 을 사용하지만 적응력이 부족하여 전체 모델 대비 성능이 저하되는 경우가 많습니다.
• 핵심 문제: 기존 Early Exit 방법들은 대부분 속도 향상과 정확도 저하 사이의 트레이드오프를 강요합니다. 즉, 더 빨리 종료하려면 정확도가 떨어집니다.

2. 방법론 (Methodology: RAEE)

저자들은 Early Exit 문제를 단순한 속도 최적화 문제가 아닌, 분포 예측 (Distribution Prediction) 문제로 재정의하고 검색 증강 (Retrieval-Augmented) 방식을 도입하여 RAEE를 제안했습니다.

2.1 핵심 통찰 (Key Observations)

1. 수정 메커니즘 (Corrective Mechanism): 전체 모델을 다 실행했을 때 틀린 예측을 하더라도, 중간 레이어에서 올바른 예측을 할 수 있는 경우가 많습니다. 즉, Early Exit 은 단순히 속도를 높이는 것을 넘어, 모델의 잘못된 예측을 **수정 (Correct)**할 수 있는 잠재력을 가집니다.
2. 유사 데이터의 일관된 행동: 의미적으로 유사한 데이터는 최적의 종료 레이어 (Exit Layer) 에서도 매우 일관된 행동을 보입니다.

2.2 RAEE 프레임워크 구조

RAEE 는 별도의 분류기 학습 없이 외부 데이터베이스를 활용하여 최적의 종료 레이어를 결정합니다.

1. 검색 데이터베이스 구축 (Build Phase):

   • 훈련 데이터에 대해 백본 모델을 실행하며, 각 레이어에서의 예측 확률과 정답 여부를 기록합니다.
   • Key: 입력 데이터의 임베딩 (Encoder 또는 백본 모델의 임베딩 사용).
   • Value: 해당 입력에 대해 **정확한 예측 (Correct Prediction)**이 이루어진 레이어 인덱스와 그 확률의 집합.
   • FAISS 등의 근사 최근접 이웃 (ANN) 검색 인덱스를 사용하여 효율적인 저장 및 검색 구조를 만듭니다.

2. 검색 증강 추론 (Inference Phase):

   • 새로운 입력 $x$가 들어오면, 먼저 임베딩을 생성합니다.
   • 검색 데이터베이스에서 $k$개의 가장 유사한 이웃 (Top-k Nearest Neighbors) 을 찾습니다.
   • 유사한 이웃들이 정확한 예측을 했던 레이어들의 확률 분포를 집계하여, 현재 입력 $x$가 가장 잘 종료될 확률이 높은 레이어를 예측합니다.
   • 종료 결정: 예측된 확률 분포가 최대가 되는 레이어에서 추론을 중단하고, 해당 레이어의 출력을 최종 예측 헤드로 전달합니다.

3. 주요 기여 (Key Contributions)

1. Early Exit 의 새로운 관점: Early Exit 문제를 분포 예측 문제로 모델링하고, 유사 데이터의 종료 정보를 통해 이를 근사할 수 있음을 증명했습니다.
2. RAEE 프레임워크 제안: 분류기 학습 없이 외부 검색 데이터베이스를 활용하여 Early Exit 을 유도하는 강건한 (Robust) 프레임워크를 제시했습니다.
3. 성능 및 효율성 동시 달성: 기존 Early Exit 방법론들이 겪던 정확도 저하 문제를 해결했습니다. 실험 결과, RAEE 는 추론 속도를 높이면서도 전체 모델 (Full Model) 보다 높은 정확도를 달성하거나, 기존 방법론들을 압도하는 성능을 보였습니다.

4. 실험 결과 (Experimental Results)

저자는 GLUE 벤치마크의 8 가지 다운스트림 태스크 (감정 분석, 문법 판단 등) 와 Llama-3-8B, Gemma-7B 와 같은 대규모 모델에 대해 실험을 수행했습니다.

• 정확도 (Accuracy):
  • RoBERTa-Large 기반에서 RAEE 는 평균 정확도 63.41을 기록하여, 비교 대상 방법론들 (HashEE, DeeBERT, AdaInfer 등) 보다 월등히 높았습니다.
  • 특히 Llama-3-8B 와 Gemma-7B 에서 RAEE 는 베이스라인 모델보다 약 15~20%p 이상 정확도가 향상되었습니다. 이는 중간 레이어의 정보를 활용해 전체 모델이 놓친 정답을 찾아냈기 때문입니다.
• 추론 지연 (Latency):
  • 백본 모델의 크기가 클수록 (수십 억 파라미터 이상) RAEE 의 속도 향상 효과가 두드러졌습니다.
  • Llama-3-8B 의 경우, AdaInfer 보다 속도는 느릴 수 있으나, AdaInfer 가 베이스라인과 유사한 성능만 내는 반면 RAEE 는 성능을 크게 향상시키면서 지연 시간을 약 50% 감소시켰습니다.
• 외부 도메인 (Out-of-Domain) 성능:
  • WikiText 데이터로 구축된 검색 DB 를 사용하여 CNN/DailyMail 및 XSum 요약 태스크에서도 성능 향상과 조기 종료를 동시에 달성했습니다.

5. 의의 및 결론 (Significance & Conclusion)

이 논문은 Early Exit 기술의 패러다임을 변화시킵니다.

• 트레이드오프의 극복: "속도를 내면 정확도가 떨어진다"는 기존 상식을 깨고, 검색 기반의 지능형 종료 전략을 통해 속도와 정확도를 동시에 개선했습니다.
• 학습 비용 절감: 별도의 분류기 학습 (Fine-tuning) 이 필요 없어, 기존 모델에 즉시 적용 가능한 비용 효율적인 솔루션을 제공합니다.
• 오류 수정 능력: Early Exit 을 단순한 가속화 도구가 아닌, 모델의 예측 오류를 보정하는 동적 수정 메커니즘으로 활용함으로써 모델의 신뢰성을 높였습니다.

결론적으로 RAEE 는 대규모 언어 모델의 효율적인 배포를 위한 강력한 프레임워크로, 추론 비용 절감과 성능 향상을 동시에 실현할 수 있는 새로운 방향성을 제시합니다.