One Learning Rate, Many Positional Steps: Position-Aware RoPE Preconditioning
Abstract
Rotary positional embeddings (RoPE) assign different frequencies to two-dimensional query/key pairs, yet standard optimizers ignore this structure. We identify a potential RoPE-induced learning-rate mismatch: the same parameter-space turn can induce different positional steps across pairs, with sensitive pairs potentially limiting the stable global learning rate. We derive an analytic sensitivity for each pair and show that the same frequency-dependent structure appears in a RoPE-specific component of the local generalized Gauss–Newton geometry. Based on this, we propose Position-Aware Preconditioning (PA), which preserves the scale-changing part of a base-optimizer update and recalibrates only its turning part using RoPE information. PA requires no extra forward/backward pass or per-parameter optimizer state and retains first-order convergence guarantees for its gradient specialization. On C4 pretraining with 130M- and 350M-parameter LLaMA-style models, PA improves AdamW and Muon, reaching baseline training-loss levels with 13–27% fewer optimizer updates.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.