프리트레인된 소프트맥스 어텐션 트랜스포머를 효율적인 하이브리드 구조로 증류하기 위한 레이어 선택 전략
📝 원문 정보
- Title:
- ArXiv ID: 2512.20569
- 발행일:
- 저자: Unknown
📝 초록 (Abstract)
프리트레인된 소프트맥스 어텐션 트랜스포머를 소프트맥스와 선형 어텐션 레이어를 교차 배치한 효율적인 하이브리드 구조로 증류하는 방법은, 대규모 사전학습 없이도 대형 언어 모델(LLM)의 추론 효율성을 크게 향상시킬 수 있는 유망한 접근이다. 핵심은 어느 레이어를 선형 어텐션 변형으로 변환할지 결정하는 레이어 선택이며, 본 논문은 일반 텍스트 데이터에 대한 소량 학습으로 얻은 레이어 중요도 점수를 활용한 간단하고 효율적인 레시피를 제시한다. 선택된 레이어에 대해서는 최신 증류 파이프라인(RADLADS; Goldstein et al., 2025)을 적용해 어텐션 가중치 전이, 은닉 상태 정렬, KL 기반 분포 매칭을 수행하고 짧은 파인튜닝을 진행한다. 실험 결과, 고정 비율로 선형 어텐션을 균등하게 삽입하는 휴리스틱이나 특수 진단 데이터셋을 이용한 복잡한 방법보다 본 접근법이 더 우수함을 확인하였다.💡 논문 핵심 해설 (Deep Analysis)
이 논문은 현재 대형 언어 모델(LLM)의 추론 비용이 급격히 증가함에 따라, 사전 학습된 거대한 모델을 그대로 사용하기보다는 구조적 경량화를 시도하는 연구 흐름에 중요한 기여를 하고 있다. 핵심 아이디어는 ‘소프트맥스 어텐션’과 계산량이 훨씬 적은 ‘선형 어텐션’을 레이어 단위로 혼합하여, 전체 모델의 연산 복잡도를 낮추면서도 성능 저하를 최소화하는 것이다. 여기서 가장 어려운 문제는 어떤 레이어를 선형 어텐션으로 교체할지 결정하는 ‘레이어 선택’이다. 기존 연구들은 단순히 일정 비율(예: 1/4, 1/2)로 레이어를 교체하거나, 특정 작업에 특화된 진단 데이터셋을 활용해 레이어 중요도를 평가하는 복잡한 절차를 사용했다. 그러나 이러한 방법은 모델 구조와 데이터 도메인에 따라 성능 변동이 크고, 진단 데이터셋을 별도로 구축해야 하는 비용이 발생한다.본 논문은 이러한 한계를 극복하기 위해 ‘소량의 일반 텍스트 데이터’를 이용해 각 레이어의 중요도를 빠르게 추정한다. 구체적으로는 레이어별 출력 변화량이나 손실 기여도를 측정해 스코어를 산출하고, 이 스코어가 낮은 레이어를 선형 어텐션으로 교체한다. 이 과정은 전체 모델을 재학습하지 않아도 되며, 몇 백 배치 정도의 짧은 학습만으로 충분히 신뢰할 만한 점수를 얻을 수 있다. 따라서 실무에서 모델을 빠르게 경량화하고자 할 때, 별도의 도메인 특화 데이터나 복잡한 분석 파이프라인 없이도 적용 가능하다.
선택된 레이어에 대해 적용된 증류 파이프라인은 RADLADS(Robust Attention Distillation with Linear Attention and Distribution Shaping)이다. 이 파이프라인은 네 단계로 구성된다. 첫째, 원본 소프트맥스 어텐션의 가중치를 선형 어텐션에 직접 전달함으로써 초기화 단계에서 정보 손실을 최소화한다. 둘째, 두 모델의 은닉 상태를 정렬(alignment)시켜 내부 표현을 일치시키고, 셋째, KL 발산을 이용해 출력 분포를 맞춤으로써 확률적 예측 특성을 보존한다. 마지막으로 짧은 파인튜닝을 수행해 남은 미세 차이를 보정한다. 이러한 다단계 접근은 단순히 파라미터를 복사하거나 한 단계만 수행하는 기존 증류 방식보다 훨씬 안정적이며, 특히 선형 어텐션이 근본적으로 다른 연산 특성을 가지기 때문에 발생할 수 있는 성능 저하를 효과적으로 억제한다.
실험 결과는 두 가지 관점에서 의미가 있다. 첫째, 동일한 변환 비율을 유지하면서도 레이어 중요도 기반 선택이 무작위 혹은 균등 선택보다 일관되게 높은 정확도와 낮은 퍼플렉시티를 기록했다. 둘째, 특수 진단 데이터셋을 활용한 복잡한 선택 방법과 비교했을 때, 일반 텍스트 기반 점수는 거의 동일하거나 더 나은 성능을 보이며, 데이터 준비 비용을 크게 절감했다. 이는 레이어 선택이 반드시 복잡한 도메인 지식이나 맞춤형 평가 지표에 의존하지 않아도 된다는 점을 실증적으로 보여준다.
전체적으로 이 연구는 (1) 레이어 선택을 위한 간단하고 재현 가능한 방법론, (2) 최신 증류 파이프라인과의 시너지 효과, (3) 실무 적용성을 고려한 비용 효율성을 동시에 달성했다는 점에서 큰 의의를 가진다. 앞으로는 다양한 모델 규모와 다른 종류의 선형 어텐션(예: Performer, Linformer)에도 동일한 프레임워크를 적용해볼 수 있을 것이며, 멀티모달 모델이나 도메인 특화 LLM에도 확장 가능성이 기대된다.