하이브리드 xLSTM 아키텍처로의 효과적인 증류 기법

기존 Transformer 모델은 문장이 길어질수록 계산량이 제곱으로 늘어나 비용과 에너지가 많이 소모된다. 이 논문은 이를 선형적으로 계산되는 xLSTM 구조로 바꾸면서도 성능 저하를 최소화하는 방법을 제시하여, 더 저렴하고 빠른 AI 서비스 구현을 가능하게 한다.

Transformer의 Self-Attention은 모든 토큰 간의 관계를 계산하므로 시퀀스 길이 N에 대해 N²의 비용이 발생한다. 이는 긴 문맥 처리 시 메모리 부족과 속도 저하의 근본적인 원인이 된다. 이 논문은 이를 해결하기 위해 연산량이 시퀀스 길이에 선형적으로 비례하는 xLSTM 아키텍처를 대안으로 삼는다.

xLSTM의 핵심인 mLSTM은 행렬 형태의 상태를 유지하며 정보를 압축하고 업데이트한다. 여기에 최근 토큰들을 정확히 기억하는 Sliding Window Attention을 병렬로 배치하여, 전역적인 흐름과 국소적인 세부 사항을 동시에 포착하도록 설계했다. 이는 마치 사람이 책을 읽을 때 전체 줄거리를 파악하면서도 방금 읽은 문장의 단어들을 정확히 기억하는 것과 유사한 원리다.

단순히 구조만 바꾸는 것에 그치지 않고, 교사 모델인 Transformer의 지식을 단계적으로 전수한다. 특히 수학이나 코드처럼 어려운 작업에서는 각 분야에 특화된 전문가 학생 모델들을 먼저 만든 뒤, 이들의 가중치를 하나로 합치는 방식을 통해 성능 손실을 최소화하고 효율성을 극대화했다.

아키텍처는 mLSTM과 SWA를 병렬로 배치하고 데이터 기반 게이팅(Data-dependent Gating)으로 출력을 결합한다. mLSTM은 행렬 형태의 상태를 유지하며 긴 의존성을 잡고, SWA는 최근 512개 토큰의 세부 정보를 보존한다. mLSTM의 상태 업데이트는 다음과 같다. [이전 상태 S_{t-1}, 망각 게이트 f_t, 입력 게이트 i_t, 그리고 현재 입력의 키 k_t와 값 v_t를 입력으로] → [이전 상태에 망각 계수를 곱하고 새로운 정보의 외적(outer product)을 더하는 연산을 수행해] → [갱신된 행렬 상태 S_t를 얻고] → [과거의 정보를 압축하여 유지하면서 새로운 정보를 효율적으로 누적한다].

증류 1단계는 레이어별 Hidden-state 정렬이다. [교사 모델의 어텐션 출력값 h와 학생 모델의 하이브리드 블록 출력값 h_hat을 입력으로] → [두 값의 차이를 제곱하여 평균을 내는 MSE 연산을 수행해] → [오차 거리를 나타내는 스칼라 값을 얻고] → [이 값이 최소화되도록 학생 모델의 게이트 파라미터를 학습시킨다].

증류 2단계는 Sparse Knowledge Distillation이다. [교사 모델이 예측한 다음 토큰 확률 분포 중 상위 256개 값과 학생 모델의 예측 분포를 입력으로] → [두 분포 사이의 정보량 차이를 측정하는 KL Divergence 연산을 수행해] → [분포의 불일치 정도를 얻고] → [학생 모델이 교사의 추론 패턴을 그대로 복제하도록 전체 가중치를 미세 조정한다].

증류 3단계는 Expert Merging이다. [수학, 코드 등 각 도메인에 특화된 전문가 모델들의 가중치 theta_i와 각 도메인의 중요도 lambda_i를 입력으로] → [가중치와 중요도를 곱해 모두 더하는 선형 합산 연산을 수행해] → [통합된 모델 파라미터 theta_merge를 얻고] → [여러 분야의 능력을 동시에 갖춘 단일 범용 모델을 완성한다].

xLSTM-Llama3.1-8B-IT는 대부분의 벤치마크에서 교사 모델 성능의 90% 이상을 회복했으며, MT-bench에서는 교사보다 높은 선호도 점수를 기록했다. 특히 수학과 코드 생성 작업에서 기존 선형화 기법인 LoLCATs보다 압도적인 성능을 보이며 교사 모델과의 격차를 크게 줄였다.

추론 효율성 측면에서 시퀀스 길이 131K 토큰 기준, Transformer 대비 지연 시간(Latency)과 GPU 메모리 사용량을 약 50% 절감했다. 배치 사이즈 8 환경에서는 생성 처리량(Throughput)이 최대 4배까지 향상되었으며, 문맥 길이가 길어질수록 Transformer는 메모리 부족(OOM)으로 실행이 불가능해지는 반면 xLSTM은 안정적인 성능을 유지했다.

Win-and-Tie rate 지표 분석 결과, 제안된 하이브리드 모델은 다양한 허용 오차 수준에서 기존의 모든 선형화 모델들을 압도하는 Pareto-front를 형성했다. 이는 특정 작업에서의 우연한 성능 향상이 아니라, 전반적인 벤치마크 세트에서 교사 모델을 대체할 수 있는 신뢰성을 확보했음을 의미한다.

mLSTM은 Matrix-valued state S_t를 사용하며, Rank-1 outer product 업데이트를 통해 정보를 축적한다. 이는 기존 Linear Attention에 입력 및 망각 게이트를 추가한 구조로, 시퀀스 길이에 따른 메모리 사용량이 일정하게 유지되는 장점이 있다.

SWA 브랜치는 최근 W=512 토큰에 대해 Softmax Attention을 수행하며, 시퀀스 시작 부분의 4개 토큰을 Sink 토큰으로 고정하여 Attention 메커니즘의 수치적 안정성을 확보한다. 데이터 기반 게이트는 mLSTM과 SWA의 출력을 동적으로 혼합하여 최적의 문맥 정보를 생성한다.

가중치 전이 시 교사의 Embedding과 MLP 레이어는 그대로 재사용하고 Attention 레이어만 하이브리드 블록으로 교체한다. 이 과정에서 Full Fine-tuning(FFT) 방식이 LoRA와 같은 매개변수 효율적 기법보다 성능 회복에 훨씬 효과적임이 확인됐다.

Expert Merging 시 TIES-Merging과 같은 복잡한 기법 대신 단순 선형 결합만으로도 긍정적인 전이(Positive Transfer)가 발생했다. 이는 선형화된 아키텍처가 도메인 지식을 모듈화하여 학습하고 통합하는 데 매우 적합한 구조임을 시사한다.

Needle-in-a-Haystack과 같은 합성 롱컨텍스트 평가에서 여전히 성능 저하가 관찰되며, 전문가 모델 통합 시 도메인 간 간섭으로 인해 STEM 분야 등 일부 특정 작업에서 성능이 하락하는 문제가 존재한다.