복사 전용 헤드를 붙여 언어모델의 학습 효율과 장문맥 검색을 개선한다

무엇인가

언어에는 어휘적 반복이 넘친다. 문서는 같은 개체와 사건을 다시 언급하고, 대화는 앞선 턴을 되풀이하며, 코드는 같은 변수와 함수를 참조한다. 실제로 인기 사전학습 코퍼스인 FineWeb-Edu의 토큰 중 약 60%가 이미 자신의 4k 토큰 문맥 안에 앞서 등장한 토큰이다. 그래서 문맥으로부터의 복사는 언어모델이 수행하는 가장 기본적인 연산 중 하나이며, 트랜스포머는 induction head나 retrieval head 같은 복사 전용 어텐션 헤드를 다수 배치한다. 그런데도 표준 신경망 언어모델 셋업은 복사를 놀랄 만큼 배우기 어렵게 만든다. 모든 예측이 어휘 크기 V보다 훨씬 좁은 D차원 선형 투영인 LM 헤드를 통과하기 때문이다. 이 저랭크 층은 모델이 표현할 수 있는 분포를 제한하고(softmax bottleneck), 역전파 시 V차원 로짓 그래디언트를 D차원으로 파괴적으로 압축해 백본에 도달하기 전에 대부분의 학습 신호를 버린다. 그 결과 바닐라 트랜스포머는 어휘가 큰 조건에서 사소한 합성 복사 과제조차 실패한다.

어떻게 동작하나

논문이 제안하는 Copy-Paste Language Model(CPLM)은 기존 LM 헤드 옆에 가벼운 Copy-Paste(CP) 헤드를 하나 더 붙인다. CP 헤드의 예측 대상은 어휘가 아니라 문맥 위치 τ ∈ {1,…,t}다. 현재 상태로 만든 쿼리 q_t = W_Q h_t와 각 과거 위치의 키 k_τ = W_K h_τ의 내적으로 어텐션처럼 위치별 점수를 매기고, α_t(τ) = softmax(k_τ^T q_t / √D)를 계산한다. 그다음 같은 토큰을 내는 위치들에 대해 이 확률을 주변화해 토큰 분포 p_CP(v) = Σ_{τ≤t, x_τ=v} α_t(τ)를 얻는다. 즉 과거 위치에 대한 분포를 만든 뒤 해당 위치의 토큰 분포를 끌어와 예측하는 방식이고, 이 분포의 지지집합은 문맥에 등장한 토큰으로만 제한된다. 쿼리·키 차원은 기본적으로 은닉 차원 D와 같게 두지만, 훨씬 작은 차원(64)도 비슷한 성능을 낸다.

무엇과 다른가

복사 경로를 여닫는 게이트가 이 설계의 핵심이다. LM 헤드에 실제 타깃이 never 되는 더미 토큰 <copy>를 추가하고, LM 헤드가 이 토큰에 부여하는 확률 c_t를 복사 메커니즘의 게이트로 쓴다. 또 CP 헤드의 키 집합에 학습 가능한 <sink> 키를 넣어, 문맥에 복사할 만한 위치가 없을 때 LM 헤드 예측으로 되돌아갈 수 있게 한다. 최종 다음 토큰 분포는 p(v) = (1−λ_t) p_LM(v) + λ_t p_CP(v), λ_t = c_t(1−s_t)로, 두 분포의 볼록 결합이므로 정의상 유효한 분포다. c_t가 복사 경로를 열고 s_t가 다시 닫는다. 학습은 이 혼합 분포에 대한 표준 다음 토큰 교차엔트로피만으로 end-to-end로 이뤄지며, 보조 손실도 복사 지도(supervision)도 없다. 모델이 언제 어디를 복사할지 스스로 배운다. 게이트를 포인터 소프트맥스 내부의 센티널이 아니라 LM 헤드가 어휘와 함께 예측한다는 점이 기존 pointer sentinel mixture와의 차이다.

어떻게 쓰나

실험은 135M에서 1.7B 파라미터까지, FineWeb-Edu 데이터로 최대 200B 토큰을 사전학습해 검증했다. 백본은 full attention, sliding-window(SW), 둘을 교차한 hybrid, Gated DeltaNet을 모두 포함한다. 결과는 모든 규모와 백본에서 일관되게 CPLM 쪽으로 기울었다. 수렴 속도는 full·hybrid에서 바닐라 대비 1.25배, SW에서는 2배 빨랐고, 1.7B SW CPLM은 4k 토큰 기준 12배 적은 메모리로 full-attention LM과 동등한 성능을 냈다(윈도 밖 토큰당 복사 키 하나만 캐시하면 되기 때문). 분포 외 문서인 FinePDFs에서 퍼플렉서티가 0.5~1.5포인트 개선됐고, HellaSwag·LAMBADA·PIQA·SciQ·ARC·WinoGrande 제로샷 평가에서는 정확도가 0.4~0.9포인트 올랐으며 특히 LAMBADA에서 +2.5~4.0포인트로 가장 컸다. 이 이득은 복사된 토큰에만 국한되지 않고 <copy> 확률이 거의 0인 토큰에서도 NLL이 감소했다.

