기술 요약: Neural Attention Search Linear (NAtS-L)

문제 정의

표준 소프트맥스(softmax) 기반 셀프 어텐션의 이차 복잡도($O(L^2)$)는 트랜스포머(Transformer)의 긴 문맥(long-context) 시나리오에서 중요한 병목 현상이 되었습니다. 선형 어텐션 모델은 과거의 Key-Value(KV) 값을 고정된 크기의 은닉 상태(hidden state)로 압축하여 선형 복잡도($O(L)$)를 달성함으로써 유망한 대안을 제시하지만, 고정된 상태 크기의 제약으로 인해 표현력(expressivity)이 제한되는 경우가 많습니다. 반면, 표준 트랜스포머는 모든 KV 문맥 정보를 보존하여 더 나은 장기 검색(long-term retrieval)을 가능하게 하지만, 높은 메모리 및 계산 비용을 초래합니다.

기존의 하이브리드 접근 방식은 소프트맥스와 선형 어텐션을 계층 간에 교차 배치하거나, 이들 사이를 전환하기 위해 고정된 스케줄을 사용하는 방식으로 이를 완화하려고 시도합니다. 그러나 이러한 방법들은 다양한 문맥에 필요한 정보를 유연하게 포착하지 못하는 고정된 구조를 유지하는 경우가 많으며, 전체 소프트맥스 레이어의 존재로 인해 효율성이 저해됩니다. 또한, 희소 어텐션(sparse attention) 변형들은 인간의 휴리스틱에 의존하는 경향이 있어, 정의된 희소 맵(sparse maps)이 커버하지 못하는 상관관계를 무시할 가능성이 있습니다.

방법론

본 논문은 동일한 레이어 내에서 서로 다른 토큰에 대해 선형 및 소프트맥스 어텐션 연산을 모두 적용하는 프레임워크인 **Neural Attention Search Linear (NAtS-L)**를 제안합니다. NAtS-L은 고정된 아키텍처 대신, 입력 문맥에 따라 각 토인 블록(token block)에 최적의 어텐션 연산을 자동으로 결정합니다.

핵심 메커니즘

1. 토큰 수준의 하이브리드화(Token-Level Hybridization): 입력 시퀀스는 청크(chunk) 단위로 분할됩니다. 어텐션 스코어 레이어(평균 풀링 후 선형 레이어로 구성)는 각 청크를 서로 다른 연산 유형(소프트맥스 어텐션 vs. 선형 어텐션의 변형인 Gated DeltaNet)을 나타내는 일련의 점수 세트로 매핑합니다. 가장 높은 점수를 받은 연산이 해당 청크에 할당됩니다.
2. 통합된 탐색 공간(Unified Search Space): 이 프레임워크는 비선형 소프트맥스와 선형 어텐션 연산을 통합하여 단일 레이어 내에서 처리할 수 있도록 합니다. 모델은 각 토큰 블록의 중요성을 학습하며, 장기적 영향력이 큰 토큰(장기 검색이 필요한 토큰)에는 소프트맥스 어텐션을 적용하고, 단기적 영향력이 있는 토큰에는 선형 어텐션을 적용합니다.
3. 청크 단위 계산(Chunk-wise Computation): 하드웨어 효율성을 유지하기 위해 NAtS-L은 청크 단위 병렬 방식을 사용합니다. 청크 내에서 선형 및 소프트맥스 상관관계가 모두 계산됩니다. 모델은 두 연산 유형에 대한 은닉 상태를 모두 유지합니다. 새로운 청크가 도착하면 모델은 두 상태를 의사 업데이트(pseudo-update)하지만, 선택되지 않은 연산 유형의 경우 은닉 상태를 청크 시작 지점으로 리셋하거나 롤백하여, 선택된 연산 유형에 필요한 KV 값만 유지되도록 합니다.
4. 출력 병합(Output Merging): 선택된 소프트맥스 및 선형 연산의 출력은 각각 정규화(RMS normalization 사용)된 후 가중 합(weighted sum)을 통해 결합됩니다. 이 결합을 위한 가중치는 쿼리($Q$) 프로젝션으로부터 동적으로 생성되며, 이를 통해 모델은 두 어텐션 유형의 기여도를 조절하는 법을 학습합니다.
5. 경사 기반 최적화(Gradient-Based Optimization): 어텐션 유형의 선택은 엔드투엔드(end-to-end)로 학습됩니다. 어텐션 스코어 레이어에 대해 기울기(gradient)가 계산되므로, 모델은 보조 손실(auxiliary loss)이나 사전 학습된 모델 없이도 연산 선택과 모델 가중치를 공동으로 최적화할 수 있습니다.

