Taylor-Calibrate: Inicialização Fundamentada para Destilação de Atenção Linear Híbrida

Resumo

Modelos de atenção linear híbridos oferecem um caminho promissor para inferência mais rápida em contextos longos: reduzem o custo quadrático e a carga do cache KV da atenção softmax completa, mantendo grande parte da qualidade dos modelos Transformer. Uma forma prática de obter tais modelos é converter um Transformer pré-treinado em vez de treinar uma nova arquitetura do zero, mas essa conversão ainda é frágil. Simplesmente copiar as projeções de atenção do professor para um aluno Gated DeltaNet (GDN) não especifica a nova dinâmica de decaimento recorrente, escrita e controle de saída. Como resultado, o modelo convertido frequentemente começa em um regime dinâmico ruim e precisa gastar muitos tokens de destilação reparando a inicialização, em vez de aprender o comportamento restante do professor. Propomos o Taylor-Calibrate, um método de inicialização leve para alunos GDN híbridos. O método utiliza estatísticas de atenção do professor guiadas por Taylor para definir a projeção de valor, a escala de tempo da memória, os controles de escrita e o controle de saída; em seguida, aplica uma etapa curta de alinhamento por camada para ajustar cada camada convertida à saída do professor. Em quatro configurações de professor e três políticas de camadas retidas, o Taylor-Calibrate produz alunos zero-shot substancialmente mais fortes, com uma melhoria de até 88x em uma ablação representativa, e atinge alvos de recuperação correspondentes com 4,9x a 9,2x menos tokens de treinamento do que a conversão ingênua.

English

Hybrid linear attention models offer an appealing path to faster long-context inference: they reduce the quadratic cost and KV-cache burden of full softmax attention while retaining much of the quality of Transformer models. A practical way to obtain such models is to convert a pretrained Transformer instead of pretraining a new architecture from scratch, but this conversion is still brittle. Simply copying the teacher attention projections into a Gated DeltaNet (GDN) student does not specify the new recurrent decay, write, and output-gating dynamics. As a result, the converted model often starts in a poor dynamical regime and must spend many distillation tokens repairing initialization rather than learning the remaining teacher behavior. We propose Taylor-Calibrate, a lightweight initialization method for hybrid GDN students. The method uses Taylor-guided teacher attention statistics to set the value projection, memory timescale, write gates, and output gate, then applies a short per-layer alignment step to match each converted layer to the teacher output. Across four teacher settings and three retained-layer policies, Taylor-Calibrate gives substantially stronger zero-shot students, with up to an 88x improvement in a representative ablation, and reaches matched recovery targets with 4.9x--9.2x fewer training tokens than naive conversion.