소개

딥러닝 아키텍처의 계산 효율성은 장기 문맥 시퀀스 처리에서 자연어 처리에 있어서 중요한 병목 현상으로 떠올랐습니다. Transformer 기반 모델들은 다양한 작업에서 최고 성능을 달성했지만, 시퀀스 길이와 관련하여 메모리 및 계산 측면에서 2차 방정식의 복잡성을 가지고 있습니다. 자가 주의 메커니즘은 $`N \times N`$ 주의 행렬을 계산하며, 여기서 $`N`$은 시퀀스 길이이고 이를 통해 $`O(N^2)`$의 메모리 소비와 추론 시간을 초래합니다. 이러한 2차 방정식적인 확장성은 Transformer를 장기 문맥 상황에 실제로 적용하는 데 제한을 두고 있으며, 사용자들은 시퀀스를 자르거나 슬라이딩 윈도우 및 계층적 처리와 같은 연산적으로 비용이 많이 드는 기술을 활용해야 합니다.

최근 State Space Models (SSMs)은 시퀀스 모델링에 있어서 Transformer의 유망한 대안으로 떠올랐습니다. 특히 Mamba 아키텍처는 선택적인 상태 공간 메커니즘을 사용하여 고정 크기의 은닉 상태를 통해 시퀀스를 순환적으로 처리하며, 이로 인해 메모리 및 계산 측면에서 선형 $`O(N)`$ 복잡성을 갖게 됩니다. 이론적 분석은 SSM들이 긴 시퀀스에 있어서 Transformer보다 크게 우수하게 수행될 것이라는 것을 제시하고 있으며, 실제 보고서는 5배 더 높은 처리량과 백만 토큰 시퀀스를 처리할 수 있는 능력을 보여주었습니다. 그러나 최근 연구에서는 기초적인 타협점이 드러났습니다: SSM들은 계산 효율성에 있어서 우수하지만, 고정 크기의 잠재 상태 때문에 문맥 학습 및 결합적 재현력에서 성능을 희생할 수 있습니다.

SSMs에 대한 관심이 증가함에도 불구하고, 실제 세계의 장기 문맥 응용 분야를 위한 SSMs와 Transformer 간의 체계적인 경험적 비교는 여전히 제한적입니다. 대부분의 기존 연구는 언어 모델링을 위한 perplexity 지표나 합성 작업에 집중하고 있으며, 이러한 아키텍처들이 복잡한 의존성을 가진 실제 발생하는 긴 시퀀스를 처리할 때 어떻게 수행되는지에 대한 중요한 질문들은 아직 해결되지 않았습니다. 또한, 계산 효율성이 이론적으로 잘 연구되었지만 다양한 시퀀스 길이에 걸쳐 세심하게 관리된 실험 조건에서 경험적인 검증은 드뭅니다.

본 논문은 이러한 간극을 메우기 위해 Mamba 아키텍처와 LLaMA Transformer를 장기 문맥 이중 치료 세션에 대해 철저한 경험적 벤치마킹 연구를 통해 비교합니다. 우리는 치료 세션이 실제 세계의 장기 문맥 처리 대표 사례로 선택되었음을 나타냅니다: 치료 세션은 자연스럽게 6,000-10,000 토큰을 포함하며 복잡한 시간 의존성을 가지고 있으며, 자르지 않고 전체 세션 동안의 동적 변화를 포착해야 합니다. 우리의 연구는 두 아키텍처가 다음과 같은 두 가지 보완적인 차원에서 체계적으로 평가됩니다:

(1) 계산 효율성: 우리는 토큰 수 512에서 8,192까지의 시퀀스 길이에 걸쳐 메모리 사용량과 추론 시간을 측정하고 경험적 확장 방정식을 도출하며 SSMs가 Transformer보다 엄격하게 효율적이 되는 중요한 교차점들을 식별합니다.

(2) 표현 효율성: 우리는 Mamba의 은닉 상태 동력학과 Transformer의 주의 패턴이라는 두 아키텍처의 내부 메커니즘을 분석하여 장기 문맥 유지 및 입력 시퀀스에서 동적 변화 감지 능력을 평가합니다.

정교하게 맞춘 모델 구성(각 50M 파라미터), 동일한 입력 시퀀스, 철저한 벤치마킹 프로토콜을 통해 제공되는 포괄적인 경험적 증거를 통해 이 연구는 장기 문맥 시퀀스 모델링을 하는 연구자와 실무자를 위한 실행 가능한 통찰력을 제공합니다. 우리의 발견은 SSMs가 Transformer보다 우월한 구체적인 조건, 이러한 우월성의 크기 및 계산 효율성을 향상시키면서 수반되는 표현력 타협점을 밝혀냅니다. 치료 세션을 테스트 케이스로 사용했지만, 이 아키텍처 통찰은 장기 문맥 처리가 필요한 모든 도메인에 일반화되며, 그 중에는 유전체학, 오디오 처리, 장문 분석 및 멀티턴 대화 시스템을 포함합니다.

