1. 문제 정의 (Problem)

현재 Transformer 기반의 대규모 언어 모델 (LLM) 은 Softmax Attention 메커니즘에 의존하고 있습니다. 이는 다음과 같은 심각한 한계를 가지고 있습니다:

• 계산 및 메모리 병목: Softmax Attention 은 시퀀스 길이 ($L$) 에 대해 이차 ($O(L^2)$) 복잡도를 가지며, 생성 (Inference) 과정에서 Key-Value (KV) 캐시가 선형적으로 증가합니다. 이로 인해 긴 문맥 (Long-context) 처리 시 메모리 부족과 계산 비용이 급증합니다.
• 기존 선형화 방법의 한계: 기존에 제안된 선형 (Linear) 또는 이차 미만 (Sub-quadratic) 어텐션 메커니즘 (예: Mamba, RWKV 등) 은 처음부터 (from scratch) 학습할 경우 학습 비용이 막대하며, 맥락 학습 (In-context learning) 및 회상 (Recall) 능력이 Transformer 에 비해 현저히 떨어집니다.
• 기존 선형화 (Linearization) 기법의 결함: 사전 학습된 Transformer 를 선형 구조로 변환하는 기존 방법들 (LoLCATs, Liger 등) 은 두 가지 주요 문제를 겪습니다:
  1. 성능 저하: Teacher 모델의 성능을 유지하지 못하며 (MMLU 에서 큰 점수 하락), 적응형 메모리 제어 메커니즘이 부족합니다.
  2. 길이 외삽 (Length Extrapolation) 실패: 고정된 위치 인코딩 (RoPE) 을 그대로 사용하여 학습 시퀀스 길이를 초과하는 문맥에서는 성능이 급격히 떨어집니다.

2. 방법론 (Methodology)

Lizard는 사전 학습된 Transformer 를 이차 미만 (Sub-quadratic) 아키텍처로 변환하면서도 Teacher 모델의 성능을 거의 손실 없이 복원하는 프레임워크입니다. 핵심 구성 요소는 다음과 같습니다.

가. 하이브리드 어텐션 메커니즘

Lizard 는 전역적 (Global) 인 문맥 압축과 국소적 (Local) 인 정밀도를 모두 확보하기 위해 두 가지 메커니즘을 결합합니다.

1. 게이팅 선형 어텐션 (Gated Linear Attention, GLA):
   • 학습 가능한 게이팅 모듈: 고정된 위치 인코딩 (RoPE) 대신 데이터에 적응적인 게이팅 ($\Gamma_i = \sigma(W_\gamma x_i)$) 을 도입합니다. 이는 과거 토큰의 감쇠 (Decay) 를 제어하여 적응형 메모리 제어와 길이 외삽을 가능하게 합니다.
   • 상수 메모리 추론: KV 캐시 전체를 저장하지 않고, 은닉 상태 ($S_i$) 를 점진적으로 업데이트하여 메모리 사용량을 상수 ($O(1)$) 로 유지합니다.
2. 앵커 윈도우 어텐션 (Anchor Window Attention):
   • GLA 가 잃을 수 있는 국소적 정밀도를 보완하기 위해 슬라이딩 윈도우 어텐션을 사용합니다.
   • 메타 메모리 토큰 (Meta-memory tokens): 학습 가능한 스칼라 파라미터 ($t_j$) 를 어텐션 분포의 분모에 추가하여, 불필요한 국소 정보에 대한 어텐션 가중치를 흡수 (Sink) 시킵니다. 이는 Massive Activations 현상을 완화하고 국소 문맥의 충실도를 유지합니다.

나. 2 단계 학습 전략

1. 1 단계 (Attention Approximation): Teacher 모델의 Softmax 어텐션 출력과 RoPE 패턴을 모방하도록 학습합니다. 게이팅 파라미터, 특징 맵 (Feature maps), 메타 메모리 토큰을 Teacher 모델의 어텐션 출력과 최소제곱법 (MSE) 으로 정렬합니다.
2. 2 단계 (Language Modeling Alignment): 선형화된 어텐션 레이어로 교체된 전체 모델을 표준 자기회귀 (Autoregressive) 언어 모델링 목적 함수로 미세 조정 (Fine-tuning) 합니다. 이때 LoRA (Low-Rank Adaptation) 를 사용하여 파라미터 효율성을 높입니다.

다. 하드웨어 인식 알고리즘 (Hardware-Aware Algorithm)

