PaperlayerView exact revision on arXiv

Source-linked AI summary

FocuSFT: Bilevel Optimization for Dilution-Aware Long-Context Fine-Tuning

Zehua Pei, Hui-Ling Zhen, Xianzhi Yu, Sinno Jialin Pan, Mingxuan Yuan, Bei Yu

arXiv:2605.09932v1cs.CL

TL;DR

Long-context models can possess large context windows yet underuse relevant information because training-time positional biases and attention sinks dilute the learning signal. FOCUSFT uses bilevel fine-tuning with fast-weight parametric memory and bidirectional context attention, improving performance across long-context benchmarks while increasing context engagement. Its scope is currently limited by training overhead, one model family, one fixed corpus, and evaluation only at the SFT stage.

  • Problem

    Positional biases and attention sinks starve relevant content tokens during SFT, weakening the gradient signal for long-context learning.

  • Method

    FOCUSFT adapts lightweight fast weights on context in an inner loop, then performs outer-loop SFT on sharpened representations with aligned bidirectional context attention.

  • Results

    +14pp on BABILong at 4K–32K, +8.2pp on RULER CWE aggregation at 16K, and +3.8pp pass@1 on GPQA agentic reasoning.

  • Takeaways & Limitations

    FOCUSFT reduces training-time attention dilution while improving long-context benchmark performance and context engagement.

  • Takeaways & Limitations

    FOCUSFT adds 1.71× training-time overhead, and its scaling beyond Qwen2.5-7B, the 3K-sample corpus, and SFT remains untested.

Abstract

from arXiv · show

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 529$\times$ and triples context engagement during training. Code: https://github.com/JarvisPei/FocuSFT

1 Introduction

Long-context capacity does not ensure reliable context utilization because positional biases and attention sinks dilute training-time attention away from relevant content. FOCUSFT addresses this bottleneck with bilevel fine-tuning and reports improvements across long-context benchmarks.

  • Motivation: Long-context models often neglect relevant information despite large advertised context windows.RULER and Lost in the Middle expose failures as task complexity rises and content appears in the middle of inputs.
  • Motivation: Training-time attention dilution starves content tokens and weakens the gradient signal used for long-context learning.Standard SFT representations are dominated by positionally privileged tokens, creating a cycle that can reinforce poor utilization.
  • Approach: FOCUSFT uses an inner fast-weight adaptation loop to build parametric memory, followed by outer-loop SFT conditioned on sharpened representations.Both loops use bidirectional context attention while preserving causal masking for responses, with inner-outer consistency aligning their behavior.
  • Results: +14pp on BABILong at 4K–32K, +8.2pp on RULER CWE aggregation at 16K, and +3.8pp pass@1 on GPQA agentic reasoning.The reported gains span retrieval, aggregation, and agentic reasoning evaluations.
  • Results: 3.1× higher context engagement accompanies FOCUSFT’s benchmark improvements.Training-time attention analysis reports increased attention directed toward context content.

2 Preliminaries and Motivation

Positional bias and attention sinks divert attention from relevant long-context content, and standard SFT reproduces these patterns during training. FOCUSFT counters the resulting dilution through bidirectional context attention and a bilevel training procedure.

  • Attention mechanisms: Scaled dot-product attention converts query-key compatibility scores into weighted value representations.The mechanism uses query, key, and value projections followed by softmax-normalized attention weights.
  • Failure modes: Causal masking restricts autoregressive attention to positions j ≤ i, making early tokens uniquely visible to later tokens.This visibility pattern supports the formation of attention sinks at initial positions.
  • Failure modes: Positional bias favors context beginnings and ends, leaving middle information systematically underattended.The resulting U-shaped attention profile can make relevant middle content effectively invisible.
  • Failure modes: Attention dilution is the starvation of semantically relevant content caused by positional bias and learned sinks consuming the attention budget.The model allocates most attention to positionally privileged tokens rather than relevant content.
  • Training-time bottleneck: Longer sequences can worsen dilution because sinks absorb more budget while additional distractors compete for the remainder.Simply increasing training context length may therefore yield diminishing gains at the lengths trained on.
  • Training-time remedy: Bidirectional context attention reduces the asymmetric visibility that drives sinks, while preserving causal masking for responses.This design targets the mask structure underlying the sink mechanism.

3 Methodology

