기술 요약: LT2: 선형 시간 루프형 트랜스포머

1. 문제 제기

루프형 트랜스포머 (LT, Universal Transformers 라고도 함) 는 디코딩 전에 여러 반복 단계 ($T$) 에 걸쳐 동일한 가중치를 재사용함으로써 파라미터 효율적인 아키텍처로 등장했습니다. 이는 파라미터 수를 증가시키지 않으면서 모델의 깊이를 $T \times N$ (여기서 $N$은 물리적 레이어 수) 으로 효과적으로 증가시킵니다. 그러나 표준 LT 는 완전한 2 차 softmax 어텐션에 의존합니다. 결과적으로 계산 비용 (FLOPs) 과 추론 메모리 (KV 캐시) 는 시퀀스 길이 ($L^2$) 에 따라 2 차적으로 증가하며, 모든 루프 반복마다 복합적으로 증가합니다. 이는 LT 를 계산적으로 비싸고 긴 컨텍스트 시나리오에서는 실용적이지 않게 만들어, 파라미터 공유의 이점을 어텐션 메커니즘의 비효율성이 상쇄하는 병목 현상을 초래합니다.

2. 방법론

저자들은 루프형 블록 내의 2 차 softmax 어텐션을 2 차 미만의 토큰 믹싱 원시 연산으로 대체하는 LT2(선형 시간 루프형 트랜스포머) 아키텍처 계열을 소개합니다. 핵심 방법론은 다음 세 가지 주요 구성 요소를 포함합니다:

2.1. 아키텍처 변형

LT2 는 멀티헤드 어텐션 (MHA) 서브레이어를 선형 또는 희소 어텐션 메커니즘으로 대체합니다. 저자들은 두 가지 주요 변형을 연구합니다:

• LT2-linear: MHA 를 선형 어텐션 메커니즘 (예: GDN, KDA, Mamba2, RetNet) 으로 대체합니다. 이들은 랭크 1 또는 대각선 + 저랭크 (DPLR) 연산을 통해 업데이트되는 재귀 상태 $S_t$를 유지하여 $O(L)$ 복잡도를 달성합니다.
• LT2-sparse: MHA 를 희소 어텐션 메커니즘 (예: 윈도우 어텐션, DSA, NSA) 으로 대체합니다. 이들은 어텐션을 로컬 윈도우나 선택된 토큰의 부분 집합으로 제한하여 2 차 미만의 복잡도를 달성합니다.

2.2. 루핑의 이론적 시너지

이 논문은 루핑이 2 차 미만의 믹서를 두 가지 방식으로 독특하게 향상시킨다고 주장합니다:

1. 선형 어텐션에서의 랭크 확장: 단일 DPLR 선형 어텐션 블록은 재귀 상태에 랭크 1 업데이트를 수행합니다. $T$번 루핑될 때, 누적 상태 전이 연산자는 이러한 인자 $T$개의 곱이 됩니다. 이론적으로 루프 반복 간 키 (keys) 가 직교한다면, 메모리 업데이트의 유효 랭크는 1 에서 $T$로 확장되어 모델의 표현력과 상태 추적 능력을 크게 향상시킵니다.
2. 희소 어텐션에서의 수용 영역 확장: 윈도우 크기 $w$를 가진 단일 희소 블록은 $O(w)$의 수용 영역을 가집니다. 블록을 루핑하면 각 반복마다 정보가 더 멀리 전파될 수 있습니다. 귀납적으로 $T$번의 루프 후 유효 수용 영역은 $O(Tw)$로 확장되어 파라미터를 증가시키지 않고 연산을 컨텍스트로 변환합니다.

2.3. 하이브리드 아키텍처 (LT2-hybrid)

효율성과 성능을 균형 있게 맞추기 위해 저자들은 루프형 설정 내에서 다양한 믹서를 하이브리드화하는 방법을 탐구합니다. 두 가지 구체적인 구성을 조사합니다:

• LT2-hybrid (GDN+DSA): 게이트드 델타 네트워크 (GDN, 선형) 와 동적 희소 어텐션 (DSA, 희소) 을 교차 배치합니다. 이 변형은 완전 어텐션 레이어가 전혀 포함되지 않습니다.
• LT2-hybrid (Full+GDN): 소량의 완전 어텐션 레이어를 GDN 레이어와 교차 배치합니다. 이는 효율성 향상의 대부분을 유지하면서 품질을 극대화하는 것을 목표로 합니다.

2.4. 증류 전략

이 논문은 처음부터 학습하지 않고 사전 학습된 완전 어텐션 루프 트랜스포머 (특히 Ouro) 를 LT2 모델로 변환하는 방법을 제안합니다. 이는 다단계 증류 과정을 포함합니다:

1. 선형 사전 정렬: 선형 믹서 출력을 교사 모델의 잔여 스트림에 정렬합니다.
2. 하이브리드 로그트 증류: KL 가이드 레이어 선택을 사용하여 완전 어텐션 레이어의 소량 부분 집합 (예: 25%) 을 유지하면서 교사 모델의 로그트를 증류합니다.
3. 긴 컨텍스트 연속 학습: 외삽 능력을 보장하기 위해 긴 컨텍스트 데이터로 추가 학습을 수행합니다.

