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

2026-06-26

AI Summary

This work addresses the “middle information loss” problem in large language models caused by positional bias in long-context inputs. The authors propose Layer-specific Positional Embedding Scaling (LPES), which assigns independent positional scaling factors to each Transformer layer to balance attention distributions across the sequence. The optimal scaling strategy is efficiently discovered via a genetic algorithm parameterized with Bézier curves, requiring no model fine-tuning or additional inference latency. Evaluated on multiple long-context benchmarks, LPES consistently improves performance, achieving up to an 11.2% absolute gain in accuracy on key-value retrieval tasks. This approach overcomes limitations of conventional handcrafted designs or methods entailing high computational overhead.

Abstract

Large Language Models (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.

Problem

• position bias
• lost-in-the-middle
• long-context
• attention distribution
• positional embedding

Innovation

• Layer-Specific Positional Embedding Scaling
• Position Bias Mitigation
• Long-Context Modeling
• Genetic Algorithm
• Bezier Curve Optimization