Sequential-Parallel Duality in Prefix Scannable Models

이 논문은 병렬 학습과 순차 추론을 동시에 지원하는 '접두어 스캔 가능 모델 (Prefix-Scannable Models)'이라는 새로운 범주를 정의하여 기존 아키텍처를 통합하고, 새로운 연산자를 통해 효율성과 표현력을 동시에 달성하는 모델을 제안합니다.

핵심

🏗️ 1. 문제 상황: 두 가지 세계의 갈등

현재 AI 모델 (특히 '트랜스포머') 은 두 가지 큰 문제를 안고 있습니다.

1. 학습할 때는 빠르지만, 추론할 때는 느리다:

   • 비유: 공장에서 제품을 대량 생산할 때는 모든 공장을 동시에 가동해서 (병렬) 아주 빠르게 만들 수 있습니다. 하지만, 고객이 하나씩 주문할 때 (추론) 마다, 그전까지 만든 모든 제품 목록을 다시 다 뒤져야 하므로 시간이 너무 오래 걸립니다.
   • 결과: 문장이 길어질수록 메모리도, 시간도 기하급수적으로 늘어납니다.

2. 기존 RNN(순환 신경망) 은 반대:

   • 비유: 공장이 아니라, 한 사람이 일일이 순서대로 제품을 만들어 나가는 방식입니다. 주문할 때는 아주 빠르고 메모리도 적게 쓰지만, 대량 생산 (학습) 을 할 때는 한 번에 하나씩만 만들 수 있어 매우 비효율적입니다.

핵심 질문: "학습할 때는 모든 공장을 동시에 가동하고, 주문할 때는 한 사람처럼 빠르고 가볍게 작동하는 완벽한 하이브리드 공장을 만들 수 있을까?"

---

🔍 2. 해답: '접두어 스캔 (Prefix Scan)'이라는 마법 도구

저자들은 이 문제를 해결하기 위해 컴퓨터 과학의 고전적인 알고리즘인 **'접두어 스캔 (Prefix Scan)'**을 활용합니다.

• 비유 (계단 오르기):
  • 기존 방식: 100 개의 계단을 오를 때, 1 단계부터 2 단계, 3 단계... 순서대로 하나씩 올라가야 합니다. (순차적)
  • 접두어 스캔: 100 명을 2 명씩 짝을 지어, "1+2", "3+4"를 동시에 계산하게 합니다. 그다음 그 결과들을 다시 짝지어 "14", "58"을 계산합니다. 이렇게 이진 트리 (Binary Tree) 구조로 계단을 올라가면, 100 계단을 아주 짧은 시간 (로그 시간) 에 정복할 수 있습니다.

이 논문의 핵심은 **"이 '접두어 스캔' 방식을 AI 의 상태 업데이트에 적용하면, 학습은 병렬로, 추론은 순차적으로 아주 효율적으로 할 수 있다"**는 것입니다.

---

🚀 3. 새로운 모델: PSM (접두어 스캔 가능 모델)

저자들은 이 원리를 확장하여 PSM이라는 새로운 모델 클래스를 정의했습니다.

A. 기존 모델들 (선형 RNN, Mamba 등)

• 비유: 이 모델들은 '접두어 스캔'을 할 때, **덧셈 (+)**처럼 규칙이 단순하고 예측 가능한 (결합 법칙이 성립하는) 연산을 사용합니다.
• 장점: 매우 빠르고 효율적입니다.
• 한계: 너무 단순한 규칙이라 복잡한 문맥 이해 (예: "누가 무엇을 했는지" 같은 복잡한 관계) 를 배우는 데는 약점이 있을 수 있습니다.

B. 새로운 모델: Transformer-PSM

• 비유: 저자들은 "단순한 덧셈만 쓸 필요는 없다!"라고 말합니다. 대신 **소프트맥스 (Softmax)**처럼 복잡한, 규칙이 조금 더 유연한 (결합 법칙이 성립하지 않는) 연산을 '접두어 스캔'에 적용할 수 있습니다.
• 어떻게? 문장을 작은 '조각 (Chunk)'으로 나누고, 그 조각들 사이를 오가는 복잡한 관계 (트랜스포머의 주의 메커니즘) 를 계산하되, 그 결과를 이진 카운터 (Binary Counter) 방식의 메모리 구조에 저장합니다.
• 효과:
  • 학습: 모든 조각을 동시에 계산 (병렬).
  • 추론: 문장이 길어져도 메모리는 로그 (Log) 수준으로만 늘어나고, 한 번에 처리하는 속도는 일정하게 유지됩니다.
  • 기적: 기존 트랜스포머처럼 복잡한 문맥을 이해하면서도, Mamba 처럼 가볍고 빠릅니다.

---

🧪 4. 실험 결과: 실제로 작동할까?

