하이브리드 모델에서 순환층이 검색 학습을 독점하는 경쟁을 BLD로 조정한다

무엇인가

이 논문은 선형 시간 순환과 소프트맥스 어텐션을 함께 쓰는 하이브리드 언어모델에서 층별 역할 분담이 어떻게 형성되는지를 다룬다. 순환층은 문맥을 고정 크기 상태로 압축해 처리량이 높고 표현력 있는 혼합을 제공하지만, 정확한 문맥 내 검색에서는 소프트맥스 어텐션보다 약하다. 그래서 하이브리드 모델은 보통 어텐션에 장거리 검색을, 순환층에 단거리 혼합을 맡기는 분업을 기대한다. 그러나 저자들은 사전학습 중 두 종류의 믹서가 같은 검색 기능을 두고 경쟁하는 Competing Learning 현상을 발견한다. 어떤 층이 검색을 먼저 배우면 검색 특화 예측 오차가 줄어들어 다른 층에 전달되는 학습 신호가 약해지고, 그 결과 기대한 분업이 만들어지지 않을 수 있다.

어떻게 동작하나

분석은 합성 검색 과제인 Spaced Retrieval에서 시작한다. 입력은 $ X1...Xt 뒤에 g개의 공백을 두고 $를 붙이고, 출력은 X1...Xt를 복원하는 형식이며, 각 Xi는 0부터 9까지 균등 샘플링한다. 실험은 t=5, g=64와 8192를 사용하고, 3.6M 파라미터의 Kimi Delta Attention(KDA) 2층과 어텐션 2층 하이브리드 모델을 g=64와 8192를 1:1로 섞어 학습한다. 손실은 출력 토큰에만 준다. 저자들은 층별 검색 기여도를 재는 Layer Retrieval Ability Score도 제안한다. 같은 길이의 두 입력 A와 B가 검색할 원본 숫자만 다를 때, 마지막 비마스킹 위치의 믹서 출력만 B의 것으로 바꿔 넣고 정답 선호 마진의 감소량을 계산한다. 이 값이 큰 층일수록 모델의 검색에 더 크게 기여한다.

무엇과 다른가

실험에서 개입 없이 학습하면 KDA 층이 검색 능력을 먼저 장악한다. 여러 학습 단계에서 KDA 전체를 동결하면 어텐션 층의 검색 능력이 즉시 급상승한다. 나중에 KDA를 다시 풀었을 때, 125스텝보다 일찍 동결한 실행은 어텐션이 충분한 우위를 쌓아 해제 후에도 검색 능력이 더 향상되지만, 늦게 동결한 실행은 어텐션이 다시 KDA에 밀려 검색 능력이 떨어진다. 그래디언트 분석에서도 동결 체크포인트의 어텐션 그래디언트 노름이 기준선의 약 10배에 이르고, KDA를 풀면 이 증폭이 완화된다. 그래디언트는 갑자기가 아니라 점진적으로 커져, 현재 파라미터 상태가 학습 신호 배분을 정하고 우세한 구성요소의 이점을 스스로 강화하는 구조임을 시사한다. 또한 장거리 예제 비율이 높아지면 어텐션이 중요해지지만, 장거리 패턴이 압도적이지 않으면 KDA가 우세해진다. KDA가 포겟 게이트를 낮은 값으로 초기화해 단거리 귀납 편향을 갖기 때문이다.

어떻게 쓰나

이에 저자들은 Better Layer Division(BLD)을 제안한다. BLD는 아키텍처를 바꾸지 않고, 사전학습 초기에 RNN 계열 층의 가중치 업데이트를 멈춰 어텐션 층이 검색 학습에서 선발 우위를 갖게 한다. 이후 동결을 풀고 표준 사전학습을 계속하면, 초기에 만들어진 분업이 유지되고 강화된다는 것이 핵심이다. 자연어 실험에서 KDA 블록은 이후 동결 해제 사전학습 구간의 약 1/10 동안 동결된다. 모델은 임베딩 포함 180M에서 1B 총 파라미터 규모이고, 어텐션 1층 뒤에 KDA 3층이 반복되는 Attention-KDA-KDA-KDA 패턴으로 KDA 대 어텐션 비율을 3:1로 고정한다. 모든 층은 NoPE를 쓰고 깊이와 은닉 폭을 바꿔 규모를 확장한다. 학습 데이터는 DCLM, 토크나이저는 Qwen3, 문맥 길이는 16,384이며, Chinchilla 원칙에 따라 대략 파라미터의 20배 토큰을 사용한다. 최적화는 초기 AdamW, 동결 해제 후 Muon과 보조 AdamW를 쓰고 BLD와 기준선의 스케줄은 동일하다.

