FreeSolo가 학습 길이의 64배까지 제로샷 검색을 유지한다

무엇인가

긴 문맥 언어모델링의 핵심 난점은 학습 길이를 넘어선 검색이다. 전체 어텐션은 시퀀스 길이에 대해 이차 비용을 쓰고, 순환 상태에 이력을 압축하는 선형 시간 모델은 고정 크기 상태 때문에 저장 용량이 묶인다. 회수율을 위해 전체 어텐션을 섞는 하이브리드는 다시 이차 비용을 부른다. 연속 사전학습(CPT)으로 희소 어텐션의 유효 문맥을 늘리는 방법도 있지만, 학습 길이 밖으로의 일반화가 어렵고 CPT 자체가 비싸다. FreeSolo는 이 지점에서 CPT를 짧게만 쓰고도 4k 학습 문맥에서 16k는 물론 256k까지 제로샷 외삽을 노린다.

어떻게 동작하나

구조는 헤드마다 두 갈래다. 단거리 갈래는 고정 창(w=512) 안에서 RoPE를 적용한 슬라이딩 윈도우 어텐션을 수행하고, 장거리 갈래는 셀 단위로 라우팅한 뒤 선택된 셀의 원본 토큰 키와 값에 NoPE 어텐션을 건다. 셀은 s=32개의 연속 토큰 묶음이고, 요약은 세 가지다. 평균 키는 라우팅용 내용 표현, 셀 중앙 위치는 상대 거리 계산용, 셀 내 최대 키 노름은 쿼리와 무관한 우선순위다. 쿼리 블록(B=128) 단위로 프리픽스가 고정되고, 셀 점수는 평균 키 내적에 노름 가중 erf 위치 편향을 더한 r_ic = q_i·k̄_c/√d + λ_h·a_c·erf((i-p_c)/(2σ_h))로 계산한다. 쿼리마다 상위 K=4개 셀을 고르고, 그 셀들의 원본 토큰에 단일 softmax를 적용해 읽는다. 두 갈래 출력은 헤드별 학습 스칼라 α_h로 합쳐진다. 중요한 설계는 RoPE를 국소 창 안에만 가두고 장거리 라우팅·판독은 회전 없는 쿼리와 키를 쓴다는 점으로, 평가 시 학습 때 보지 못한 회전 오프셋이 등장하지 않는다.

무엇과 다른가

변형은 두 가지다. FS-L은 쿼리가 도착하기 전에 프리픽스를 잘라낸다. 각 헤드가 최근 W=8개 셀과 최대 키 노름 상위 P=64개 셀만 남기고, 그 안에서 쿼리별 상위 K개를 고른다. 비용은 O(N)이고 상태 크기가 고정된다. 이 제약에 적응하도록 4096 토큰 문맥에서 짧은 CPT를 돌리는데, 340M에서 0.73B, 800M에서 1.57B 토큰으로 사전학습의 약 5% 수준이며 더 긴 문맥으로 학습하지 않는다. FS-B는 셀을 버리지 않는다. 헤드마다 전체 프리픽스를 담는 포레스트와 중앙값 이상 고노름 셀만 담는 포레스트 두 개를 만들고, 각 노드에 평균 키·반경·최대 노름·최소 위치를 저장한 뒤 빔 서치(레벨당 W_b=16)로 내려간다. 노드 점수 U_v는 하위 셀 점수의 상한 역할을 해 가지치기 근거가 되고, 도달한 후보 셀은 토큰 키로 logsumexp 재랭킹한 뒤 상위 K=4개를 읽는다. radix sort 구성을 포함한 알고리즘 비용은 O(N log²N)로 초선형이지만 이차보다는 낮다. 800M FS-B는 recall mask와 switch mask로 가중한 다음 토큰 손실로 0.42B 토큰 CPT를 추가한다.

어떻게 쓰나