3. 주요 기여

• LT2 아키텍처: 2 차 어텐션을 선형 및 희소 믹서로 대체하여 선형 시간 복잡도 ($O(L)$) 를 달성하는 새로운 루프형 트랜스포머 계열.
• 이론적 분석: DPLR 선형 어텐션을 루핑하면 상태 업데이트의 랭크가 $T$로 증가하고, 희소 어텐션을 루핑하면 수용 영역이 $O(Tw)$로 확장됨을 보여주는 공식 증명.
• 하이브리드 설계: 품질과 효율성 사이의 파레토 프론티어를 탐색하는 두 가지 유망한 하이브리드 구성 (GDN+DSA 및 Full+GDN) 식별.
• 증류 파이프라인: 기존 사전 학습된 루프형 모델을 선형 시간 변형으로 변환하여 성능은 유지하면서 추론 비용을 획기적으로 줄이는 성공적인 레시피.
• 안정성 메커니즘: 가중치 공유 루프에서의 학습 안정성에 중요한 "어텐션 싱크" 누적 및 루프 반복 전반에 걸친 거대한 활성화 증가를 완화하기 위한 SDPA 출력 게이트 도입.

4. 실험 결과

저자들은 언어 모델링, 합성 작업, 그리고 긴 컨텍스트 검색 벤치마크에서 LT2 를 평가했습니다.

• 언어 모델링:
  • 13 억 파라미터에서 LT2-hybrid (Full+GDN) 는 평균 제로샷 성능 **62.89%**를 달성하여 선형 시간 효율성을 유지하면서 표준 루프형 트랜스포머 (59.27%) 를 능가했습니다.
  • LT2-hybrid (GDN+DSA) 는 표준 루프형 트랜스포머의 품질 (평균 59.3%) 과 일치했지만, 8k 컨텍스트에서 약 5.7 배 더 높은 디코드 처리량 (초당 125 토큰 vs 22 토큰) 을 제공했습니다.
• 효율성:
  • LT2 변형은 표준 루프형 트랜스포머에서 관찰된 "디코드 절벽"을 제거했습니다. 표준 LT 는 4k 에서 32k 컨텍스트 사이에서 디코드 처리량의 50% 이상을 잃는 반면, LT2 변형은 모든 시퀀스 길이와 배치 크기에서 평평한 처리량을 유지했습니다.
  • LT2 모델은 배치 크기 8 에서 32k 시퀀스를 처리할 수 있었지만, 표준 루프형 트랜스포머는 8k 에서 메모리 부족 (OOM) 이 발생했습니다.
• 합성 작업 (상태 추적 및 회상):
  • 루핑은 2 차 미만의 믹서 성능을 크게 향상시켰습니다. 표준 루프형 트랜스포머는 커리큘럼 단계 4 ($n=64$) 에서 정체되었지만, 2 차 미만의 하이브리드 (예: 루프형 GDN+NSA) 는 단계 5 ($n=128$) 에 도달하여 해결 가능한 최대 문제 크기를 두 배로 늘렸습니다.
• 증류 결과:
  • 약 10 억 토큰의 연속 학습만으로 훈련된 변환된 Ouro-Hybrid-1.4B 모델은 산업 수준의 10 억 파라미터 모델을 능가했으며, 제로샷 벤치마크에서 40 억 파라미터 모델과 경쟁력을 보였으며, 선형 시간 어텐션의 속도 이점을 유지했습니다.

5. 중요성과 주장

이 논문은 LT2 가 루프형 트랜스포머를 언어 모델링을 위한 확장 가능한 아키텍처로 만드는 명확한 경로를 제공한다고 주장합니다. "유효 깊이" (루핑이 제공) 를 "2 차 비용" (어텐션이 제공) 에서 분리함으로써 LT2 는 다음을 가능하게 합니다:

1. 확장 가능한 재귀: prohibitive 한 계산 비용 없이 긴 컨텍스트를 위한 깊은 반복적 추론을 실현 가능하게 함.
2. 효율적인 소형 모델: 효율적인 토큰 믹서와 루프형 아키텍처를 활용한다면 소형 모델 (10 억~40 억 파라미터 범위) 이 훨씬 큰 산업용 모델과 경쟁 가능한 성능을 달성할 수 있음을 입증.
3. 실용적 채택: 기존 사전 학습된 모델을 LT2 변형으로 증류할 수 있는 능력은 진입 장벽을 낮추어, 실무자들이 처음부터 재학습하지 않고도 이러한 효율적인 아키텍처를 채택할 수 있게 함.

저자들은 효율적인 토큰 믹서가 재귀적 깊이를 미래 언어 모델을 위한 실용적인 확장 축으로 만들 수 있으며, 효율적이고 능력이 뛰어난 소형 언어 모델의 개발을 진전시킬 것이라고 결론지었습니다.