기술 요약: 저복잡성 정보 풍부 시퀀스 모델로서의 토플리츠 MLP 믹서

1. 문제 제기

트랜스포머 기반 대규모 언어 모델 (LLM) 은 자기주의 (self-attention) 메커니즘의 이차 시간 ($O(n^2)$) 및 공간 복잡도에 근본적으로 제약받습니다. 선형화 attention, Mamba 와 같은 상태 공간 모델 등 이 복잡성을 이차 미만 또는 선형 수준으로 낮추기 위한 다양한 시도가 있었으나, 이러한 접근법들은 종종 상당한 트레이드오프를 초래합니다. 구체적으로 선형 복잡도 모델은 훈련 효율성 저하, 수치적 불안정성, 그리고 입력 정보 유지 능력의 현저한 결여로 고통받습니다. 이러한 정보 유지 부족은 정밀한 복사, 인-컨텍스트 학습, 장문맥 검색이 필요한 작업에서 낮은 성능으로 나타나며, 이는 현대 LLM 능력의 특징입니다. 핵심 과제는 이러한 신흥 능력을 저하시키는 토큰 정보 접근에 대한 심각한 제한을 부과하지 않으면서 계산 복잡도를 낮추는 것입니다.

2. 방법론: 토플리츠 MLP 믹서 (TMM)

저자들은 자기주의를 시퀀스 차원에서의 삼각 마스크가 적용된 토플리츠 행렬 곱셈으로 대체하는 트랜스포머와 유사한 아키텍처인 **토플리츠 MLP 믹서 (TMM)**를 소개합니다.

핵심 아키텍처

