기술 요약: AMOR (적응형 메타인지 출력 라우터)

문제 제기

재귀적 어텐션 하이브리드 모델은 선형 재귀 상태 공간 모델 (SSM) 의 추론 효율성과 어텐션 메커니즘의 표현력을 결합하는 것을 목표로 합니다. 그러나 기존 하이브리드 아키텍처는 일반적으로 고정 스케줄 전략에 의존하여, 모든 토큰에 대해 또는 사전에 정의된 레이어 (예: 매 $N$번째 레이어) 에서 어텐션을 균일하게 적용합니다. 이러한 접근 방식은 재귀 상태만으로도 다음 토큰을 높은 확신으로 예측할 수 있는 경우에도 어텐션의 완전한 이차 비용 ($O(N^2)$) 을 발생시킵니다. 결과적으로 이러한 모델들은 SSM 의 효율성 향상을 완전히 활용하지 못하며, 특히 불필요한 어텐션 블록의 계산 오버헤드가 병목 현상이 되는 긴 컨텍스트 시나리오에서 그렇습니다. 또한, 고정 하이브리드 모델은 특정 입력 난이도에 맞게 보정되지 않을 수 있는 어텐션 메커니즘을 강제로 적용하기 때문에 분포 변화 (예: 길이 외삽) 하에서 성능이 저하되는 경향이 있습니다.

방법론: AMOR 아키텍처

저자들은 AMOR(Adaptive Metacognitive Output Router, 적응형 메타인지 출력 라우터) 를 제안합니다. 이는 고정된 위치가 아닌 예측 불확실성에 기반하여 어텐션을 동적으로 할당하는 사후 (post-hoc) 하이브리드 아키텍처입니다.

핵심 아키텍처

AMOR 은 표준 재귀 백본 (Mamba2 및 Gated DeltaNet 으로 테스트됨) 을 기반으로 하며, 재귀 레이어 이후이지만 최종 LM 헤드 이전에 $K$(일반적으로 3) 개의 엔트로피 게이트 어텐션 블록을 추가합니다.

• 백본: 재귀 믹서는 고정된 크기의 상태를 유지하면서 전체 시퀀스를 처리합니다.
• 사후 블록: 재귀 레이어를 대체하는 직렬 하이브리드와 달리, AMOR 블록은 가산적입니다. 이들은 백본의 정규화된 출력에서 작동합니다.
• 비대칭성: 첫 번째 AMOR 블록 (블록 0) 은 백본의 최종 정규화된 잔차 (post-norm) 에 직접 작동하는 반면, 후속 블록 (1 및 2) 은 표준 pre-norm 공식을 사용합니다. 이 설계는 백본을 독립적이고 유효한 언어 모델로 유지하면서 AMOR 블록을 조건부 정제 단계로 취급합니다.

엔트로피 게이트 메커니즘

핵심 혁신은 인지 과학의 시스템 1/시스템 2 구분 (Kahneman, 2011) 에서 영감을 받은 경성 (hard), 무기울동 (gradient-free) 게이트 메커니즘입니다.

1. 불확실성 추정: 각 토큰에 대해 공유된 LM 헤드가 로짓 (logits) 을 계산합니다. 정규화된 출력 엔트로피 ($H_t$) 는 소프트맥스 분포에서 계산됩니다.
2. 적응형 임계값: 동적 임계값 $\tau$는 배치 중앙값 엔트로피 ($\mu$) 의 지수 이동 평균 (EMA) 에 스케일링된 표준 편차 ($\kappa\sigma$) 를 더하여 계산됩니다:
   $$\tau = \mu + \kappa\sigma$$
   이는 학습된 라우터 파라미터 없이도 학습 중 모델의 불확실성 분포 변화를 추적할 수 있게 합니다.
3. 조건부 실행:
   • $H_t > \tau$인 경우, 게이트가 발동 ($g_t=1$) 합니다: 토큰은 RoPE 를 포함한 인과적 어텐션 (SDPA) 을 거치며, 출력은 학습된 채널별 스케일링 인자 $\alpha$와 함께 잔차에 추가됩니다.
   • $H_t \le \tau$인 경우, 게이트는 비활성화 ($g_t=0$) 됩니다: 어텐션 계산이 건너뛰어지고 토큰은 변경 없이 통과합니다.
4. 기울기 흐름: 게이트는 미분 불가능합니다. 기울기는 발동 위치에서만 어텐션 가중치 ($W_Q, W_K, W_V, W_O$) 를 통해 흐릅니다. 중요한 점은, 현재 토큰이 어텐션을 건너뛰더라도 KV 캐시가 검색을 위해 유효하도록 유지되도록, $W_K$와 $W_V$는 미래의 발동 위치에서 기울기를 받습니다.

학습 및 추론

