MERLIN이 다중 상태 선형 어텐션의 세그먼트 집계를 개선한다

무엇인가

이 논문은 긴 시퀀스 모델링에서 self-attention의 이차 시간 복잡도와 선형으로 커지는 KV 캐시를 줄이려는 흐름을 다룬다. 단일 고정 크기 메모리 상태를 쓰는 선형 어텐션은 긴 문맥을 압축할 때 정보 손실이 생기고, 이를 완화하려고 문맥을 여러 세그먼트로 나눠 각각 메모리 상태에 압축하는 다중 상태 선형 어텐션(MSLA)이 제안됐다. MSLA에는 문맥을 어떻게 나눌지(separation)와 세그먼트별 출력을 어떻게 합칠지(aggregation)라는 두 문제가 있는데, 저자들은 상대적으로 덜 연구된 aggregation을 핵심 병목으로 본다. 기존 방법은 위치나 질의와 각 세그먼트 키의 평균 정렬만으로 가중치를 정해, 세그먼트 안에서 정렬이 어떻게 분포하는지 놓친다.

어떻게 동작하나

MERLIN은 각 세그먼트 키의 1차 모멘트와 2차 모멘트를 저장해 이 문제를 푼다. 토큰마다 기존 메모리 상태용 키·질의와 별도로 더 작은 차원 d_r의 키 k'_t와 질의 q'_t를 투영한다. 세그먼트 i에 대해 1차 모멘트 μ_i는 세그먼트 내 키들의 평균 벡터, 2차 모멘트 Σ_i는 키 외적 k'_j k'_j^T의 평균이다. 이 둘은 고정 크기라 세그먼트가 길어져도 요약 비용이 일정하고, μ_i는 키의 평균 방향, Σ_i는 키가 세그먼트 안에서 얼마나 퍼져 있는지를 담는다.

무엇과 다른가

질의가 주어지면 MERLIN은 μ_i와 Σ_i로부터 평균 정렬 rμ_t,i = μ_i^T q'_t와 분산 rΣ_t,i = q'_t^T Σ_i q'_t - (rμ_t,i)^2를 계산한다. 평균만 보면 세그먼트 안에 강하게 맞는 키 하나가 있는 경우와 모든 키가 약하게 맞는 경우를 구분하기 어렵지만, 분산은 강한 정렬이 존재하는지를 드러낸다. 최종 집계 가중치는 λ_t,i = γ_t softmax_i(τμ rμ_t,i) + (1-γ_t) softmax_i(τΣ rΣ_t,i)로 정한다. γ_t는 토큰 은닉 상태에서 나오는 스칼라 게이트이고 τμ, τΣ는 학습되는 온도다. 이 가중치로 각 세그먼트 출력을 가중합해 최종 출력을 만든다. 저자들은 토큰 수준 softmax attention의 질량을 세그먼트별로 묶으면 세그먼트의 로그 증거가 평균과 분산, 즉 정렬의 첫 두 누적률로 근사된다는 점도 보인다. MERLIN은 고정 계수 대신 평균 점수와 분산 점수를 각각 softmax로 정규화한 뒤 학습 게이트로 섞어, 2차 모멘트를 쓰는 일반화된 라우팅으로 만든다.

어떻게 쓰나

실험은 Gated DeltaNet-1.3B와 Mamba-2-780M의 다중 상태 변형에서 이뤄졌다. Gated DeltaNet-1.3B는 FineWeb-edu 50B 토큰으로 사전학습했고, Mamba-2-780M은 Pile 300B 토큰으로 사전학습된 공개 체크포인트를 썼다. 단일 상태 모델을 다중 상태로 바꾸기 위해 256토큰 균일 세그먼트와 16K 문맥, 최대 64개 세그먼트를 두고 mid-training을 한다. 1단계는 Long-Data-Collections 2B 토큰 중 10%(0.2B)를 RULER NIAH 유사 과제로 증강하고, 2단계는 60M 토큰의 절반을 일반 LDC, 절반을 RULER HotpotQA와 SQuADQA 합성 샘플로 구성한다. 비교 대상은 단일 상태 그대로인 Vanilla, 위치 기반 가중치의 Position, 평균 정렬 기반의 Meanpool이다.

