Mitigating Position Bias in Transformers via Layer-Specific Positional Embedding Scaling

Technical Summary: Mitigating Position Bias in Transformers via Layer-Specific Positional Embedding Scaling

Problem Statement
Large Language Models (LLMs) based on Transformers, particularly those utilizing Rotary Position Embeddings (RoPE), suffer from a "lost-in-the-middle" phenomenon. As context length increases, these models tend to over-attend to tokens at the beginning and end of the input sequence while neglecting critical information located in the middle. Existing solutions, such as Attention Buckets and MoICE, attempt to mitigate this by combining multiple RoPEs with different bases or scaling factors. However, these approaches typically rely on manually designed rules and require multiple forward passes (model-level or module-level ensembling) during inference. This results in significant computational overhead, increased latency, and high inference costs, limiting their practical applicability. Furthermore, determining optimal scaling factors is a combinatorial optimization problem that is computationally intractable for exhaustive search and difficult to solve via gradient-based methods due to the discrete nature of layer-specific assignments.

Methodology
The authors propose Layer-Specific Positional Embedding Scaling (LPES), a training-free method that assigns distinct scaling factors to each layer of the Transformer without modifying model parameters or increasing inference latency.

1. Layer-Specific Scaling: Unlike methods that apply a single global scaling factor or multiple bases per attention head, LPES assigns a unique scaling factor $s_h$ to each layer $h$. This allows the model to combine information from multiple scaled RoPEs within a single forward pass.
2. Bézier Curve Constraint: To address the intractability of searching for $L$ independent scaling factors (where the search space size is $M^L$), the authors model the relationship between layer depth and scaling factors using Bézier curves. By defining the curve with a small set of discrete control points $P_k$, the search space is drastically reduced from $M^L$ to $(M \times L)^C$, where $C$ is the number of control points. This approach leverages the observation that smooth, continuous variations in scaling factors serve as a beneficial inductive bias, preserving the model's layer-wise representational structure.
3. Curve-Constrained Genetic Algorithm: The authors employ a specialized genetic algorithm to optimize the control points of the Bézier curve.
   • Fitness Function: The algorithm evaluates individuals based on model performance across three sub-datasets where relevant information appears at the beginning, middle, and end of the input.
   • Constraints: The mutation and crossover operators enforce monotonicity in the x-coordinates of control points and limit the range of y-coordinate changes to ensure smoothness and prevent abrupt scaling shifts.
   • Efficiency: The optimal scaling configuration can be found using only a few hundred examples (e.g., 200 instances) on standard hardware (four H100 GPUs) within 3–4 hours.

Key Contributions

• LPES Framework: A novel method that mitigates positional bias by assigning layer-specific scaling factors, achieving balanced attention distribution without fine-tuning model weights or incurring additional inference latency.
• Efficient Optimization: The introduction of a genetic algorithm constrained by Bézier curves, which efficiently solves the combinatorial optimization problem of finding optimal scaling factors, reducing the search space and enabling rapid convergence.
• Empirical Validation: Comprehensive experiments demonstrating that LPES preserves general capabilities while significantly improving long-context performance across multiple benchmarks and model architectures.

Results
Extensive experiments were conducted on models including Vicuna-7B, StableBeluga-7B, and Qwen2.5-7B across benchmarks such as MDQA, Key-Value Retrieval, ZeroSCROLLS, and L-Eval.

• Performance Gains: LPES consistently outperformed baselines (including Positional Interpolation, Attention Buckets, Ms-PoE, and MoICE). Notably, it achieved up to an 11.2% accuracy gain on the Key-Value Retrieval dataset for Vicuna-7B.
• Positional Bias Mitigation: The method significantly improved accuracy for information located in the middle of the context (50% position), addressing the "lost-in-the-middle" issue more effectively than existing methods.
• Inference Efficiency: LPES requires only a single forward pass. It is reported to be 2.42× faster than MoICE and 1.45× faster than Ms-PoE, while maintaining the same inference time as the baseline.
• Generalization: Scaling factors learned on one dataset (MDQA) transferred effectively to other long-text tasks (ZeroSCROLLS, L-Eval) and different models, without degrading general capabilities on MMLU and C-Eval benchmarks.
• Inductive Bias: Analysis of representational structure deviation showed that smooth Bézier-based scaling causes minimal disruption to the model's internal layer-wise representations compared to noisy or shuffled scaling strategies.

Significance
The paper claims that LPES offers a broadly applicable and efficient solution to the positional bias problem in long-context LLMs. By avoiding the computational overhead of multi-pass inference and the complexity of full model fine-tuning, LPES enables the deployment of long-context capabilities in already trained models. The authors emphasize that their approach transforms a difficult combinatorial search into a tractable optimization problem via Bézier curves, providing a practical path to balanced attention distribution without the latency penalties associated with previous ensemble-based methods. The work highlights that smooth, layer-specific scaling acts as a powerful inductive bias, naturally emerging as a high-performing configuration.