저자들은 이 모델을 여러 가지 테스트에 적용했습니다.

1. 상태 추적 (State Tracking):

   • 상황: "컵을 A 에서 B 로, B 에서 C 로 옮겼다"는 식의 복잡한 순서를 기억하는 게임입니다.
   • 결과: 기존 모델 (트랜스포머, Mamba) 은 문장이 길어지면 실수가 잦아졌지만, Transformer-PSM은 훈련된 길이보다 훨씬 긴 문장에서도 거의 완벽하게 기억해냈습니다. (길이에 대한 일반화 능력이 뛰어남)

2. 연관 기억 (Associative Recall):

   • 상황: "사과"라는 단어가 나왔을 때, 문장 앞부분에 있던 "빨간색"을 기억해내는 작업입니다.
   • 결과: 문장 조각 (Chunk) 의 크기를 조절하면, 트랜스포머의 정밀한 기억 능력과 RNN 의 빠른 속도를 모두 잡을 수 있었습니다.

3. 언어 모델링 (WikiText):

   • 결과: 기존 GPT-2 와 비슷한 수준의 언어 이해 능력을 보이면서, 문장이 길어질수록 지연 시간 (Latency) 이 거의 늘어나지 않았습니다.

---

💡 5. 결론: 왜 이것이 중요한가?

이 논문은 **"학습은 무작정 빠르게, 추론은 가볍게"**라는 AI 의 오랜 딜레마를 해결할 수 있는 새로운 설계 도면을 제시합니다.

• 핵심 메시지: 우리는 더 이상 "빠르지만 멍청한 모델"과 "똑똑하지만 느린 모델" 사이에서 선택할 필요가 없습니다. 접두어 스캔이라는 수학적 도구를 통해, **두 마리 토끼를 모두 잡을 수 있는 모델 (Transformer-PSM)**을 만들 수 있습니다.

한 줄 요약:

"AI 가 문장을 읽을 때, 처음부터 끝까지 한 줄로 읽는 게 아니라, 작은 블록으로 나누어 동시에 계산하되, 기억할 때는 아주 작은 메모리만 쓰는 새로운 방식을 제안합니다. 이는 더 길고 복잡한 문장을 처리하면서도 속도가 느려지지 않는 차세대 AI 의 핵심 열쇠가 될 것입니다."

기술 요약

1. 문제 정의 (Problem)

현대 신경망 시퀀스 모델은 병렬 학습 (Parallelizable Training) 과 고속 순차 추론 (Fast Sequential Inference) 이라는 두 가지 상충되는 요구사항을 동시에 충족해야 합니다.

• Transformer: 훈련 시 병렬 처리가 가능하고 장기 의존성을 잘 포착하지만, 추론 시 메모리 복잡도가 시퀀스 길이의 제곱 ($O(N^2)$) 에 비례하여 확장됩니다. 또한, KV 캐시를 저장해야 하므로 메모리 사용량이 선형 ($O(N)$) 으로 증가합니다.
• 전통적 RNN 및 상태 공간 모델 (SSM, 예: Mamba): 추론 시 $O(1)$ 의 메모리와 선형 시간 복잡도를 가지지만, 훈련 시 순차적 의존성으로 인해 병렬화가 어렵거나 표현력 (Expressivity) 이 제한적입니다.

이러한 한계를 극복하기 위해 최근 Gated Linear Attention (GLA), Mamba 등 "순차 - 병렬 이중성 (Sequential-Parallel Duality, SPD)"을 가진 모델들이 등장했습니다. 그러나 어떤 모델 클래스가 거의 상수 시간의 병렬 평가와 선형 시간/상수 공간의 순차 추론을 모두 지원할 수 있는지에 대한 이론적 체계화가 부족했습니다.

2. 방법론 (Methodology)

저자들은 기존 모델들의 공통점을 병렬 접두사 스캔 (Parallel Prefix Scan) 알고리즘 관점에서 재정의하고, 이를 일반화한 새로운 모델 클래스를 제안합니다.

A. 접두사 스캔 가능 모델 (Prefix-Scannable Models, PSMs)

• 핵심 아이디어: 시퀀스 모델의 상태 업데이트를 Blelloch 접두사 스캔 알고리즘으로 계산 가능한 형태로 정의합니다.
• 기존 모델 (SPD-(n, 1)): 기존 Mamba, Linear Transformer 등은 상태 업데이트 연산자가 결합법칙 (Associativity) 을 만족합니다. 이 경우 스캔 알고리즘을 통해 $O(N)$ 작업량과 $O(\log N)$ 깊이의 병렬 학습, 그리고 $O(1)$ 메모리의 순차 추론이 가능합니다.
• 새로운 일반화 (SPD-(n, log n)): 저자들은 결합법칙을 만족하지 않는 연산자 (예: Softmax Attention) 도 허용하는 PSM을 정의합니다.
  • Chunking: 시퀀스를 $c$ 크기의 청크 (chunk) 로 나눕니다.
  • Aggregation: 각 청크를 인코딩한 후, 결합법칙이 성립하지 않는 연산자 (Agg) 를 사용하여 청크 간 상태를 병합합니다.
  • Online Inference: 훈련 시 정적 (Static) 스캔을 사용하고, 추론 시 이진 카운터 (Binary Counter) 기반의 온라인 스캔을 사용하여 $O(\log N)$ 메모리 내에서 동일한 계산 순서 (Parenthesization) 를 재현합니다.

