[Paper] Sessa: 선택적 State Space 어텐션

Source: arXiv - 2604.18580v1

위에 제공된 Source 링크만으로는 번역할 본문 내용이 포함되어 있지 않습니다. 번역을 원하는 텍스트(예: 초록, 본문, 결론 등)를 제공해 주시면, 요청하신 대로 한국어로 번역해 드리겠습니다.

개요

이 논문은 Sessa라는 새로운 디코더 아키텍처를 소개합니다. 이 아키텍처는 어텐션과 순환 상태‑공간 모델을 결합합니다. 선택적 어텐션 메커니즘을 피드백 루프 내부에 직접 삽입함으로써, Sessa는 과거의 많은 위치에서 정보를 동시에 집계할 수 있으며, 트랜스포머의 고전적인 (1/\ell) 드롭‑오프와 Mamba와 같은 모델의 지수적 감소보다 느린 멱법칙 메모리 감소를 달성합니다. 이는 신뢰할 수 있는 장거리 컨텍스트 처리가 필요한 작업에 특히 적합합니다.

주요 기여

• Selective State‑Space Attention (Sessa): 재귀 피드백 경로 안에 어텐션 연산을 배치하는 디코더로, 레이어당 다중 경로 집합을 가능하게 합니다.
• Theoretical memory analysis: 완화된 가정 하에 Sessa가 (0<\beta<1)인 파워‑법칙 메모리 꼬리 (O(\ell^{-\beta}))를 보이며, 이는 일반 Transformer의 (O(1/\ell)) 감소보다 느리고 기존 상태‑공간 모델의 지수 감소보다 확실히 우수함을 증명합니다.
• Flexible retrieval profiles: Sessa가 감소하지 않는(평탄한) 어텐션 가중치를 구현할 수 있음을 보여주며, 이는 다른 아키텍처가 명시적인 기억 트릭 없이 달성할 수 없는 특징입니다.
• Empirical validation: 장기 컨텍스트 언어 모델링 벤치마크에서 Sessa는 동일 규모의 Transformer와 Mamba‑스타일 모델을 능가하며, 단기 컨텍스트 작업에서도 경쟁력을 유지합니다.
• Unified view of retrieval mechanisms: Sessa를 “단일‑읽기” 검색(예: 어텐션)과 “단일‑피드백” 전파(예: 상태‑공간) 사이의 다리로 위치시켜, 두 접근이 상호 배타적이지 않음을 보여줍니다.

Methodology

1. Core architecture – Sessa는 표준 상태‑공간 (SSM) 재귀를 기반으로 합니다:
   [ h_t = A h_{t-1} + B x_t,\qquad y_t = C h_t, ]
   여기서 (h_t)는 은닉 상태, (x_t)는 입력 토큰 임베딩이며, (A,B,C)는 학습되는 행렬입니다.

2. Embedding selective attention – 은닉 상태를 업데이트하기 전에, Sessa는 선택된 과거 은닉 벡터들의 주의 가중 요약을 계산합니다:
   [ \tilde{h}{t-1} = \sum{k \in \mathcal{S}t} \alpha{t,k}, h_k, ]
   여기서 주의 점수 (\alpha_{t,k})는 경량 쿼리‑키 메커니즘에 의해 생성되며 입력에 의존할 수 있습니다 (예: 작은 MLP).

3. Feedback with many paths – 이제 재귀는 혼합된 상태를 사용합니다:
   [ h_t = A,\tilde{h}_{t-1} + B x_t, ]
   이는 모델이 단일 선형 체인 대신 하나의 단계에서 여러 과거 위치를 통해 정보를 라우팅할 수 있게 합니다.

4. Theoretical analysis – “확산 균일‑라우팅” 가정 하에 (주의 질량이 효과적 지원 (S_{\text{eff}}(t))에 대략 고르게 퍼져 있다고 가정), 저자들은 지연 (\ell)에 있는 토큰의 영향에 대한 닫힌 형태의 경계를 도출합니다:
   [ \text{influence}(\ell) = \Theta(\ell^{-\beta}),\quad 0<\beta<1. ]
   이 경계는 타이트함을 보이며, 라우팅 전략을 변경하지 않으면 개선될 수 없음을 의미합니다.