관련 연구

State Space Models과 Transformers의 벤치마킹

최근 경험적 연구들은 SSMs와 Transformer를 여러 성능 차원에서 체계적으로 비교해 왔습니다. Waleffe et al.은 소비자 등급 GPU에서 Transformer, Mamba 및 Mamba-2, 그리고 하이브리드 모델을 포괄적인 벤치마킹하여 SSMs가 24GB 메모리 제한 내에서 220K 토큰 길이를 달성했다는 것을 밝혔습니다. 이는 Transformer보다 약 3배 더 긴 시퀀스입니다. 그들의 연구 결과에 따르면, 8B 파라미터 규모에서는 순수한 Mamba 모델들이 표준 언어 모델링 작업에서 Transformer와 동등하거나 이를 초월했지만 강력한 결합적 재현력을 요구하는 벤치마크에서는 크게 뒤처졌습니다. 하이브리드 아키텍처(Mamba-2-Hybrid)는 단지 7-8%의 자가 주의 계층을 포함하여 이 성능 차이를 줄였습니다.

Jelassi et al.의 이론적 및 경험적 분석은 두 아키텍처 사이에 근본적인 표현력 차이를 드러냈습니다. 그들의 연구는 Transformer가 2계층 아키텍처로 지수 길이의 문자열을 복사할 수 있지만, SSM은 고정 크기 잠재 상태 때문에 이에 제한적이라는 것을 보여주었습니다. 160M 파라미터 모델들의 경험적 연구는 Transformer가 복사 작업에 Mamba보다 100배 더 적은 훈련 데이터를 필요로 함을 나타냈습니다. 사전 학습된 모델들은 비슷한 perplexity 점수에도 불구하고, Transformer가 SSM보다 문맥 복사와 정보 검색 작업에서 크게 우수하다는 것을 드러냈습니다. Wang et al.은 표현력 유동 동역학을 통해 이러한 차이를 추가로 분석하여, Transformer는 초기 과평탄화를 겪고 나중에 회복하는 반면 SSM은 초기 토큰 독특성을 유지하지만 나중에 균일화되어 장기 문맥이 늘어날수록 체계적으로 다른 실패 모드를 보인다는 것을 설명했습니다.

아키텍처 효율성 및 확장법칙

Gu와 Dao가 도입한 Mamba 아키텍처는 선택적인 상태 공간을 통해 Transformer보다 5배 더 높은 추론 처리량을 달성하고, 백만 길이 시퀀스까지 선형적으로 확장됩니다. 그 후에 개발된 Mamba-2 아키텍처는 구조화된 준분리 가능 행렬을 통해 SSMs와 주의 메커니즘 사이의 이론적 연결을 수립하고, 추론 속도를 2-8배 빠르게 했습니다. 경험적인 평가 결과 Mamba-3B는 동일한 크기의 Transformer보다 우수하며 두 배 크기의 모델과 비교할 때 언어 모델링 벤치마크에서 대등한 성능을 보였습니다.

LLaMA 모델 시리즈는 Transformer 아키텍처의 최첨단으로, 그룹화된 쿼리 주의(GQA), RoPE 위치 인코딩, RMSNorm 계층 정규화 및 SwiGLU 활성화를 포함한 최적화를 사용합니다. 그러나 자가 주의 메커니즘의 제곱 복잡성이 기본적인 병목 현상입니다. LLaMA 2는 문맥 길이가 4,096 토큰으로 한정되어 있으며, 메모리 및 계산 요구사항은 제곱적으로 증가합니다—시퀀스 길이가 두 배로 늘어날 때마다 메모리 소비량과 추론 시간은 네 배씩 증가합니다. FlashAttention 및 효율적인 주의 구현은 큰 개선을 제공하지만 근본적으로 $`O(N^2)`$ 확장 행동을 바꾸지는 않습니다.

임상 NLP 응용 프로그램에서 계산 효율성

