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

Abstract: LLMs still struggle with the ``lost-in-the-middle'' problem, where critical information located in the middle of long-context inputs is often underrepresented or lost. While existing methods attempt to address this by combining multi-scale rotary position embeddings (RoPE), they typically suffer from high latency or rely on suboptimal hand-crafted scaling strategies. To overcome these limitations, we introduce a layer-specific positional embedding scaling~(LPES) method that assigns distinct scaling factors to each layer. LPES achieves a more balanced attention distribution without fine-tuning model parameters or increasing inference delay. A specially designed genetic algorithm is employed to efficiently select the optimal scaling factors for each layer by incorporating Bézier curves to significantly reduce the search space. Extensive experiments demonstrate that LPES effectively mitigates positional attention bias and delivers consistent improvements across multiple long-context benchmarks, yielding up to an $11.2$\% accuracy gain on the key-value retrieval dataset.

• The paper introduces LPES, a method that assigns unique RoPE scaling factors per transformer layer to counter the lost-in-the-middle problem.
• LPES uses Bézier curve parameterization and a genetic algorithm to efficiently optimize scaling factors, achieving up to 11.2% accuracy improvement in mid-sequence token retrieval.
• The approach is training-free, preserves model representational structure, and delivers significant speedups while generalizing across various transformer models and tasks.

Layer-Specific Positional Embedding Scaling for Mitigating Transformer Position Bias

Background and Motivation

Transformer-based LLMs leveraging RoPE exhibit pronounced position bias as context length increases, manifesting in degraded attention to mid-sequence tokens (the "lost-in-the-middle" problem). Standard RoPE encodes relative positional information but inherently biases attention toward sequence boundaries. Previous solutions, notably Attention Buckets and MoICE, attempt to blend multi-scale RoPE variants for more balanced attention distributions but incur substantial computational and latency overheads due to either model-level or module-level ensembling and rely on manual or heuristic scaling factor selection.

The discussed work introduces Layer-specific Positional Embedding Scaling (LPES), addressing core inefficiencies of these approaches by shifting the locus of optimization to the layer level, using distinct RoPE scaling factors per transformer layer. This configuration enables uniform context utilization without modifying the model's weights or incurring additional inference costs.

Methodology

Layer-Specific Positional Scaling

LPES assigns a unique scaling factor $s_\ell$ to the RoPE computation at each transformer layer $\ell$. By adjusting these factors, LPES flattens the RoPE-induced long-range attention decay and disruptive "attention waves," systemically reducing position bias. The optimal set of scaling factors is not straightforward to determine given the exponential configuration space ($M^L$ for $M$ candidate factors and $L$ layers).

Bézier Curve Parameterization

To reduce the search dimensionality, LPES parameterizes the sequence of scaling factors as a low-degree Bézier curve over layers, with the y-coordinates of the control points encoding the scaling factors. This constraint imposes smoothness and continuity across layers—a property empirically observed to be a favorable inductive bias for layer-wise modifications—while compactly representing performant scaling assignments. The Bézier specification drastically shrinks the search space from $M^L$ (brute force) to roughly $(M \times L)^C$ for $C$ control points, enabling efficient search.

Genetic Algorithm Search

LPES employs a domain-constrained genetic algorithm to identify optimal Bézier control points. The population encodes candidate curves; fitness reflects model performance on proxy tasks using datasets where the relevant information appears at varied input positions, favoring solutions that address position bias directly. Offspring are generated via crossover and mutation with monotonicity constraints on the control points, ensuring interpretable, smooth parameterizations. Fitness evaluation leverages a weighted sum of accuracies over beginning, middle, and end-context tasks, with parameter insensitivity validated empirically.

This strategy converges typically within hours on several hundred training samples and can be executed efficiently with modest compute resources.

Empirical Results

Position Bias Mitigation and Efficiency

LPES demonstrates substantial reductions in position bias across evaluated LLMs (Vicuna-7B-v1.5, StableBeluga-7B, Qwen2.5-7B) and long-context QA or retrieval benchmarks (MDQA, key-value retrieval, ZeroSCROLLS, L-Eval):

• LPES achieves up to 11.2% higher accuracy on key-value retrieval at middle-context positions relative to RoPE-based baselines.
• Layer-specific scaling consistently outperforms both single-scale (PI) and previous multi-scale methods, both in terms of accuracy and uniformity of attention distribution.
• LPES achieves model- and module-level ensembling performance with a single forward pass, yielding speedups of 2.42x over MoICE and 1.45x over Ms-PoE.

Generalization and Transferability

Scaling factors optimized for one task (e.g., MDQA) transfer robustly to other tasks (ZeroSCROLLS, L-Eval) and models, demonstrating that LPES captures architectural, not task-specific, position bias.

Preservation of Representational Structure

The paper provides substantial quantification of internal representational stability by measuring layer-wise representational similarity matrices. Bézier-parameterized scaling yields lower global representational deviation compared to uniform, noisy, shuffled, or random curve parameterizations, indicating minimal distortion of the learned hierarchical structure.

Optimization Stability and Ablation

LPES is robust with respect to the choice and number of curve control points (with four often sufficient), the precise hyperparameterization of the genetic search, and dataset sampling. Non-smooth or high-frequency scaling assignments are suboptimal, and direct gradient-based optimization of scaling factors is ineffective.

Theoretical and Practical Implications

LPES decisively demonstrates that transformer position bias can be neutralized by introducing a parameter-efficient, training-free, layer-wise RoPE scaling regimen optimized over smooth curve manifolds, avoiding weight updates, catastrophic forgetting, and inference delays. As LPES operates purely as a plug-and-play transformation at the embedding level, it is compatible with post-deployment models, low-resource adaptation scenarios, and efficient production settings—addressing deployment bottlenecks associated with previous methods.

From a theoretical perspective, the success of smooth layer-wise modulation motivates further research into continuous parameterization of layer operations, especially for long-context and hierarchical attention problems.

Limitations and Future Work

The current instantiation of LPES is “training-free,” optimizing only hyperparameters and not model weights. While this ensures deployment simplicity, future work could examine integration with gradient-based fine-tuning protocols or dynamic, sample-adaptive scaling schedules. The extension of LPES to even larger models and contexts, or to cross-modal architectures, is a natural direction for investigation.

Conclusion

LPES establishes that position bias in transformers—especially the lost-in-the-middle effect—can be rectified efficiently by introducing layer-level scaling of RoPE, parameterized via Bézier curves and optimized using a constrained genetic algorithm. This method achieves marked accuracy improvements with minimal computational overhead, preserves model generalization, and supports straightforward deployment. The methodology offers a significant step toward robust, context-length-agnostic transformer architectures.

For comprehensive technical details, refer to "Mitigating Position Bias in Transformers via Layer-Specific Positional Embedding Scaling" (2606.27705).