실험은 FineWeb-Edu로 학습한 LLaMA 스타일 디코더 341M·743M, 사전학습 15B·32B 토큰, 4096 토큰 문맥에서 수행했다. 평가는 NVIDIA RULER를 태스크·길이당 100 예제, 그리디 디코딩으로 돌렸다. 800M·16k(학습 길이의 4배)에서 FS-L은 S-NIAH-2/3에서 81/40을 기록해 Raven의 26/0, FS의 72/34와 비교된다. 64k에서는 FS-L이 53/13으로 FS의 28/4, NAtS-L의 0/0을 앞선다. FS-B는 800M·16k에서 S-NIAH-1/2/3에 100/88/84를 내고 RULER 평균 30.4를 기록했다. RULER 평균은 800M·16k에서 FS/FS-L이 23.1/22.3인 반면 Raven/NAtS-L은 12.8/2.8이다. 단일 바늘을 넘어서도 FS-L은 MK-NIAH-1에서 23(두 선형 베이스라인은 각 3), MV-NIAH에서 17(Raven 9, NAtS-L 0.5)을 냈다. 4k 대비 RULER 평균 하락은 800M 기준 FS-L이 32.0에서 22.3, FS-B가 37.7에서 30.4였다. 340M에서 FS는 256k까지 98~100을 유지하고 FS-L은 256k에서 84를 기록한 반면 Raven은 131k에서 20으로 떨어졌다. 800M에서는 FS/FS-L이 256k에서 96/85, NAtS-L은 0이었다. 짧은 문맥 품질도 유지된다. 340M에서 FS 51.9, FS-L 51.0으로 TF++ 51.8, Raven 51.4, NAtS-L 50.8과 비슷하고, 800M에서 FS 56.5는 TF++ 56.8, SWA 56.2에 근접하며 FS-L 55.2, Raven 54.6이다.

전제와 한계

효율은 A100-40GB 한 장에서 4k에서 256k로 문맥을 늘려 측정했다. FS-L의 어텐션 버퍼는 340M에서 766MiB, 800M에서 1149MiB로 일정하게 유지됐고, CUPT 그래프 디코딩은 각각 약 89·75 토큰/초, 프리필 처리량은 약 2.1~2.2K 토큰/초로 입력 길이에 거의 선형으로 늘었다. 다만 마스크 구현은 프리필이 34.6~44.2K 토큰/초로 빠르지만 KV 캐시가 계속 자라고, 340M·4k 학습은 8개 GPU에서 98.4K 토큰/초로 전체 어텐션의 303K 토큰/초보다 느리다. 저자들은 프루닝·수집·희소 판독용 커널 최적화를 향후 과제로 남긴다.

실무 관점에서 이 논문이 주는 신호는 명확하다. 문맥을 늘리기 위해 긴 시퀀스로 재학습하는 대신, 학습 시점의 짧은 문맥을 유지한 채 라우팅 요약(평균 키·키 노름)만 학습시키고 추론 시 검색 예산을 조절하는 설계가 가능하다는 것이다. 특히 RoPE를 국소 창에만 묶고 장거리 판독을 NoPE로 처리하는 방식은 위치 외삽 문제를 구조적으로 회피하는 패턴이라, 자체 롱컨텍스트 모델을 만드는 팀이 참고할 만하다. 도입 전에 확인할 것은 두 가지다. 첫째, 자신의 워크로드가 단일 바늘형 검색인지 아니면 서로 독립적으로 질의 가능한 다수의 사실을 요구하는지. 전자면 FS-L의 고정 예산이 잘 맞지만 후자면 FS-B처럼 메모리를 계속 키우는 쪽이 필요하다. 둘째, 프리필 처리량과 학습 처리량이 아직 전체 어텐션보다 불리하다는 점이므로 지연 시간이 민감한 서비스에서는 실제 커널 성능을 먼저 재봐야 한다.

저자가 밝힌 한계도 분명하다. 실험은 800M 규모에서 멈췄다. FS-L의 고정 유지 예산은 쿼리가 도착하기 전에 증거를 버릴 수 있고, CPT가 즉시 프루닝으로 잃은 S-NIAH-1을 62에서 76~90까지 회복시키긴 하지만 800M FS-L은 평가한 모든 길이에서 기본 FS보다 RULER 평균이 낮다. FS-L의 메모리 예산 안에서 스위트 평균을 끌어올리는 CPT 목적 함수를 찾는 일은 미해결로 남는다. FS-B는 메모리를 키워 RULER 평균을 개선하지만 800M 다중 질의·다중 값 검색에서는 Raven-full에 뒤진다. 고정된 추론 알고리즘 대신 검색 정책 자체를 학습하는 것도 향후 과제다. 부록 수준의 소거 실험에서는 국소 NoPE 대조군이 하류 정확도를 51.9에서 38.3으로 떨어뜨렸고, 장거리 erf 항을 제거하면 S-NIAH-1은 유지되지만 S-NIAH-2는 나빠지고 S-NIAH-3은 좋아지는 비단조적 양상이 나타났다.