Taylor-Calibrate: 하이브리드 선형 어텐션 증류를 위한 원칙적 초기화

하이브드 선형 어텐션(GDN)은 long-context 추론에서 KV-cache 비용을 줄이지만, 새로 도입된 recurrent dynamics의 초기화가 성능에 큰 영향을 미친다. 교사 모델의 투영만 복사하는 기존 전이는 교환된 동적 파라미터의 decay, write 게이트, output 게이트를 제대로 설정하지 못하면 초기 제로샷 품질이 저하되고 추가 distillation 토큰이 필요하다. Taylor-Calibrate은 교사 attention 통계를 활용해 decay, write 게이트, value scale, output 게이트를 초기화하고, 레이어 단위 정렬으로 빠른 보정을 수행함으로써 초기 성능을 향상시키고 downstream recovery를 가속한다.

• 전체 설정: 교사 Transformer를 부분적으로 Softmax-attention으로 유지하고 나머지는 GDN으로 대체하는 하이브리드 구성에서, 교사 투영(Q/K/V/O)을 재사용하는 표준 전이(backbone)를 기반으로 한다. - Phase 1: Taylor-derived Calibration - dh( head별 평균 look-back distance)와 eh( entropy) 등을 교사 어텐션에서 계산해 decay(dt_bias)와 write gate(β⋆)를 설정하고, value-path의 스케일(WV)과 output-gate의 초기 값을 결정한다. 또한 g_proj를 작은 게이트로 초기화하고, W_V를 sigma⋆로 스케일한다. - Phase 2: Per-layer Alignment - 각 변환된 GDN 레이어에서 z_S^(ℓ)(x; θℓ)와 z_T^(ℓ)(x)을 calibration 입력으로 비교해 L_init(ℓ)를 최소화하도록 θℓ를 업데이트한다. - 최적화 대상은 W_Q/W_K/W_V/W_O, A_log, dt_bias, a_proj, b_proj, g_proj 등이며 calibration 데이터는 소규모 코퍼스다.

• 메인 벤치마크에서 Taylor-Calibrate은 zero-shot Avg를 일관되게 향상시키고, 토큰 예산이 같은 경우보다 빠른 회복을 보인다. 예를 들어 100M 토큰 시나리오에서 Qwen2.5-1.5B의 Uniform 정책에서 Baseline 대비 Avg가 크게 개선되고, Qwen3-8B에서도 비슷한 추세를 보인다. - Recovery: Stage 1에서 초기 시작점이 더 낮은 손실 상태를 제공하며, Stage 2의 토큰 수를 크게 줄여 교사 수준의 품질에 빠르게 도달한다. 100M 토큰에서의 개선 폭은 모델에 따라 2–3배 이상, 일부 설정에서 4–9×의 토큰 절감으로 나타난다. - Ablation: Phase 1_ANALYTICAL calibration만으로는 충분치 않으며, Phase 2의 layer-local alignment가 필요하다. Alignment만으로도 상당한 개선이 가능하지만, 두 단계의 결합이 최적의 Avg를 제공한다. - Long-context: RULER 지표는 여전히 교사에 도달하기까지 시간이 더 필요하며, Stage 2 이후에도 추가 recovery 학습이 필요하다.