• 토플리츠 파라미터화: 토큰 혼합 가중치가 임의의 행렬인 표준 MLP 믹서와 달리, TMM 은 이러한 가중치를 토플리츠 행렬로 파라미터화합니다. 토플리츠 행렬에서 주어진 대각선 위의 모든 값은 동일하므로, 전체 행렬을 단일 벡터 (첫 번째 행과 열) 로 표현할 수 있습니다.
• 인과적 마스크: 다음 토큰 예측을 위한 인과성을 강제하기 위해, 구현된 가중치 행렬은 삼각 마스크가 적용됩니다. 구현에서 활성화 값은 이 마스크가 적용된 토플리츠 행렬 ($T'$) 과 왼쪽에서 곱해지는데, 이는 상삼각 행렬입니다.
• FFT 를 통한 계산: 토플리츠 행렬을 순환 행렬 (circulant matrix) 에 임베딩함으로써, 저자들은 행렬 - 벡터 곱을 계산하기 위해 고속 푸리에 변환 (FFT) 을 활용합니다. 이는 훈련 및 추론 프리필 (prefill) 동안 계산 복잡도를 $O(n^2)$에서 $O(dn \log n)$으로 줄이고, **$O(dn)$**의 공간 복잡도를 가집니다.
• 멀티헤드 및 멀티커널 변형: 이 아키텍처는 투영된 헤드 차원에 독립적인 토플리츠 행렬을 적용하는 멀티헤드 구성을 지원합니다. 또한, "커널" 변형은 1D 합성곱과 유사하게 제한된 수의 숨겨진 차원 요소들을 따라 혼합을 허용합니다.

이론적 기반

본 논문은 **연산자 지수 이론 (operator index theory)**의 관점에서 TMM 을 분석합니다. 구체적으로 훈련된 토플리츠 층의 프레드홀름 지수 (Fredholm index) 를 검토합니다. 저자들은 TMM 이 모든 가능한 입력에 대해 엄격하게 가역적이지는 않지만, 훈련된 인과적 TMM 층은 종종 -1 과 같은 지수를 나타내어 시프트 연산자와 유사하며, 이는 다른 아키텍처에서 발견되는 편향과 구별되는 순차적 정보 이동 메커니즘을 시사한다고 지적합니다.

3. 주요 기여 및 주장

이 논문은 Hyena, Mamba, 윈도우 attention 트랜스포머와 같은 다른 이차 미만 아키텍처와 비교하여 TMM 의 효과성에 대해 세 가지 주요 주장을 제기합니다:

1. 훈련 효율성: 일정한 계산 조건 하에서 TMM 은 다른 이차 미만 모델과 유사하거나 더 낮은 평가 손실을 달성합니다. Hyena 및 Mamba 와 같은 선형 복잡도 기준 모델에 비해 우수한 수렴 속도를 보여줍니다.
2. 정보 보존: TMM 은 비교 가능한 이차 미만 모델보다 훨씬 더 많은 입력 정보를 유지합니다. 이는 다음과 같은 증거로 입증됩니다:
   • 정보 추출: 최종 숨겨진 상태에서 입력 시퀀스를 재생성하는 데 더 높은 충실도.
   • 복사 능력: TMM 은 Mamba 및 Hyena 모델이 일반적으로 어려움을 겪는 복사 작업의 훈련 효율성에서 이차 복잡도 트랜스포머에 근접하거나 이를 초과합니다.
   • 정보 용량: 명시적으로 정보 유지를 최대화하도록 훈련될 때, TMM 은 Hyena 나 Mamba 모델보다 더 높은 정보 용량을 보여주며, 이는 훈련에 필요한 계산량이 1/2 에서 1/8 로 훨씬 적음에도 불구하고 해당됩니다.
3. 벤치마크 성능: 우수한 정보 유지 능력으로 인해 TMM 은 정보 검색 (IR) 및 인-컨텍스트 학습 (ICL) 과 같이 정밀한 입력 표현이 필요한 벤치마크에서 다른 저복잡도 모델보다 우수한 성능을 발휘하며, 일반 언어 이해 작업에서는 동등한 성능을 유지합니다.

4. 실험 결과

저자들은 FineWeb-edu 의 100 억 토큰 부분 집합에서 모델을 훈련하고 트랜스포머, Mamba 2, Hyena 모델과 비교 평가했습니다.

• 훈련 효율성: 계산 및 메모리가 일치할 때 TMM 은 Hyena 및 Mamba 모델보다 빠르게 수렴했습니다. 저자들은 선형 attention 구현이 수치적 불안정성과 느린 수렴으로 고통받았음을 지적하며 TMM 에 초점을 맞췄습니다.
• 복사 작업: 512 토큰을 입력에서 출력으로 복사해야 하는 작업에서 TMM 은 트랜스포머와 비교 가능한 훈련 효율성을 달성했습니다. 주목할 점은 Hyena 모델의 모델 차원 ($d_m$) 의 1/32 에 불과한 TMM 이 거의 동일한 복사 효율성을 달성했다는 것입니다. 연구 결과 훈련 가능한 토플리츠 파라미터가 결정적이었으며, 이러한 행렬을 고정하면 복사 실패가 완전히 발생했으나 훈련 가능한 헤드 투영만으로는 부족함이 확인되었습니다.
• 벤치마크:
  • IR/ICL: TMM 은 SQuAD, SQuAD v2, LongBench, IFEval 에서 Hyena 를 능가했습니다.
  • 일반 작업: ARC-Easy, GLUE, HellaSwag 와 같은 작업에서 TMM 은 Hyena 및 트랜스포머와 동등한 성능을 발휘했습니다.
• 절대 연구 (Ablation Studies): 저자들은 비선형성으로 분리된 두 개의 순차적 토플리츠 층을 사용하는 것 (표준 MLP 믹서 스타일) 이 단일 층보다 효율성이 낮음을 발견했습니다. 또한, 헤드의 수를 늘리는 것은 점근적 수렴을 개선하는 반면, 커널 크기를 일정 지점 이상 늘리는 것은 효율성을 감소시켰습니다.

5. 의의 및 한계

의의

본 논문은 TMM 이 경쟁력 있는 훈련 효율성과 빠른 디코딩을 위한 "간단하고 attention 이 없는 경로"를 제공한다고 주장합니다. 그 주요 의의는 정보 유지를 저하시키는 심한 아키텍처 편향 (예: 최근성 편향 또는 상태 공간 선택성) 이 저복잡도 모델에 본질적으로 필요하지 않음을 입증했다는 점에 있습니다. 명시적인 토플리츠 파라미터화를 사용하여 모델은 Mamba 와 Hyena 에서 관찰되는 "복사 격차"를 피하며, 이는 구조화된 접근 패턴 (토플리츠) 이 프리필 단계의 이차 비용 없이 트랜스포머의 정보 풍부한 능력을 보존할 수 있음을 시사합니다.

저자들은 또한 역설적인 이론적 발견을 강조합니다: 전체 입력 공간에서 비가역적인 훈련된 인과적 TMM 층은 종종 가역성을 위해 명시적으로 훈련된 모델 (예: 오토인코더) 보다 제로 (또는 구체적으로 -1 로 작용하는 시프트 연산자) 에 가까운 프레드홀름 지수를 가집니다. 이는 인과적 모델에서 층의 "가역성"이 수학적 가역성보다는 시프트를 통한 정보 흐름의 보존과 더 관련이 있음을 시사합니다.

한계

저자들은 명시적으로 몇 가지 한계를 인정합니다:

• 규모: 실험은 약 130 억260 억 토큰으로 훈련된 2 천만3 억 파라미터 범위의 모델에서 수행되었습니다. 이러한 발견이 계산량이 몇 배 더 많은 모델로 확장될 것이라는 보장은 없습니다.
• 자기회귀 추론: TMM 은 프리필 동안 $O(dn \log n)$이지만, 첫 번째 토큰이 계산된 후 자기회귀 생성 시에는 $O(n^2)$ 복잡성으로 돌아갑니다 (입력 행렬의 행을 캐시하고 순차적으로 처리해야 함). 따라서 TMM 은 긴 프롬프트에 대한 첫 번째 토큰 도달 시간을 단축하지만, 트랜스포머에 비해 후속 토큰 생성 속도를 높이지는 못합니다.
• 구성 공간: 가능한 하이퍼파라미터 조합의 지수적 수로 인해 연구는 제한된 수의 아키텍처 구성을 테스트했습니다.

결론적으로, TMM 은 자기회귀 생성의 추론 병목 현상이 관리된다면 복잡한 추론 및 검색 작업에 필요한 높은 정보 유지와 계산 효율성 사이의 균형을 맞추는 시퀀스 모델링을 위한 유망한 방향을 나타냅니다.