FocuSFT: Otimização Bilevel para Ajuste Fino de Contexto Longo Consciente de Diluição

Resumo

Modelos de linguagem de grande porte agora conseguem processar entradas cada vez mais longas, porém sua capacidade de utilizar efetivamente informações dispersas em contextos extensos ainda é limitada. Atribuímos essa lacuna à forma como o orçamento de atenção é gasto durante o ajuste fino supervisionado (SFT) em sequências longas: vieses posicionais e sumidouros de atenção fazem com que o modelo aloque a maior parte de sua atenção a tokens posicionalmente privilegiados, em vez de conteúdo semanticamente relevante. Essa diluição da atenção durante o treinamento (a inanição de tokens de conteúdo na distribuição de atenção) enfraquece o sinal de gradiente, limitando a capacidade do modelo de aprender capacidades robustas para contextos longos. Apresentamos o FocuSFT, uma estrutura de otimização bilevel que aborda esse problema no momento do treinamento. Um laço interno adapta parâmetros leves e rápidos ao contexto de treinamento, formando uma memória paramétrica que concentra a atenção no conteúdo relevante; o laço externo realiza o SFT condicionado a essa representação refinada. Ambos os laços aplicam atenção bidirecional sobre os tokens de contexto, preservando o mascaramento causal para as respostas, reduzindo a assimetria causal que origina sumidouros de atenção e alinhando o comportamento interno e externo. No BABILong, o FocuSFT melhora a precisão em até +14 pontos percentuais para comprimentos de contexto de 4K a 32K; no RULER, ele eleva a agregação CWE de 72,9% para 81,1% em 16K; e no GPQA com uso agêntico de ferramentas, obtém um ganho relativo de 24% no pass@1. A análise de atenção mostra que o FocuSFT reduz a massa de sumidouros de atenção em 529 vezes e triplica o engajamento de contexto durante o treinamento. Código: https://github.com/JarvisPei/FocuSFT

English

Large language models can now process increasingly long inputs, yet their ability to effectively use information spread across long contexts remains limited. We trace this gap to how attention budget is spent during supervised fine-tuning (SFT) on long sequences: positional biases and attention sinks cause the model to allocate most of its attention to positionally privileged tokens rather than semantically relevant content. This training-time attention dilution (the starvation of content tokens in the attention distribution) weakens the gradient signal, limiting the model's ability to learn robust long-context capabilities. We introduce FocuSFT, a bilevel optimization framework that addresses this problem at training time. An inner loop adapts lightweight fast-weight parameters on the training context to form a parametric memory that concentrates attention on relevant content, and the outer loop performs SFT conditioned on this sharpened representation. Both loops apply bidirectional attention over context tokens while preserving causal masking for responses, reducing the causal asymmetry that gives rise to attention sinks and aligning inner-outer behavior. On BABILong, FocuSFT improves accuracy by up to +14pp across 4K--32K context lengths; on RULER, it raises CWE aggregation from 72.9\% to 81.1\% at 16K; and on GPQA with agentic tool use, it yields a 24\% relative gain in pass@1. Attention analysis shows that FocuSFT reduces attention sink mass by 529times and triples context engagement during training. Code: https://github.com/JarvisPei/FocuSFT