심리치료 연구에서 자연어 처리 응용 프로그램은 치료 세션의 장시간으로 인해 고유한 계산적 도전에 직면합니다. Imel et al.은 Transformer 기반 모델들이 전체 치료 세션을 처리할 때 심각한 메모리 제약을 겪는다는 것을 보여주었습니다. 입력 길이와 함께 요구사항이 지수적으로 증가하여 연구자들은 세션을 짧은 단위로 분석해야 하며, 장기 의존성을 분석하는 데 실패하게 됩니다. 그들의 RoBERTa를 사용한 386명의 클라이언트로부터 1,235개의 세션에 대한 작업에서는 입력 길이가 증가함에 따라 계산 복잡성이 지수적으로 증가하여 연구자들이 세션을 짧은 단위로 처리하고 장기 의존성을 분석하는 데 실패하게 됩니다.

Flemotomos et al.은 세션 텍스트로부터 치료 동맹을 추정하기 위한 Working Alliance Transformer(WAT)를 개발했지만, Transformer 아키텍처가 긴 문맥을 처리하더라도 장기 의존성을 유지하는 데 어려움이 있다는 것을 지적했습니다. 그들의 분석은 평균 주의 거리가 67 토큰에 머물러 있음을 드러내며, 이는 전체 세션 동안의 동적 변화를 포착하는 데 근본적인 제약을 의미합니다. 치료 상호작용의 종합 분석을 위한 효율적인 장기 문맥 아키텍처 개발은 여전히 중요하며, 여기서 전체 세션에 걸친 “치유 곡선"을 포착하는 것이 합의 형성 및 세션 결과 이해에 필수적입니다.

이론적 기반

이 섹션에서는 Transformer와 State Space Model(SSM) 아키텍처를 비교하기 위한 이론적인 기초를 설정하며, 그들의 계산과 표현적 특성을 중점으로 합니다.

계산 효율성 이론

Transformer 메모리: 제곱 확장

자가 주의 메커니즘은 $`N \times N`$ 행렬 $`A = \operatorname{softmax}(QK^T / \sqrt{d_k})`$, 여기서 $`Q, K \in \mathbb{R}^{N \times d_k}`$를 통해 주의 점수를 계산합니다. 단일 헤드에 대해 필요한 메모리는 $`M_{\text{attn}} = N^2 w`$이며 여기서 $`w`$는 바이트 단위의 단어 크기입니다. 이는 $`h`$ 개의 헤드와 $`L`$ 층으로 확장되었을 때 전체 주의 메모리 및 결과적인 경험적 적합 모델은 다음과 같이 정의됩니다:

\begin{align}
    M_{\text{total-attn}} &= L \cdot h \cdot N^2 \cdot w \label{eq:transformer_memory_total} \\
    M_T(N) &= \alpha N^2 + \beta N + \gamma \label{eq:transformer_memory_fit}
\end{align}

이 모델에서 $`\alpha N^2`$ 항은 주의 행렬 저장을 나타내고, $`\beta N`$ 항은 임베딩 및 피드포워드 네트워크(FFN) 활성화와 같은 선형 구성 요소를 설명합니다. 상수 $`\gamma`$는 모델 매개변수로부터의 고정된 오버헤드를 포착합니다.

State Space Model 메모리: 선형 확장

Transformer와 달리 SSM은 상태 업데이트 방정식 $`h_t = \bar{A}h_{t-1} + \bar{B}x_t`$ 및 $`y_t = Ch_t`$를 통해 토큰을 순환적으로 처리합니다. 은닉 상태 $`h_t \in \mathbb{R}^{d_{\text{state}}}`$가 고정되어 있기 때문에, 상태 메모리 $`M_{\text{state}} = L \cdot d_{\text{state}} \cdot w`$는 시퀀스 길이와 무관하게 일정합니다. $`N`$에 의존적인 메모리는 입력 임베딩과 중간 활성화로 제한되며, 이들은 모두 $`O(N \cdot d_{\text{model}})`$로 확장됩니다. 따라서 SSM의 경험적 메모리 모델은 다음과 같이 단순화됩니다:

\begin{equation}
    M_M(N) = \alpha' N + \gamma' \label{eq:ssm_memory_fit}
\end{equation}

여기서 $`\alpha' N`$ 항은 시퀀스에 의존적인 활성화를 포착하고, $`\gamma'`$는 고정된 매개변수 오버헤드를 포착합니다.

효율성 지표 및 추론 확장

메모리 효율 비율 $`\rho_M(N) = M_T(N) / M_M(N)`$은 Transformer의 메모리 오버헤드를 양화하며, 이는 선형적으로 $`N`$에 따라 증가합니다. 이러한 확장성 차이는 추론 지연에도 반영됩니다. Transformer의 자가 주의는 층당 $`O(N^2 \cdot d_k)`$ 연산을 수행하지만, SSM의 선택적 스캔은 단지 $`O(N \cdot d_{\text{state}})`$ 시간이 필요합니다. 길이 $`N`$에 대한 예상 시간 모델은 다음과 같습니다:

\begin{align}
    \text{Transformer:} \quad T_T(N) &= aN^2 + bN + c \label{eq:transformer_time} \\
    \text{SSM:} \quad T_M(N) &= a'N + c' \label{eq:ssm_time}
\end{align}

이는 Transformer의 경우 시퀀스 길이가 두 배로 늘어날 때 추론 시간이 네 배로 증가하지만, SSM은 그저 두 배로 증가한다는 것을 의미하며, 장기 문맥 응용 프로그램에 대한 효율성 격차를 강조합니다.

표현 효율성 이론

State Space Models의 은닉 상태 동력학

SSMs에서는 세션 내에서 은닉 상태 $`h_t`$의 진화를 추적하여 모델이 치료 상호작용 내의 역동적인 변화를 어떻게 나타내는지 측정합니다. 주요 지표들은 다음과 같습니다:

은닉 상태 속도: 연속된 은닉 상태 간의 변경율을 측정합니다:

\begin{equation}
v_t = \|h_t - h_{t-1}\|_2 \label{eq:state_velocity}
\end{equation}

높은 속도는 주요 치료 순간에 해당할 수 있는 빠른 표현 변화를 나타냅니다.

은닉 상태 이동: 초기 상태로부터의 누적 편차를 측정합니다:

\begin{equation}
d_t = \|h_t - h_0\|_2 \label{eq:state_drift}
\end{equation}

이동은 세션 동안 총 표현 변화를 반영하며 치유 곡선을 나타냅니다.

계층별 분석: 각 계층에서 $`v_t`$와 $`d_t`$를 살펴봄으로써 역동적인 처리가 가장 적극적으로 일어나는 네트워크의 어느 부분인지 식별할 수 있습니다.

Transformer의 주의 메커니즘

Transformer에서는 주의 행렬을 분석하여 컨텍스트 활용 방식을 이해합니다:

주의 엔트로피: 주의 분포의 확산성을 측정합니다:

\begin{equation}
H(A_i) = -\sum_{j=1}^{N} A_{ij} \log A_{ij} \label{eq:attention_entropy}
\end{equation}

높은 엔트로피는 전역 주의를 나타내며 낮은 엔트로피는 로컬 주의를 나타냅니다.

평균 주의 거리: 모델이 얼마나 먼 과거까지 주목하는지를 측정합니다:

\begin{equation}
\bar{d}_i = \sum_{j=1}^{N} A_{ij} \cdot |i - j| \label{eq:attention_distance}
\end{equation}

이 지표는 모델이 장기 의존성을 유지하는지 또는 주로 로컬 컨텍스트에 초점을 맞추는지를 나타냅니다.

효과적인 컨텍스트 윈도우

“효과적인 컨텍스트 윈도우"를 모델의 예측에 의미 있게 기여하는 토큰 범위로 정의합니다:

• SSMs: 과거 토큰이 현재 예측에 어떤 영향을 미치는지 식별하기 위해 그래디언트 기반 속성을 사용
• Transformers: 주목 범위 분석을 통해 일반적인 주목 토큰 범위를 결정

방법론

데이터셋

공개된 장기 문맥 이중 치료 데이터셋의 제한성 때문에, 우리는 OpenAI GPT-4.5를 사용하여 치료 세션 대본을 합성적으로 생성했습니다. 총 4개의 세션이 생성되었으며 각 세션은 약 50분 간의 의사와 두 명의 클라이언트 간의 상호작용을 시뮬레이션합니다. 세션 설계는 실제 치료 동적을 포착하기 위해 감정 조율, 커뮤니케이션 패턴, 관계 동적, 불안 관리를 포함합니다. 생성 프롬프트는 자연스러운 역할 교대, 다양한 감정 강도, 반영적인 청취 및 공감적 응답과 같은 증거 기반 치료 기술, 그리고 오프닝 랩, 작업 단계에서부터 마무리까지의 실제 세션 진행을 지정합니다.

모든 세션은 HuggingFace transformers의 GPT-2 토크나이저를 사용하여 토큰화되었습니다. 각 말하기는 화자 역할(예: “Therapist:” 또는 “Client:")으로 접두사가 붙여졌습니다. 세션은 JSON 형식으로 저장되었으며 구조화된 대화 배열에는 화자와 텍스트 필드가 포함됩니다.

모델 구성

연구팀은 모델 크기의 차이보다 아키텍처 효율성을 공정하게 비교하기 위해 각각 약 50만 개의 파라미터를 가진 두 모델을 구성했습니다. 이를 통해 벤치마킹 결과는 본질적인 아키텍처 차이를 반영하며 모델 크기의 불일치를 제외할 수 있습니다.