• 수치적 불안정성 해결: 기존 게이팅 선형 어텐션의 병렬 계산 시 누적 곱 (Cumulative product) 이 저정밀도 (bfloat16) 에서 언더플로우 (Underflow) 를 일으켜 Full Precision (float32) 연산을 강제하는 문제가 있었습니다.
• 로그 공간 재파라미터화: Hedgehog 특징 맵의 비음수 성질을 활용하여 게이팅 항을 로그 공간으로 이동시켜 Query 와 Key 프로젝션에 흡수시킵니다. 이를 통해 Tensor Core를 활용한 표준 GEMM 연산이 가능해지고, 저정밀도 연산이 안정적으로 수행되어 학습 속도가 크게 향상됩니다.

3. 주요 기여 (Key Contributions)

1. Lizard 프레임워크 제안: 사전 학습된 Transformer 를 이차 미만 아키텍처로 변환하여 거의 손실 없는 (Near-lossless) 성능 회복을 달성했습니다.
2. 성능 및 효율성 동시 달성:
   • MMLU 벤치마크: 기존 선형화 방법 (LoLCATs, Liger) 보다 9.4~24.5 점 더 높은 성능을 기록했습니다.
   • 하이브리드 설정: 원본 Softmax 레이어의 50% 만 유지하는 하이브리드 구성에서도 Teacher 모델 (LLaMA-3-8B) 과 유사한 성능 (MMLU 65.1 vs 66.6) 을 보였습니다.
3. 강력한 길이 외삽 및 회상 능력: 학습 시퀀스 길이 (2K) 를 훨씬 초과하는 64K 토큰까지 안정적인 성능을 보이며, "Needle-in-a-Haystack" 및 RULER 벤치마크에서 Teacher 모델과 유사한 회상 정확도를 달성했습니다.
4. 학습 속도 향상: 하드웨어 인식 알고리즘을 통해 학습 처리량 (Throughput) 을 **32%**까지 향상시켰습니다.

4. 실험 결과 (Results)

• 언어 이해 벤치마크: Mistral-7B 와 LLaMA-3-8B Teacher 모델을 기반으로 한 실험에서, Lizard 는 기존 Sub-quadratic 모델들 (Mamba, RWKV 등) 과 선형화 모델들보다 모든 벤치마크 (PiQA, ARC, MMLU 등) 에서 압도적으로 우수한 성능을 보였습니다. 특히 MMLU 5-shot 에서 LLaMA-3-8B-Lizard 는 72.4 점 (Teacher: 72.0) 을 기록했습니다.
• 긴 문맥 회상 (Long-Context Recall):
  • Needle-in-a-Haystack: 2K~8K 길이의 시퀀스에서 완벽한 회상 정확도를 보였으며, LoLCATs 는 학습 윈도우를 초과하면 실패한 것과 대조적입니다.
  • RULER: 32K 문맥 길이에서도 Teacher 모델과 유사한 회상 정확도 (81.3 vs 80.5) 를 유지했습니다.
• 생성 효율성: 32K 시퀀스 길이에서도 메모리 부족 (OOM) 없이 상수 메모리와 안정적인 처리량을 유지하는 반면, FlashAttention-2 를 사용하는 Teacher 모델은 OOM 에 직면했습니다.

5. 의의 및 결론 (Significance)

Lizard 는 **고성능 (Transformer 기반)**과 고효율 (Sub-quadratic) 사이의 간극을 해소하는 실용적인 솔루션을 제시합니다.

• 확장성: 기존에 수조 개의 토큰으로 학습된 최첨단 LLM 들을, 긴 문맥 처리가 가능하고 메모리 효율적인 인ference 엔진으로 변환할 수 있는 경로를 제공합니다.
• 실용적 가치: 제한된 하드웨어 자원 환경에서도 무한한 문맥 (Infinite Context) 을 처리할 수 있게 하여, LLM 의 배포 가능성을 크게 확장합니다.
• 한계점: 여전히 고품질의 사전 학습된 Teacher 모델에 의존하며, 완전한 어텐션에 비해 회상 (Recall) 과 메모리 간의 트레이드오프는 완전히 제거되지 않았습니다.

요약하자면, Lizard 는 적응형 게이팅 메커니즘과 하드웨어 최적화 알고리즘을 통해 기존 Transformer 의 지능을 유지하면서 계산 및 메모리 효율성을 획기적으로 개선한 획기적인 프레임워크입니다.