전제와 한계

RULER NIAH와 QA를 16K 문맥에서 평가한 결과, Gated DeltaNet에서 MERLIN은 평균 62.7로 Meanpool 57.9와 Position 54.6을 앞섰다. Mamba-2에서는 76.2로 Meanpool 71.2와 Position 69.6을 앞섰고, Meanpool 대비 7.0%, Position 대비 9.5% 상대 개선이다. LongBench 16K 이하 과제에서는 Gated DeltaNet이 21.9 대 19.3으로 13.5%, Mamba-2가 20.6 대 18.4로 12.0% 상대 개선을 냈다. 특히 MERLIN은 4배 더 많은 메모리를 쓰는 베이스라인보다 나은 성능을 보였다. 기본 세그먼트 길이 256에서 MERLIN이 64토큰 세그먼트를 쓰는 Meanpool과 Position보다 높은 점수를 낸 것이다. HELMET TriviaQA 검색 증강 생성 과제의 위치 편향 분석에서도 Vanilla는 문맥 끝에 답이 있을 때만 잘 맞히는 강한 최신성 편향을 보였고, Position은 약한 최신성 편향이 남았으며, Meanpool과 MERLIN은 위치 전반에 걸쳐 비교적 고르게 작동하되 MERLIN이 전체적으로 더 높았다. 세그먼트 길이를 64, 128, 256, 512로 바꿔도 MERLIN은 모든 길이에서 베이스라인을 앞섰고, 세그먼트가 길어져도 평균 기반 집계에 대한 우위가 유지됐다.

추가 메모리와 파라미터 부담은 크지 않다고 저자들은 보고한다. d_r=128에서 2차 모멘트를 위한 추가 메모리는 Gated DeltaNet에서 최대 12.6%, Mamba-2에서 4.2%이고, d_r이 작아지면 메모리 비용은 제곱으로 줄어든다. MERLIN의 파라미터 오버헤드는 Gated DeltaNet 4.6%, Mamba-2 2.2%로 Meanpool과 비슷한 수준이다. 다만 d_r을 16, 32, 64, 128로 바꾼 실험에서는 d_r이 커질수록 LongBench 평균이 올라가, 표현력과 메모리 사이의 트레이드오프가 있다. 키·질의를 기존 메모리 상태와 공유하거나 1·2차 모멘트에 감쇠와 쓰기 강도(decay, write-strength)를 넣는 변형은 성능을 크게 떨어뜨려, 전용의 작은 키·질의와 단순한 모멘트가 중요하다는 ablation도 제시된다.

개발자 관점에서 이 논문은 MSLA 계열의 긴 문맥 모델에서 세그먼트 집계 가중치를 더 내용 인식적으로 만드는 비교적 값싼 레버를 보여준다. 긴 문서 QA나 검색 증강 생성처럼 여러 구간 중 어디에 답이 있는지 찾아야 하는 작업에서, 위치 기반 최신성 편향을 줄이고 세그먼트 선택을 개선하려는 경우 참고할 만하다. 다만 단일 상태 체크포인트를 그대로 쓰는 것이 아니라 NIAH 유사 데이터와 추론 데이터를 섞은 mid-training이 필요하고, 세그먼트 길이와 d_r에 따라 메모리와 표현력이 달라지므로 자신의 문맥 길이와 메모리 예산에서 재검증해야 한다.

저자들이 별도의 한계 절을 두지는 않았지만, 원문에서 확인되는 전제는 분명하다. 이 연구는 aggregation만 다루고 separation은 256토큰 균일 분할로 고정했으며, 평가는 16K 문맥과 1.3B·780M 두 모델 규모에 집중한다. 16K를 넘는 문맥이나 다른 모델 규모에서의 일반화는 원문에 제시되지 않았다. 또한 2차 모멘트의 차원 d_r이 성능과 메모리의 트레이드오프를 만들고, 다중 상태로의 전환에 mid-training 비용이 든다는 점을 실무에서는 전제로 삼아야 한다.