• 학습: AMOR 은 효율성을 위해 FlashAttention-2 를 사용하여 "마스크가 있는 밀집 (dense with mask)" 모드로 학습되며, 이는 진정한 희소 실행과 수학적으로 동일합니다. 임계값 $\tau$는 추론 시 고정됩니다.
• 추론: 모델은 조건부 건너뛰기를 수행합니다. 게이트가 꺼져 있을 때, 비용이 많이 드는 $Q$ 프로젝션, 어텐션 계산, $O$ 프로젝션이 우회되어 계산 자원이 절약됩니다. 그러나 컨텍스트 연속성을 유지하기 위해 $K$ 및 $V$ 프로젝션과 캐시 업데이트는 항상 발생합니다.

주요 기여

1. 불확실성 기반 라우팅: AMOR 은 고정된 레이어 스케줄에 의존하는 대신, 재귀 백본이 높은 예측 불확실성 (높은 엔트로피) 을 보일 때만 선택적으로 어텐션을 호출하는 간단하고 사후적인 메커니즘을 도입합니다.
2. 파라미터 효율적 하이브리드화: 이 방법은 게이트 및 어텐션 블록에 대한 최소한의 파라미터 오버헤드 (약 <5%) 를 추가하면서도 고정 스케줄 하이브리드와 비교할 수 있거나 더 나은 성능을 달성합니다.
3. 분포 변화에 대한 강건성: 낮은 불확실성 토큰에 대해 재귀 백본에 의존함으로써, AMOR 은 어텐션 메커니즘이 종종 위치 드리프트로 고통받는 길이 외삽 중 고정 하이브리드에서 관찰되는 성능 저하를 완화합니다.
4. 무기울동 게이팅: 이 접근법은 학습된 라우터나 강화 학습의 복잡성을 피하고, 모델 자체의 출력 분포에서 유도된 통계적 임계값을 사용합니다.

실험 결과

저자들은 Mamba2 및 Gated DeltaNet 백본을 사용하여 세 가지 규모 (180M, 440M, 1.5B 파라미터) 에서 AMOR 을 평가했으며, 순수 재귀 모델, 순수 Transformer, 고정 스케줄 직렬/퓨즈드 하이브리드와 비교했습니다.

• 상식 추론: AMOR 은 모든 규모와 벤치마크 (LAMBADA, HellaSwag, PIQA, ARC 등) 에서 순수 재귀 백본과 고정 스케줄 하이브리드 모두와 일치하거나 능가하는 성능을 일관되게 보여주었습니다.
• 컨텍스트 내 검색: 순수 재귀 모델은 정확한 복사 검색에 어려움을 겪는 반면, AMOR 은 백본 대비 크게 개선되어 고정 하이브리드와 경쟁력 있는 결과를 달성하면서도 **약 22%**의 토큰에서만 어텐션을 활성화했습니다.
• 긴 컨텍스트 성능: LongBench 벤치마크에서 AMOR 은 컨텍스트 길이가 증가함에 따라 안정적인 성능을 유지했습니다. 반면, 고정 스케줄 하이브리드와 순수 Transformer 은 RoPE 분포 변화와 메모리 제약으로 인해 성능이 크게 저하되었습니다. AMOR 의 예측 가능한 토큰에서 어텐션을 건너뛰는 능력은 더 우아하게 확장되도록 했습니다.
• 추론 효율성: AMOR 은 긴 컨텍스트 영역에서 더 나은 디코드 처리량 확장을 보여주었습니다. 이론적 분석과 실증적 벤치마크 (최대 448K 컨텍스트) 는 컨텍스트 길이가 특정 교차점 (440M 모델의 경우 약 256K) 을 초과하면 어텐션 건너뛰기로 인한 절감분이 엔트로피 계산 오버헤드를 능가함을 보여주었습니다.

중요성과 주장

이 논문은 언제 어텐션을 적용하는지가 얼마나 적용하는지만큼 중요하다고 주장합니다. 저자들은 다음과 같이 주장합니다:

1. 효율성과 표현력은 상충 관계가 아님: 입력 난이도에 기반한 라우팅 문제로 계산 할당을 처리함으로써, 모델은 어텐션의 전체 비용을 치르지 않고도 높은 표현력을 달성할 수 있습니다.
2. 간단함이 효과적이다: 복잡한 학습된 라우터나 보조 목적 함수는 필요하지 않습니다. 단순하고 출력 기반의 엔트로피 게이트만으로도 하이브리드 아키텍처를 효과적으로 조정할 수 있습니다.
3. 적응형 계산: AMOR 은 입력의 난이도에 맞춰 계산을 동적으로 조정하는 모델 클래스를 지향하며, 경량 상태 공간 백본의 강점을 유지하면서 모델이 불확실할 때만 무거운 어텐션 메커니즘을 선택적으로 배포합니다.

이 결과는 예측 불확실성에 기반한 어텐션의 선택적 할당이, 특히 긴 컨텍스트 생성 및 검색 작업에 있어 균일하거나 고정된 라우팅 전략에 대한 강건한 대안을 제공함을 시사합니다.