FOCUSFT uses bilevel optimization to adapt temporary fast weights that sharpen context representations before standard SFT. Inner–outer consistency and bidirectional context attention align the adaptation and training passes while preserving causal response masking.

  • Bilevel Optimization: FOCUSFT decomposes each training step into an inner fast-weight adaptation loop and an outer SFT loop conditioned on the adapted parameters.The fast weights are re-initialized for each training step, while the outer loop updates the base model.
  • Inner Loop: Parametric Memory: The inner loop adapts LoRA parameters on response-token prediction, causing them to encode context information useful for generating accurate responses.After K gradient steps, the adapted parameters modify intermediate representations and reshape later attention toward salient sample content.
  • Outer Loop: SFT: The outer loop performs autoregressive cross-entropy training on response positions using the combined base and adapted parameters, while discarding fast weights afterward.Only the base parameters receive gradient updates, with the adapted weights serving as temporary attention-sharpening parameters.
  • Bidirectional Context Attention: Bidirectional attention is applied over context tokens, while response tokens retain causal masking in both loops.The shared mask gives context tokens a complete view of one another and removes the causal asymmetry that makes initial tokens attention sinks.
  • Inner–Outer Consistency: Inner–outer consistency requires matching objectives and attention masks so adapted representations remain compatible with the outer loop rather than causing distortion.Both loops minimize the same response-prediction objective and use the same attention pattern.

4 Experiments

FOCUSFT is evaluated across synthetic reasoning, retrieval-aggregation, real-world question answering, and agentic reasoning benchmarks. The experiments use Qwen2.5-7B and compare against pretrained and standard-SFT baselines under matched training data and budget.

  • Evaluation: FOCUSFT is evaluated on benchmarks spanning synthetic reasoning, retrieval-aggregation, real-world QA, and agentic reasoning.The benchmark suite includes BABILong, RULER, LongBench, and GPQA with agentic tool use.
  • Training Setup: Qwen2.5-7B is trained on 3K multi-turn agentic SFT samples with a maximum sequence length of 4096 tokens for 5 epochs.Training uses an effective batch size of 32 across 8 GPUs with gradient accumulation of 4.
  • Bilevel Hyperparameters: The inner loop uses K=2 gradient steps with learning rate 1.0 on rank-32 LoRA adapters applied to FFN layers in the top 35% of transformer layers.Inner gradients are clipped at norm 1.0.
  • Baselines: Baselines include pretrained Qwen2.5-7B and standard SFT using identical data, model, and training budget without bilevel optimization.Ablations separately remove bilevel optimization or bidirectional context attention.

BABILong

FOCUSFT improves long-context performance across BABILong, RULER, LongBench, and GPQA while reducing attention sinks and redistributing attention toward relevant context. Ablations indicate that bilevel optimization is the primary contributor, with bidirectional attention providing additional benefit when combined with it.

  • BABILong results: +14.2pp, +10.2pp, +10.2pp, and +9.6pp over Standard SFT at 4K, 8K, 16K, and 32K on BABILong.Standard SFT provides essentially no improvement over the pretrained model, while FOCUSFT maintains its advantage beyond the 4K training length.
  • Cross-benchmark results: +8.2pp on RULER CWE at 16K, reaching 81.1% versus 72.9%.CWE aggregates information spread across the full context; FOCUSFT shows a smaller +1.0pp edge on NIAH-MV at 16K.
  • Cross-benchmark results: +2.4pp average F1 on LongBench and +3.8pp GPQA pass@1, reaching 19.4% versus 15.6%.The largest LongBench gain is +5.2pp on MultifieldQA.
  • Ablation study: Bilevel optimization supplies the primary ablation gain, while bidirectional attention alone degrades short-context performance but adds benefit when combined with bilevel optimization.At 32K, the combined gain is +9.6pp and exceeds the sum of individual effects, indicating positive interaction.
  • Ablation study: Performance peaks at lf=0.35 as the inner-loop adaptation fraction balances memory capacity and base-model stability.Too few adapted layers under-capacitate memory, whereas too many over-specialize and disrupt pretrained representations.
  • Efficiency: FOCUSFT costs 1.71× wall time but adds zero inference overhead because inner-loop adapters are discarded after training.The inner loop contributes the dominant 1.52× overhead; bidirectional attention adds negligible cost.
  • Attention analysis: 529× reduction in per-layer attention sink mass, from 30.1% under Standard SFT to 0.06% under FOCUSFT.The reduction is pervasive across all 28 layers and is attributed to bidirectional context attention removing causal asymmetry.
  • Attention analysis: 3.1× higher total context engagement, increasing from 13.5% to 41.4%.System and user prompt attention rises from 7.0% to 27.0%, while tool response attention rises from 6.5% to 14.3%.

5 Related Work