전제와 한계

평가는 파인튜닝 없이 진행한다. BLD 모델은 언어모델링 손실이 눈에 띄게 낮고, 상식 추론에서는 완만한 향상을, 검색 중심 과제에서는 큰 향상을 보인다. 상식 추론은 ARC-Easy, ARC-Challenge, HellaSwag, PIQA, BoolQ, WinoGrande, LAMBADA 평균으로 측정한다. 검색은 BASED 평가군의 SWDE, FDA, SQuAD, TriviaQA, DROP, Natural Questions 여섯 과제로 측정한다. RULER는 1K에서 16K 토큰에 걸친 57개 과제-길이 조합 평균으로 평가한다. 언어모델링은 WikiText-2 단어 수준 퍼플렉시티와 LongPPL 선택 토큰 퍼플렉시티로 본다. BLD는 Recall과 RULER에서 상당한 이득을, LongPPL 같은 장문맥 지표에서도 큰 우위를 보인다. 다만 본문 서술에는 Table 1의 구체적인 수치가 제시되지 않아 개선폭의 정확한 값은 확인할 수 없다. 문맥 활용 실험에서는 Books3와 BookCorpus2에서 16K까지 토큰 위치별 다음 토큰 손실을 재는데, BLD가 특히 뒤쪽 위치에서 더 낮은 손실을 보인다. 규모 확장 곡선도 BLD의 손실 우위가 평가 규모 밖에서도 유지될 수 있음을 시사한다.

전이성도 확인한다. 같은 구성에서 KDA를 Gated DeltaNet(GDN)이나 Gated Linear Attention(GLA)으로 바꿔도 BLD의 낮은 학습 손실과 다운스트림 성능 이점이 유지된다. 또한 510M KDA-어텐션 하이브리드에서 어텐션 헤드별 alpha 값을 1,000스텝 학습해 장거리 검색에 의존하는 retrieval head와 단거리 streaming head를 구분한다. BLD 모델의 어텐션 헤드는 기준선보다 높은 alpha 값을 보여 장거리 검색 헤드가 더 많고, 낮은 alpha 헤드를 마스킹해도 정확도가 크게 유지되다가 가장 높은 alpha의 마지막 세 헤드를 제거할 때 급락한다. 기준선은 가장 높은 alpha 헤드 하나를 제거할 때만 큰 하락을 보여 검색 헤드가 하나뿐인 것으로 해석된다.

개발자 관점에서 이 논문은 하이브리드 LLM을 사전학습하거나 층별 역할을 설계할 때 손실 하나만 보면 안 된다는 점을 보여준다. 순환층과 어텐션의 검색 기여도, retrieval head의 alpha 분포, 장문맥 위치별 손실, RULER와 Recall 계열 성능을 함께 확인해야 한다. BLD는 아키텍처 변경 없이 초기 동결 스케줄만 추가하는 방법이라 실무에서 시도하기 쉽지만, 동결 범위와 기간, 동결 해제 시점, 옵티마이저 전환은 데이터와 모델 규모에 맞춰 조정해야 한다. 저자들도 BLD가 경쟁 학습을 다루는 예비적 접근이며 더 깊은 이해가 필요하다고 밝힌다. 실험은 최대 1B 파라미터, DCLM, 16K 문맥, KDA 중심이며 GDN과 GLA로의 전이를 확인했지만, 더 큰 규모와 다른 데이터 분포에서의 일반화는 추가 검증이 필요하다.