기술 요약: O(1) 셀프 어텐션을 위한 단 하나의 벡터: 오션 스테이트 (The Ocean State)

1. 문제 정의

표준 트랜스포머는 구조적 한계를 겪고 있다: 셀프 어텐션은 $O(N^2)$의 계산 비용을 발생시키며, 컨텍스트 길이에 따라 선형적으로 증가하는 $O(N)$ 메모리 캐시(Key-Value 캐시)를 필요로 한다. 최근의 연구들(RetNet, RWKV, Mamba, Titans)은 선형 어텐션, 선택적 게이팅(selective gating), 또는 외부 메모리 모듈과 같은 새로운 메커니즘을 도입하여 이를 완화하려 시도하지만, 이들은 근본적으로 아키텍처를 변경한다. 또한, 표준 트랜스포머는 매 포워드 패스(forward pass) 이후에 전체 컨텍스트를 흡수한 결과물인 최종 은닉 상태(hidden state)를 폐기한다.

본 논문은 $O(N^2)$의 비용이 컨텍스트 모델링이라는 개념 자체에 내재된 것이 아니라, 현재의 구현 방식에 구조적으로 귀속된 것이라고 상정한다. 질문은 다음과 같다: 만약 전체 컨텍스트를 흡수하는 은닉 상태를 폐기하지 않고 유지하며, 다음 토큰을 위한 유일한 어텐션 대상으로 사용한다면 어떻게 될 것인가?

2. 방법론: 오션 스테이트 (The Ocean State)

제안된 해결책인 "오션 스테이트"는 새로운 구성 요소를 추가하지 않고(예: 게이팅, 외부 메모리 없음), 표준 트랜스포머 아키텍처를 두 가지 방식으로 재배치하는 것을 포함한다.

2.1 핵심 메커니즘

1. 오션 스테이트 (재배치 1):

   • 포워드 패스 이후 은닉 상태를 폐기하는 대신, 현재 토큰의 처리 결과가 이전 은닉 상태를 대체하며, 이를 "오션 스테이트(ocean state)"라고 부른다.
   • 어텐션 메커니즘: 어텐션은 엄격하게 현재 토큰과 지속적인 오션 스테이트(단일 밀집 벡터) 사이에서 계산된다.
   • 결과: 이는 Key-Value(KV) 캐시, 위치 인코딩(positional encodings), 그리고 인과적 마스크(causal masks)의 필요성을 제거한다. 상태는 누적이 아닌 "릴레이"(1:1 교체)를 통해 업데이트되므로, 정보가 압축되거나 버려지는 것이 아니라 재배열된다.
   • 읽기 전용 (Read-Only): 오션 스테이트는 현재 토큰에 의해 읽히지만, 피드 포워드 네트워크(FFN)나 레이어 스택을 통과하지 않는다. 오직 토큰만이 스택을 통과한다.

2. 오션 BPTT (재배치 2):

   • 오션 스테이트는 순환 변수(recurrent variable)이므로, 시간 축을 따라 언롤링(unrolling)되는 역전파(Backpropagation Through Time, BPTT)가 필요하다.
   • 그래디언트 흐름: 백워드 패스(backward pass) 동안, 단계 $t$에서의 오션 스테이트의 그래디언트는 단계 $t$의 토큰을 기반으로 계산된다. 이 그래디언트는 단계 $t-1$의 은닉 상태 그래디언트에 더해진다.
   • 리드-패스 연산자 (Read-Path Operator): 특정 연산자가 포워드 패스에서 저장된 어텐션 가중치를 사용하여 토큰의 그래디언트를 오션 스테이트로 매핑한다. 토큰이 오션을 더 많이 "읽을수록", 오션 상태를 업데이트하기 위해 전달되는 그래디언트 신호는 더 강해진다.

2.2 알고리즘 구현

학습 루프는 다음과 같이 구성된다:

• 포워드 패스 (Forward Pass): 토큰 $t$는 $ocean_{t-1}$에 어텐션을 수행한다. 토큰은 레이어 스택(Attention + FFN)을 통과한다. 출력값은 $ocean_{t-1}$을 대체하여 $ocean_t$가 된다.
• 백워드 패스 (Backward Pass): 그래디언트는 단계 $t$에서의 손실(loss)로부터 단계 $t$의 레이어를 거쳐 역방향으로 흐른다. 결정적으로, 오션 스테이트의 그래디언트는 시간 단계에 걸쳐 누적되어, 미래 토큰으로부터의 에러가 과거 컨텍스트의 표현을 업데이트할 수 있게 한다.

2.3 엔지니어링 변형

본 논문은 비용, 시간, 품질의 균형을 맞추기 위해 세 가지 학습 레짐(regime)을 탐구한다:

• 직렬 오션 BPTT (Serial Ocean BPTT): 전체 시간 체인을 언롤링한다. 가장 높은 품질의 이득을 얻지만, 계산 비용이 높고 실행 시간(wall-clock time)이 길다.
• 프레임워크 상의 오션 (Ocean on Framework): 표준 오토그래드(autograd) 프레임워크에서 전체 체인을 실행한다. 완전한 품질 이득을 유지하면서 배칭(batching)을 통해 처리량을 개선한다.
• 오션 디태치 (Ocean Detach): 시간 그래디언트 체인을 끊는다(단계 간 BPTT 없음). 표준 트랜스포머와 같이 완전한 병렬 학습이 가능하지만, 시간 그래디언트로부터 얻는 성능 이득은 희생된다.

3. 주요 기여

