FPA는 경로 합 대신 최저점으로 키를 소거한다
First Passage Attention
무엇인가
이 논문은 소프트맥스 어텐션에 망각 게이트를 붙이는 기존 방식이 키와 쿼리 사이 경로의 감쇠를 합산한다는 점을 문제로 잡는다. 이런 합산은 단조적이라 중간 문맥을 지우면서 더 오래된 구조적으로 중요한 토큰을 남길 수 없다. 백스페이스, 닫는 괄호, 변수 재할당처럼 최근 키를 무효화하면서 오래된 키는 살려야 하는 사건에서 이 한계가 드러난다. FoX의 토큰 단위 망각 게이트는 나이에 따라 키를 단조 감쇠시키고, Selective Attention과 DCP의 쌍별 누적은 억제가 중간 단계 수에 비례해 사건의 깊이와 지속 시간을 혼동하며 쌍별 계산 비용도 든다.
어떻게 동작하나
FPA는 경로 합 대신 경로 최솟값을 쓴다. 각 헤드는 모든 위치에 스칼라 필드 φ를, 각 키에 스칼라 레벨 ψ를 두고 편향을 b_ts = β min(0, min_{s<u≤t} φ_u − ψ_s), b_tt = 0으로 정의한다. 키 s는 이후 필드가 자기 레벨 이상이면 살아 있고 b=0이며, 필드가 레벨 아래로 내려가면 죽는다. 죽은 키의 로짓은 그 구간에서 가장 깊은 딥의 깊이에 비례해 낮아지고, 딥이 얼마나 오래 지속됐는지는 보지 않는다. 구현은 헤드별 선형 사상으로 증분 δ와 오프셋 g, h를 만들고, 걷기 W_t = Σδ_u에 대해 φ_t = W_t + g_t, ψ_t = W_t + h_t로 필드와 레벨을 정한다. β는 헤드별 양의 학습 스칼라다. 이 선형 매개변수화는 FoX, 병목 편향, 학습된 만료를 특수 사례로 포함한다. g=h=0인 스택 형태에서는 살아 있는 키가 스택처럼 변해, 나중 토큰이 더 새로운 키를 지우면서 오래된 키를 남길 수 있다. 학습은 순전파에서 범위 최솟값을 쓰고 역전파에서는 τ=1인 소프트 최솟값 −τ log Σ exp(−φ_u/τ)의 기울기를 쓰는 직통 추정기로 한다. 블록 크기 B에서 편향은 블록 쌍마다 O(B^2) 연산과 길이 T에 대해 O(T/B) 블록 최솟값을 쓰고, 자기회귀 디코딩에서는 캐시된 키·헤드마다 두 스칼라를 저장하며 각 단계에서 O(t)로 편향 행을 갱신한다.
무엇과 다른가
이론에서는 사다리 가족을 통해 한 개의 상수 차원 FPA 헤드가 이후 딥이 만드는 모든 살아 있는 키 집합을 실현한다고 주장한다. 반면 위치 편향이나 토큰 단위 망각 게이트를 가진 소프트맥스 헤드는 같은 집합을 실현하려면 헤드 차원이 r−2 이상 필요하거나, 키가 쓰일 때 고정되는 경우 d·p ≥ r−2가 필요하다. FPA 헤드는 O(log r) 비트 스칼라와 O(log r) 크기의 상수로 충분하다. 또한 범위 극값만이 지속 시간에 불변한 억제를 주며, Selective Attention처럼 딥을 합산하면 지속 시간이 달라질 때 목표를 정확히 맞추지 못한다는 명제가 제시된다. 헤드별 필드가 필요한 이유도 다룬다. Selective Attention과 DCP처럼 층 공유 마스크를 쓰면 차원 요구가 최대 절반으로 줄 뿐 r 의존성은 사라지지 않는다. 다만 이 구분은 한 층에 대한 것이고, Selective Attention과 DCP도 하드 극한에서는 같은 살아 있는 집합을 실현하므로 FPA가 그 극한에서 더 표현력이 높다는 주장은 아니다. FPA의 차이는 편향이 단일 위치 스칼라에만 의존하고, 지속 시간이 아니라 깊이에 의존한다는 점이다.
어떻게 쓰나
1층 프로브는 15개 시드로 진행됐다. d=4에서 RoPE, NoPE, ALiBi, 토큰 단위 FoX 헤드는 r=6에서 14 또는 15개 시드로 성공했지만 r=10, 14, 18에서는 모든 시드에서 실패했다. FPA 헤드는 r=10, 14, 18에서 각각 11, 12, 13개 시드로 성공했고 열두 비교 모두 p=0.00049로 기각됐다. 같은 헤드를 d=r+2로 키우면 모든 시드에서 성공했으므로 실패가 학습 문제는 아니었다. 딥의 가장 깊은 단계가 1 또는 2스텝으로 유지되는 가변 지속 시간 과제에서는 FPA가 13개 시드로 성공했고, Selective Attention과 최솟값 대신 딥 합을 쓴 FPA 변형은 모두 0개 시드로 실패했다(p=0.00012). 고정 지속 시간에서는 두 경쟁자도 15개 시드 모두 맞췄다. 두 걷기 가족의 상보 부분에서 차원 2의 두 헤드를 쓰면 FPA는 11개 시드로 성공했고, 층 공유 마스크를 쓰는 Selective Attention과 DCP는 각각 0개 시드로 실패했다(p=0.00049). 푸시·팝 과제에서는 FPA가 14개 시드, 딥 합 FPA가 11개 시드로 성공했고 1층 FoX, 소프트맥스, Selective Attention은 모두 실패했다. 살아 있는 키 중 이름의 출현을 찾는 과제에서는 FPA가 15개 시드 모두 성공했고, 키 레벨을 빼거나 쿼리 레벨을 쓰거나 내용 검색을 빼면 모두 실패했으며 Selective Attention은 13개 시드로 성공했다. 이전 연구 과제에서는 flip-flop 언어 모델링, 복사, 카운팅에서 FPA가 모든 시드로 성공했고 변수 할당에서는 9/15로 FoX의 전 시드 성공에 못 미쳤다. 다중 쿼리 연관 회상에서는 FPA가 모든 시드에서 실패했고 소프트맥스는 모두 성공했다. 저자 진단은 1층 편향이 학습 초기에 각 값에서 자기 키로 가는 어텐션을 억제해 회상 회로의 첫 단계가 형성되지 않았다는 것이며, 1층에 편향을 넣지 않은 FPA는 15개 시드 모두 성공했고 탐색적 확인에서 3·4층 FPA도 2층처럼 실패했다.
전제와 한계
언어 모델 실험은 12층, 폭 768, 12헤드, 헤드 차원 64의 디코더 넷을 SlimPajama 2.7B 토큰, Mistral 7B의 32k 토큰 어휘, 문맥 2048, 단일 시드, 혼합 정밀도로 학습했다. FPA와 FoX는 위치 인코딩 없이, Selective Attention과 소프트맥스는 RoPE를 썼다. SlimPajama 테스트에서 FPA 손실은 FoX보다 0.0192 nats 낮았고 차이의 상한 신뢰구간 −0.0186이 0.00995 마진 안에 들어왔다. PG-19 99권의 2048~4095 위치에서는 FPA 손실이 0.034 nats 낮았고 상한은 −0.028이었다. 패스키 검색은 길이 3072와 4096에서 FPA가 100개를 모두 찾았고 FoX는 73개였다(p=7.5×10^-9). 16,384 토큰 창에서 FPA와 FoX는 학습 길이를 넘어서도 손실이 평평했고, FPA는 8192까지 FoX보다 0.02~0.03 nats 낮았지만 8192~16,384에서는 0.026 nats 높았다(95% 구간 [0.019, 0.032]). 그래도 8192와 16,384에서 FPA는 모든 패스키를 찾았고 FoX는 20%를 찾았다. RULER에서는 8192 길이에서 FPA가 0.202로 FoX의 0.117보다 높았지만 4096과 16,384에서는 약간 낮았고, BABILong에서는 어떤 모델도 0.09에 도달하지 못했다. 여섯 다운스트림 제로샷 정확도는 FoX와 가까웠고 대부분 차이의 95% 구간이 0을 포함했다. 관찰 분석에서는 1024 이후 쿼리에서 중앙 헤드가 편향 −1 이상인 키의 0.3%만 남겼고, 144개 헤드 중 19개는 살아 있는 집합의 10% 이상이 접미사가 아니었다. 문서 경계에서의 죽음 비율은 토큰 비율의 6.6배였고, 56개 헤드에서 걷기가 여는 괄호와 닫는 괄호에서 반대 방향으로 움직였으며 괄호 깊이와의 상관은 이동 대조군의 99% 밴드를 벗어났다. 다만 최대 상관이 0.25로 작아 저자들은 괄호 추적의 증거가 아니라 설명으로 읽는다.
비용 측정에서 FPA 한 층은 순전파·역전파 기준 SDPA의 0.97배, FoX의 1.50배 시간이 걸렸다. 다섯 가지 형상에서 그 비율은 각각 0.95~1.08배, 1.42~1.56배였고 Selective Attention은 FPA의 2.4배, 형상별로 2.0~2.5배였다. FPA 순전파만은 FoX의 1.03~1.07배였고, 한 커널 안에서 범위 최솟값이 FoX 편향 시간의 14%, 필드 기울기가 26%를 더했다. 전체 FPA 모델은 FoX 처리량의 0.914배, 소프트맥스의 1.103배, Selective Attention의 1.47배로 학습됐다. 배치 1, 길이 8192 디코딩은 토큰당 22.0ms로 FoX 20.2ms, 소프트맥스 23.9ms와 비교됐고 1024에서 16,384까지 시간이 늘지 않아 고정 오버헤드가 지배했다. FPA 상태는 키·값 캐시에 1.56%를 더했다. 편향이 다시 오르지 않으므로 −10 아래로 떨어진 키를 버릴 수 있고, 이때 2048에서 손실이 0.00044 nats, 8192에서 0.00047 nats 올랐으며 키의 34%와 26%를 유지했다. 캐시의 24~27%만으로도 8192와 16,384에서 모든 패스키를 찾았다. 같은 규칙은 FoX에서 0.0195~0.0227 nats 손실을 내며 키의 15%와 7%를 유지했고, 같은 평균 크기의 슬라이딩 윈도는 2048에서 소프트맥스에 0.115 nats 손실을 냈다.
개발자에게 FPA는 긴 문맥 추론에서 KV 캐시를 줄이면서 오래된 키를 선택적으로 보존하는 어텐션 커널 후보로 의미가 있다. 편향이 단조 증가하지 않아 임계값 아래 키를 버리는 정책과 잘 맞고, 논문의 Triton 커널은 학습에서 SDPA와 비슷한 속도를 보였다. 다만 첫 층 편향이 다중 쿼리 연관 회상 같은 과제를 망칠 수 있고, 110M 파라미터·단일 시드·2.7B 토큰·12GB Turing GPU 한 파티션에서의 비교라 규모가 커지면 결론이 달라질 수 있다. 이론도 주로 한 층에 대한 것이며, 더 깊은 모델은 다루지 않고 정밀도 한계도 두지 않는다. 가변 지속 시간 명제는 정확한 적합에 대한 것이고 허용 오차 내 도달 범위를 제한하지 않는다. 실무 적용 전에는 대상 과제에서 첫 층 편향 유무, 캐시 축소 임계값, 긴 문맥 검색 성능을 직접 검증해야 한다.