전제와 한계

지도 미세조정과 장문맥, 검색 실험도 뒷받침한다. Tulu-3로 1.7B 모델을 SFT했을 때 CPLM은 평균 점수에서 1.2~5.1포인트 앞섰고, 수학 과제에서는 full attention +3.9, sliding window +7.9포인트였다. LongBench에서 평균 +4.1포인트, 합성 장문맥 과제에서 +7.5포인트를 기록했다. 7자리 숫자를 중립 텍스트 속 8개 깊이에 심는 단일 니들 검색 실험(문맥 1k~64k, 사전학습 문맥의 최대 16배)에서 full attention은 49.7→54.8, hybrid 3:1은 58.2→83.6, SW는 14.3→99.4로 개선됐다. 윈도가 256에 불과한 SW CPLM이 거의 완벽하게 니들을 찾았다는 뜻이다. HotpotQA distractor 설정에서는 문서를 백본에서 완전히 숨기고 CP 키로만 접근하게 해도 5개 distractor에서 정확도의 90~97%를 유지했고, 500개에서도 34~38% EM으로 문서를 KV 캐시에 담은 모델들과 대등하거나 앞서면서 캐시는 50배 작았다.

저자들은 이 성공의 원인을 LM 헤드 그래디언트 병목 완화로 설명한다. 135M 모델 분석에서 LM 헤드는 그래디언트 노름의 99% 이상을 잃는 반면, CP 헤드는 로짓이 어휘 V가 아니라 문맥 위치 t개 위에 정의되므로(t ≪ V) 초기 위치에서 최대 85%의 그래디언트 노름을 보존한다. 어휘 크기 8k~512k의 SpamLang 합성 과제에서 바닐라 모델은 512k 어휘에서 34% 정확도에 그쳤지만 CPLM은 모든 어휘 크기에서 과제를 풀었다. CP 헤드에서 백본으로 가는 그래디언트를 차단하는 stop-gradient 어블레이션을 하면 이득 대부분이 사라졌고, softmax 병목이 작동하지 않는 바이트 수준 모델(D > V = 257)에서는 CPLM이 바닐라보다 나아지지 않았다. 또한 CPLM 백본은 T=128에서 induction head를 4개만 갖는 반면 바닐라 모델은 13~15개를 가졌다. CP 헤드가 여러 헤드에 흩어져 있던 복사 작업을 대신 떠맡는다는 해석이다.

실무 관점에서 CPLM은 두 곳에 바로 쓸 수 있다. 하나는 sliding-window나 hybrid 백본을 쓰면서 메모리를 줄이려는 경우다. CP 헤드는 1.7B 기준 파라미터 0.5%, 연산 3%만 추가하고 9B·30B에서는 학습 처리량의 1% 미만을 쓰면서, 윈도 밖 정보를 복사 키로 끌어온다. 다른 하나는 RAG다. 문서를 백본에 넣지 않고 CP 키만으로 인코딩하면 문서별 키를 한 번 계산해 여러 질의에 재사용할 수 있어, full KV 캐시보다 훨씬 작은 메모리로 검색 정확도를 유지한다. 저자들은 135M~1.7B 체크포인트를 HuggingFace에 공개했고, NanoGPT 스피드런 기록도 11.3% 개선했다고 밝혔다. 다만 도입 전에 확인할 것은 어휘 크기가 은닉 차원보다 훨씬 큰, 즉 병목이 실제로 존재하는 셋업인지다. 바이트 수준처럼 V가 작은 경우에는 이득이 없다고 논문이 직접 보고한다.

한계도 분명하다. CP 헤드는 문맥에 등장한 토큰에 대해서만 LM 헤드 병목을 완화하며, 어휘 전체에 대한 표현력 문제를 해결하지는 않는다. 실험 규모는 1.7B 파라미터에서 멈췄고, 저자들은 복사-붙여넣기 메커니즘을 토큰이 아닌 스팬이나 외부 코퍼스로 확장하는 것, 그리고 더 큰 규모에서 이득이 어떻게 변하는지가 자연스러운 다음 단계라고 밝힌다.