주요 기여

1. 통합 프레임워크: 저자들은 비선형 소프트맥스와 선형 어텐션에 대한 통합된 설명을 제공하여, 이들이 단일 레이어 내에서 처리되고 선택될 수 있도록 했습니다.
2. 적응형 아키텍처: NAtS-L은 고정된 레이어별 비율에 의존하는 대신, 입력 문맥 정보에 따라 토큰 수준(청크 수준)에서 최적의 어텐션 연산을 자동으로 결정하는 프레임워크로 제안되었습니다.
3. 효율성 및 표현력: NAtS-L은 성능과 효율성 사이의 강력한 균형을 달성하여, 순수 선형 모델보다 더 나은 장기 문맥 모델링 능력을 제공하면서도 전체 소프트맥스 트랜스포머보다 낮은 계산 지연 시간을 유지함을 입증했습니다.

실험 결과

저자들은 Fineweb-Edu 데이터셋을 사용하여 380M 및 800M 파라미터 규모의 모델로 학술적 규모의 사전 학습 태스크에 대해 NAtS-L을 평가했습니다.

• 언어 모델링 및 추론: 제로샷 상식 추론 벤치마크(예: LAMBADA, PIKA, Hellaswag)에서 NAtS-L 및 그 하이브리드 변형은 순수 Gated DeltaNet (GDN), Mamba2, Transformer 베이스라인과 대등하거나 약간 더 나은 성능을 보였습니다.
• 검색 태스크: 실제 검색 태스크(예: SWDE, SQD, DROP)에서 NAtS-L Hybrid는 가장 우수한 평균 점수를 기록하며, 순수 선형 모델(특정 태스크인 FDA 등에서 어려움을 겪음)과 순수 트랜스포머(SQD에서 어려움을 겪음)를 모두 능가했습니다. 이는 모델이 필요한 문맥 정보를 보존하는 데 있어 견고함을 보여줍니다.
• 길이 외삽(Length Extrapolation): 최대 65,536 토큰 문맥(PG19, NarrativeQA, CodeParrot)에서 평가했을 때, NAtS-L Hybrid는 입력 문맥이 훈련 길이(4096 토큰)를 초과하더라도 퍼플렉시티(perplexity) 수준을 유지한 반면, 순수 선형 하이브리드(GDN Hybrid)는 효과적인 외삽에 실패했습니다.
• 장기 문맥 벤치마크: RULER 및 LongBench에서 NAtS-L Hybrid는 비교된 모델 중 특히 8k 및 16k 문맥 길이에서 최첨단(SOTA) 성능을 달-성했습니다.
• 추론 지연 시간: NAtS-L Hybrid는 트랜스포머 대비 프리필(pre-filling) 단계에서 5.4배, 128k 문맥 길이에서의 디코딩(decoding) 단계에서 2.3배의 속도 향상을 보였으며, 순수 GDN 모델보다는 1.66배 느린 수준이었습니다.

의의 및 주장

본 논문은 NAtS-L이 특정 정보 요구 사항에 따라 토큰을 선형 또는 소프트맥스 어텐션에 적응적으로 할당하는 유연한 하이브리드 아키텍처를 제공한다고 주장합니다. 이를 엔드투엔드로 학습함으로써, 모델은 고정된 하이브리드 구조의 경직성과 전체 소프트맥스 어텐션의 비효율성을 피할 수 있습니다.

저자들은 이 접근 방식이 다음과 같은 이점을 제공한다고 강조합니다:

• 필요한 경우 소프트맥스 어텐션을 통해 장기 문맥 관련 정보를 보존합니다.
• 단기적 영향력이 있는 토큰에 대해 선형 어텐션을 사용하여 전반적인 계산 비용을 줄입니다.
• 어텐션 유형의 비율에 대한 별도의 사전 학습된 모델이나 부가적인 제약 없이 더 나은 정확도-효율성 트레이드오프를 달성합니다.

이 연구는 토큰 수준의 선형 및 소프트맥스 어텐션 혼합이 검색 능력을 저해하지 않으면서도 확장 가능한 장기 문맥 언어 모델을 향한 실행 가능한 경로임을 시사합니다.