Transformers improve position embedding by adjusting layer usage
Where Position Matters: Profiling-Guided Layer-Wise RoPE Allocation in Transformers
Machine LearningArtificial Intelligence
Summary
Transformers use position information to understand the order of words, but they usually treat all their layers the same when applying this information. The authors found that different layers actually need different amounts of position info. They developed a way to measure and assign how much position data each layer should get, improving performance without using more resources. This method works across different model sizes and shows better results when training from scratch.
What this means in practice
- •For machine learning engineers: Optimize transformer models by assigning position embeddings unevenly across layers to improve efficiency and accuracy under fixed memory budgets.
- •For ai hardware designers: Design transformer accelerators that allocate rotary embedding resources physically per layer, enhancing compute-storage trade-offs.
Authors
Zhiwei Liu, Boao Kong, Bing Liu
Abstract
Rotary Position Embedding (RoPE) typically uses uniform coverage across Transformer layers despite heterogeneous positional requirements. Phase-off interventions reveal non-uniform, model-specific rotary-phase dependence, while coverage interventions relate it to coverage-upgrade benefits. We propose profiling-guided layer-wise RoPE allocation, assigning high/low coverage under a fixed budget using early coverage-benefit profiles. The framework supports fixed-width query/key representations and physical rotary-dimension allocation in Multi-head Latent Attention (MLA). Normalized-depth mapping transfers profiles within model families, enabling small-model measurements to guide larger-model allocations before pretraining. A brief analysis relates selection and transfer to score calibration and inter-layer interactions. From-scratch experiments demonstrate validation-loss improvements over available budget-matched baselines, while analytical resource accounting identifies favorable quality--storage--compute trade-offs for physical MLA. These results connect positional demand to transferable, efficient architecture design.