5. Training setup – 실험에서는 토크나이제이션, 옵티마이저 (AdamW), 학습률 스케줄, 그리고 계산 예산을 Sessa, Transformer, Mamba 베이스라인 모두에 동일하게 적용하여 공정한 apples‑to‑apples 비교를 보장합니다.

결과 및 발견

• 메모리 테일: 토큰 영향에 대한 실증 측정이 예측된 멱법칙 감소와 일치하여 이론적 분석을 확인합니다.
• 선택적 검색: 어텐션 온도를 조정함으로써 Sessa는 평탄한(감소하지 않는) 영향 곡선을 생성할 수 있으며, 모든 토큰이 동일하게 접근 가능해야 하는 작업에 유용합니다(예: 대규모 임포트를 포함한 코드 완성).
• 효율성: 추가 어텐션 단계에도 불구하고, 토큰당 전체 FLOP 수는 어텐션이 압축된 은닉 벡터 집합(대개 저차원 투영)에서 작동하기 때문에 표준 SSM 레이어와 비교할 만합니다.

실용적 함의

• Long‑document processing – LLM 기반 어시스턴트, 법률 문서 분석기, 과학 논문 요약기를 개발하는 개발자들은 Sessa의 느린 메모리 소멸을 활용할 수 있습니다. 이는 외부 검색 모듈 없이도 수 킬로바이트 규모 컨텍스트의 초반 정보를 보존합니다.
• Hybrid architectures – Sessa는 트랜스포머의 병렬성(attention)과 상태공간 모델의 재귀성을 결합하는 깔끔한 방법을 보여주며, 지연 시간과 메모리 유지 사이의 균형을 맞추는 맞춤 레이어 개발의 가능성을 열어줍니다.
• Memory‑efficient inference – 주의(attention)가 재귀 루프 내부에서 압축된 은닉 표현에 대해 수행되므로, Sessa는 제한된 GPU 메모리를 가진 엣지 디바이스에서도 동일한 크기의 일반 트랜스포머보다 더 긴 컨텍스트를 처리할 수 있습니다.
• Fine‑tuning flexibility – 선택적 attention 가중치를 독립적으로 고정하거나 재학습할 수 있어, 전체 모델을 재학습하지 않고도 도메인별 요구에 맞게 “검색 프로필”(예: 보다 균일하거나 보다 뾰족한)을 조정할 수 있습니다.

제한 사항 및 향후 연구

• 확산 라우팅 가정 – 이론적 보장은 비교적 균일한 어텐션 분포에 의존한다; 매우 피크가 높은 어텐션은 더 빠른 감쇠로 돌아갈 수 있으며, 이는 현재 분석에 포함되지 않는다.
• 수십억 파라미터로의 확장성 – 실험은 ~300 M‑파라미터 모델에 제한되었으며, Sessa가 최첨단 LLM 규모로 확장될 때 어떻게 동작하는지는 아직 확인되지 않았다.
• 하드웨어 친화성 – FLOP‑기준으로는 비슷하지만, 혼합된 attention‑in‑recurrence 패턴이 기존 transformer‑최적화 커널에 완벽히 매핑되지 않을 수 있어, 최대 처리량을 위해 맞춤형 커널이 필요할 수도 있다.
• 다양한 작업 – 논문은 언어 모델링에 초점을 맞추고 있으며, Sessa를 비전, 음성, 혹은 멀티모달 스트림에 적용하려면 추가 엔지니어링이 필요하고 새로운 과제가 드러날 수 있다.

저자들이 제시한 향후 연구 방향으로는 비균일 어텐션에 대한 이론적 프레임워크 확장, 학습된 희소성 패턴을 통합하여 계산량을 추가로 감소시키는 방안, 그리고 기존 대규모 LLM에 플러그인 모듈로 Sessa를 적용해 전체 재학습 없이 장거리 메모리를 복원하는 연구 등이 있다.

저자

• Liubomyr Horbatko

논문 정보

• arXiv ID: 2604.18580v1
• 카테고리: cs.LG, cs.AI, cs.CL
• 출판일: 2026년 4월 20일
• PDF: Download PDF