• 아키텍처 재배치: 기존의 KV 캐시와 위치 인코딩을 제거하고, 버려지는 은닉 상태를 지속적인 "오션" 벡터로 재사용함으로써 $O(1)$ 어텐션과 $O(1)$ 메모리가 가능함을 입증했다.
• 창발적 능력 (Emergent Capability): 트랜스포머가 명시적인 지시 없이도 단일 밀집 벡터로부터 정보를 읽고 추출하는 법을 학습할 수 있음을 보여준다.
• 극한의 컨텍스트에 대한 확장성: 이 아키텍처가 성능 저하 없이 최대 1,000만 단계의 컨텍스트 길이를 처리할 수 있음을 증명했다. 이는 표준 트랜스포머가 메모리와 계산 제약으로 인해 불가능한 수치이다.
• 그래디언트 언롤링: "오션 BPTT"를 도입하여 전체 레이어 스택의 출력을 순환 상태(recurrent state)로 적용하였으며, 이것이 성능 향상의 주요 원천임을 확인했다.

4. 실험 결과

실험은 커스텀 NumPy 구현(float64)으로 수행되었으며 PyTorch를 통해 검증되었고, 셰익스피어 코퍼스(문자 및 단어 수준)를 대상으로 진행되었다.

4.1 성능 및 학습 역학

• 우월성: 테스트된 모든 구성(너비 $d \in \{32, 64, 128, 256\}$, 깊이 $L \in \{2, 4, 6\}$, 다양한 헤드 수)에 대해, "Ocean Full" 레짐이 표준 트랜스포머보다 일관되게 우수한 성능을 보였다.
• 안정적인 이득: 성능 차이($\Delta = Loss_{std} - Loss_{full}$)는 양수이며 안정적이었다(약 +0.15의 손실 감소).
• 이득의 원천: 이 이득은 주로 구조적이다. 시간 체인이 분리된 경우(Ocean Detach), 성능이 크게 떨어지는데, 이는 시간 그래디언트 언롤링이 핵심 요소이지 단순히 어텐션 메커니즘 때문이 아님을 나타낸다.

4.2 장기 컨텍스트 및 메모리

• 망각 없음: 1,000만 단계의 프리필(prefill) 테스트에서 손실값은 평탄하게 유지되었다(1M 단계에서 1.532, 10M 단계에서 1.531). 오션 스테이트는 용량 포화나 망각의 징후를 보이지 않았다.
• 에코 태스크 (Echo Task): 200만 단계의 지연(delay)이 있는 에코 태스크에서 모델은 거의 0에 가까운 손실(0.0006)을 달 기록한 반면, 표준 트랜스포머는 완전히 실패했다.
• 메모리 효율성: 오션 스테이트는 컨텍스트 길이에 독립적인 $O(1)$ 메모리(차원이 $d$인 단일 벡터)만을 요구한다. 이와 대조적으로, 표준 트랜스포머는 $O(N)$ 메모리를 위한 KV 캐시가 필요하며, 이는 100만 토큰 규모(13B 모델 기준 약 1TB 필요)에서 실행 불가능해진다.

4.3 추론 및 생성

• 외삽 불가 (No Extrapolation): 오션은 위치 인코딩 외삽(예: RoPE, ALiBi)을 필요로 하지 않는 일관된 텍스트를 생성하는데, 이는 컨텍스트 길이가 어텐션 메커니즘의 입력 변수가 아니기 때문이다.
• 실행 시간 (Wall-Clock Time): 순환의 직렬적 특성 때문에 학습 시 단계당 표준 트랜스포머보다 느리지만, 레이어별 파이프라이닝(Ocean A-2)을 통해 실행 시간을 1.9배 단축할 수 있으며, 중첩 병렬성(nested parallelism, Ocean A-3)은 이론적으로 $10^3$의 이득을 제공한다. 저자는 학습 과정에서 비용, 시간, 품질 사이의 명시적인 트레이드오프가 존재하며, 실제 구현을 위해서는 특정 레짐(예: 산업적 규모를 위한 detach, 연구 품질을 위한 full chain)을 선택해야 함을 강조한다.

5. 의의 및 주장

논문은 다음 세 가지 핵심 사항을 제시하며, 이를 검증 가능하다(falsifiable)고 정의한다:

1. 학습 일관성: 오션 스테이트는 표준 트랜스포머와 동일한 학습 곡선 형태를 가지며 학습되므로, 이것이 근본적으로 다른 학습 패러다임이 아니라 유효한 재배치임을 확인한다.
2. 보편적 이득: 동일한 조건 하의 모든 트랜스포머 아키텍처에 대해 오션 스테이트는 양(+)의 성능 이득($\Delta > 0$)을 제공한다.
3. 무한한 용량: 충분한 차원($d$)이 있다면, 오션 스테이트는 관찰 가능한 어떤 규모에서도 용량 제한, 망각, 또는 검색 실패를 보이지 않는다.

저자는 "오션 스테이트"가 임의의 컨텍스트 길이에 대해 $O(1)$ 추론 및 학습의 원칙을 해결한다고 결론지으며, "우리는 이 원칙이 해결되었다고 믿는다"라고 밝혔다. 단일 벡터 내에 정보를 조직화하는 트랜스포머의 창발적 능력에 의존함으로써, 이 방법은 이차적(quadratic) 어텐션과 선형적 메모리 증가라는 구조적 병목 현상을 제거하여, "작고 개인적인" 트랜스포머와 연속적인 인지 상태를 가진 인간형 로봇을 가능하게 할 잠재력을 가진다. 본 연구는 기존 구성 요소들의 속성을 존중하면서 "만약 ~한다면(what-if)" 실험을 통해 도출된 결과임을 강조하며, 특정 구성(예: 왜 detach가 일부 사례에서 표준을 앞서는지 등)에 대한 열린 질문도 인정하고 있다.