Prior long-context work expands context capacity, reduces attention cost, or addresses inference-time failures, but does not specifically target attention-budget allocation during fine-tuning. FOCUSFT is positioned as complementary to these approaches and draws on bilevel optimization and fast-weight methods.

  • Long-context architectures: Long-context architectures extend sequence reach or efficiency through positional encoding, distributed mechanisms, sparse attention, and noise-reduction designs.These methods expand context visibility or efficiency but do not address attention dilution limiting context utilization.
  • Inference-time methods: Inference-time methods use positional calibration, dynamic score weighting, or test-time parameter adaptation to address long-context failures during deployment.Test-time training is methodologically closest to FOCUSFT because both adapt parameters on the input context via gradient steps.
  • Long-context fine-tuning: Long-context fine-tuning methods improve efficient adaptation, training data, or alignment but do not specifically address attention-budget distribution during fine-tuning.The cited approaches include shifted sparse attention, instruction-following data, and alignment-data construction.
  • Bilevel optimization and fast weights: Bilevel optimization and fast-weight research provides foundations for inner-loop adaptation, task generalization, and secondary parameters updated on faster timescales.FOCUSFT applies these ideas to construct context-dependent parametric memory during fine-tuning.

6 Limitations

FOCUSFT has a 1.71× training-time overhead, and its evaluation is limited to one model family and a fixed 3K-sample corpus. Its extension beyond SFT to RLHF or DPO remains untested.

  • Computational cost: FOCUSFT incurs a 1.71× training-time overhead from inner-loop gradient steps, while inference cost is unchanged.The overhead is attributed to the bilevel inner loop.
  • Evaluation scope: Experiments use a single model family, Qwen2.5-7B, and a fixed training corpus of 3K samples.Scaling behavior across model sizes and larger data regimes remains to be explored.
  • Method scope: The bilevel formulation targets SFT, and integration with RLHF or DPO has not been investigated.The paper identifies these integrations as natural extensions rather than evaluated capabilities.

7 Conclusion

FOCUSFT treats the gap between long-context visibility and utilization as a training-time attention-dilution problem. Its bilevel training combines fast-weight memory, sharpened representations, bidirectional context attention, and inner-outer consistency to mitigate this cycle.

  • Training-time positional biases and attention sinks starve relevant content tokens, corrupting the learning signal and contributing to poor long-context capabilities.
  • FOCUSFT uses fast-weight adaptation in an inner loop to form parametric memory that concentrates attention on semantically relevant content.
  • The outer loop performs SFT on representations sharpened by the inner-loop memory, while bidirectional context attention reduces causal asymmetry linked to attention sinks.
  • Inner-outer consistency aligns sharpened representations with downstream use and helps mitigate the cycle in which diluted training produces poor long-context models.

A Additional Experimental Details

The appendix specifies the main FOCUSFT training configuration and the GPQA Diamond evaluation protocol. GPQA uses multi-turn agentic tool-use rollouts with 32 independent samples per problem and reports both pass@1 and pass@32.

  • Table 5 lists all hyperparameters for the main FOCUSFT configuration.
  • GPQA Diamond evaluation uses the Open-AgentRL framework with multi-turn tool-use rollouts in Hermes chat format.
  • Each of the 198 GPQA problems receives n=32 independent rollouts with temperature 1.0, top-p 0.6, and up to 16 assistant turns.
  • The evaluation reports pass@1 by majority vote and pass@32 as an oracle across all rollouts.

B.1 Per-Task BABILong Breakdown

The per-task BABILong breakdown shows that FOCUSFT’s largest gains occur on tasks requiring multiple dispersed facts, especially QA2 and QA3. These improvements are strongest at 4K for QA2 and across both reported lengths for QA3.

  • +26pp at 4K is the largest reported QA2 gain, on two-fact reasoning over dispersed information.
  • +31pp at 4K and +30pp at 16K are reported for QA3, which requires temporal reasoning over dispersed facts.
  • QA2 and QA3 require locating and connecting multiple dispersed facts, the scenario identified as especially vulnerable to attention dilution.

B.2 RULER 2×2 Ablation

The RULER ablation attributes the primary aggregation-task benefit to bilevel optimization, while bidirectional attention alone slightly degrades performance at 16K. Additional inference-time adaptation is unnecessary because the trained base model already internalizes attention sharpening.

  • RULER 2×2 ablation: Bilevel optimization provides the primary benefit on RULER aggregation tasks, specifically CWE, in the 2×2 ablation.
  • RULER 2×2 ablation: Bidirectional attention alone slightly degrades performance at 16K in the RULER ablation.
  • Inference-time adaptation: Inference-time adaptation does not consistently improve over the base FOCUSFT model and slightly hurts on most benchmarks.
  • Inference-time adaptation: FOCUSFT incurs zero inference overhead because bilevel training improves the base model weights directly and discards inner-loop adapters after training.
  • Attention analysis: The attention analysis uses a representative 4096-token multi-turn agentic sample, with heatmaps from layer 14 and sink mass summed over positions [0:5].
Loading 2605.09932v1…