B. Transformer-PSM

• 저자들은 PSM 프레임워크를 구체화한 Transformer-PSM을 제안합니다.
• 구조:
  1. Encoder: 토큰을 청크 단위로 인코딩.
  2. Aggregation: GPT-2 스타일의 양방향 어텐션 (Bidirectional Attention) 을 사용하여 청크 간 상태를 병합 (결합법칙이 성립하지 않음).
  3. Inference: 인코딩된 이전 상태와 현재 청크를 결합하여 다음 토큰을 예측 ( causal mask 사용).
• 이 모델은 Transformer 의 표현력을 유지하면서도 PSM 의 효율적인 추론 구조를 따릅니다.

3. 주요 기여 (Key Contributions)

1. SPD 모델의 이론적 체계화: 현대의 선형 RNN 과 상태 공간 모델들이 결합법칙을 만족하는 연산자를 가진 접두사 스캔 알고리즘의 특수한 경우임을 증명하고 통일된 관점을 제시했습니다.
2. PSM (Prefix-Scannable Models) 의 정의: 결합법칙이 성립하지 않는 일반적 연산자 (Softmax 등) 를 허용하는 모델 클래스를 정의했습니다. 이는 훈련 시 $O(N)$, 추론 시 $O(\log N)$ 메모리, $O(1)$ 평균 추론 시간을 보장합니다.
3. Transformer-PSM 구현 및 검증: PSM 프레임워크를 기반으로 Transformer-PSM 을 설계하고, 상태 추적 (State Tracking), 연관 회상 (Associative Recall), 언어 모델링 등 다양한 태스크에서 기존 모델 대비 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

저자들은 WikiText-103 언어 모델링과 합성 알고리즘 태스크 (S5 상태 추적, MQAR) 를 통해 Transformer-PSM 을 평가했습니다.

• 상태 추적 (State Tracking - S5):
  • 훈련 시 18 토큰까지 학습했으나, 160 토큰 이상의 시퀀스에서도 Transformer 와 Mamba 보다 훨씬 낮은 오류율을 보이며 뛰어난 길이 일반화 (Length Generalization) 능력을 입증했습니다.
• 연관 회상 (Associative Recall - MQAR):
  • 청크 크기 (Chunk Size) 를 64 로 설정한 Transformer-PSM 은 전체 컨텍스트를 사용하는 Transformer 와 유사한 완벽한 정확도를 달성하면서도, Mamba 나 슬라이딩 윈도우 Transformer 보다 우수한 성능을 보였습니다.
• 언어 모델링 (WikiText-103):
  • 청크 크기가 커질수록 (32 → 256) 퍼플렉시티 (Perplexity) 가 감소하여 Vanilla GPT-2 에 근접하는 성능을 보였습니다.
  • 추론 효율성: GPT-2 는 컨텍스트 길이가 길어질수록 추론 지연 시간이 선형적으로 증가하는 반면, Transformer-PSM 은 청크 요약 정보를 재사용하여 상수 시간 ($O(1)$) 에 가까운 지연 시간을 유지했습니다.

5. 의의 및 결론 (Significance)

• 이론적 통합: 기존에 별개로 연구되던 선형 RNN, 상태 공간 모델, 선형 어텐션 모델들을 접두사 스캔 알고리즘이라는 단일 프레임워크 하에 통합했습니다.
• 설계 공간 확장: 결합법칙을 강제하지 않음으로써 Softmax 어텐션과 같은 강력한 연산자를 효율적인 순차 추론 구조에 통합할 수 있음을 보였습니다.
• 실용적 가치: Transformer 의 높은 표현력과 RNN/SSM 의 효율적인 추론을 모두 갖춘 새로운 모델 아키텍처를 제시하여, 긴 시퀀스 처리가 필요한 실제 응용 분야에서 Transformer 의 병목 현상을 해결할 수 있는 가능성을 열었습니다.

요약하자면, 이 논문은 PSM이라는 새로운 모델 클래스를 정의하고, 이를 통해 병렬 학습의 효율성과 순차 추론의 속도를 동시에 달성하면서도 Transformer 의 표현력을 유지하는 모델을 성공적으로 구현하고 검증했습니다.