Scalable LLM Reasoning Acceleration with Low-rank Distillation

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 의 수학적 추론 능력은 종종 긴 체인 오브 씽킹 (Chain-of-Thought, CoT) 생성을 필요로 합니다. 이로 인해 다음과 같은 두 가지 주요 문제가 발생합니다.

• 효율성과 성능의 트레이드오프: 기존에 개발된 효율적인 추론 방법들 (희소성 기반의 프루닝, 양자화 등) 은 일반 언어 작업 (요약, 독해 등) 에서는 성능 저하가 적지만, 수학적 추론 작업에서는 성능이 급격히 저하됩니다.
  • 예시: Gemma 2 2B 모델에 CATS(희소 임계값 방법) 를 적용하면 언어 생성에는 영향이 없으나, GSM8K 정확도가 51.02% 에서 34.42% 로 떨어집니다. DeepSeek-R1-Distill-Qwen 1.5B 에 GRIFFIN(적응형 구조적 가지치기) 을 적용하면 MATH-500 정확도가 79.40% 에서 42.00% 로 반토막 납니다.
• 오류 누적: 긴 생성 과정에서 발생하는 작은 근사 오차들이 누적되어 전체 생성 경로를 잘못된 방향으로 이끄는 문제가 발생합니다.
• 계산 비용: 추론 단계 (Generation) 는 사전 채우기 (Prefill) 단계보다 계산 비용이 훨씬 많이 들며, CoT 가 길어질수록 이 비용은 기하급수적으로 증가합니다.

따라서, 추론 효율성을 높이면서도 수학적 추론 능력을 유지하거나 복원할 수 있는 방법이 절실히 필요합니다.

2. 방법론 (Methodology: Caprese)

저자들은 Caprese (CAPability REcovery with Scalable Efficiency) 라는 새로운 증류 (Distillation) 방법을 제안합니다. 이 방법은 효율적인 추론 알고리즘 (예: GRIFFIN, CATS) 으로 인해 발생하는 Feedforward (FF) 블록의 잔차 (Residual) 오차를 저랭크 (Low-rank) 선형 계층으로 학습하여 복원하는 방식입니다.

핵심 기술적 요소:

1. 잔차의 저랭크 구조 발견:

   • FF 블록의 희소화 (Sparse) 또는 압축 과정에서 발생하는 오차 (잔차) 를 분석한 결과, 이 오차들은 매우 작은 저랭크 (Low-rank) 구조를 가진다는 것을 발견했습니다.
   • Figure 2 에서 보듯, FF 출력의 일부만 유지하는 Top-k 방식에 저랭크 근사 (Rank 256 등) 를 추가하면 밀도를 높이는 것보다 훨씬 적은 파라미터로 오차를 효과적으로 줄일 수 있습니다.

2. 계층별 및 종단간 (E2E) 증류:

   • Layer-wise Distillation: 각 FF 블록의 압축된 출력과 원본 모델 출력 간의 오차를 저랭크 행렬 ($L \in \mathbb{R}^{D \times r}, R \in \mathbb{R}^{r \times D}$) 로 학습합니다. 이때 원본 가중치는 고정 (Frozen) 됩니다.
   • End-to-end (E2E) Distillation: 학습된 저랭크 계층들을 모델 전체에 적용하여 최종 임베딩 오차를 최소화하도록 미세 조정합니다.

3. 병렬 추론 (Parallel Inference):

   • Caprese 의 저랭크 계층은 기존 FF 가중치와 병렬로 계산되거나, 업/다운 프로젝션 행렬에 직접 연결 (Concatenation) 되어 추가 지연 시간 (Latency) 을 거의 발생시키지 않습니다.
   • 이는 LoRA 와 달리, 희소성 기반의 효율성을 해치지 않으면서도 추가 계산을 병렬화할 수 있는 구조적 이점이 있습니다.

4. 저비용 훈련:

   • 전체 모델 가중치를 수정하지 않고, 약 1% 의 추가 파라미터 (저랭크 계층) 만 학습합니다.
   • 2 만 개 (20K) 의 합성 수학 데이터만으로 충분한 성능 복원이 가능합니다.

3. 주요 기여 (Key Contributions)

1. 추론 능력 복원: 효율적인 추론 알고리즘 (GRIFFIN, CATS) 을 적용했을 때 손실된 수학적 추론 능력을 거의 완벽하게 복원하거나, 경우에 따라 원본 모델보다 더 높은 성능을 달성합니다.
2. 효율성 유지: 추가 파라미터가 적고 병렬 계산이 가능하여, 기존 효율화 방법의 지연 시간 감소 효과를 유지합니다 (Time-to-next-token 16% 이상 단축).
3. 활성화 파라미터 감소: Caprese 를 적용하면 활성 파라미터 (Active Parameters) 가 약 20 억 개 (2B) 감소하여 (Gemma 2 9B, Llama 3.1 8B 기준), 메모리 및 계산 부하를 크게 줄입니다.
4. 응답 길이 단축: Caprese 는 명시적인 정규화 없이도 더 간결한 응답을 생성하는 경향이 있어, 불필요한 토큰 생성을 줄여 추가적인 효율성을 제공합니다.

4. 실험 결과 (Results)

다양한 모델 (Llama 3.1/3.2, Gemma 2, DeepSeek-R1-Distill-Qwen 등) 과 작업 (MATH-500, GSM8K, AMC 2023 등) 에서 검증되었습니다.

• 수학적 추론 성능:
  • DeepSeek-R1-Distill-Qwen 7B: CATS 적용 시 AMC 2023 정확도가 75.00% → 62.50% 로 하락했으나, Caprese 를 적용하면 78.13% 로 원본보다 높은 성능을 회복했습니다.
  • Llama 3.2 3B Instruct: Best-of-N (여러 응답 중 선택) 방식에서 원본 모델 대비 Pass@100 을 7.0% 향상시키면서, 계산량은 15.8% 줄였습니다.
• 언어 작업 성능: 수학 데이터로만 증류되었음에도 CoQA, QASPER, XSum 등 일반 언어 작업에서는 성능 저하가 없었습니다.
• 지연 시간 (Latency): Qwen 2.5 14B 모델에서 전체 지연 시간을 16% 이상 단축했습니다. LoRA 와 결합한 GRIFFIN 은 지연 시간 절감 효과를 40% 이상 상쇄하는 반면, Caprese 는 그 효과를 유지했습니다.
• 응답 길이: CATS 는 긴 응답을 생성하는 반면, Caprese 는 원본 모델 수준 또는 그보다 짧은 응답을 생성하여 토큰 수를 최대 8.5% 감소시켰습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 효율적인 추론 (Efficient Inference) 과 복잡한 추론 능력 (Reasoning Capability) 사이의 간극을 해소하는 획기적인 접근법을 제시합니다.

• 실용성: 기존 모델의 가중치를 변경하지 않고 소량의 파라미터와 데이터만으로 배포 가능한 효율적 모델을 만들 수 있어, 실제 서비스 환경에 적용하기 용이합니다.
• 확장성: 추론 단계에서의 테스트 타임 스케일링 (Test-time Scaling) 이 중요해지는 시점에서, 긴 CoT 생성 시 발생하는 오류 누적과 계산 비용을 동시에 해결할 수 있는 길을 열었습니다.
• 미래 전망: Caprese 는 추론 영역뿐만 아니라 다른 도메인으로도 확장 가능하며, 저랭크 계층에 담긴 지식의 본질을 탐구함으로써 더 지능적이고 효율적인 LLM 추론의 미래를 제시합니다.

요약하자면, Caprese는 "효율성을 위해 성능을 포기해야 한다"는 기존 통념을 깨고, 저랭크 증류를 통해 효율성과 고성능 추론을 동시에 달성할 수 있음을 입증한 중요한 연구입니다.