AdaRoPE

Code for AdaRoPE: Not All Attention Heads Should Rotate and Scale Equally (ICML 2026).

AdaRoPE learns per-head rotary frequencies (AdaFreq) and length-dependent attention scales (AdaScale).

git clone https://github.com/Outsider565/AdaRoPE.git
cd AdaRoPE
pip install -e ".[test]"

Requires Python 3.10 or later and Transformers 4.57.1.

python examples/llama_step.py

The example builds a small Llama model and takes one optimizer step. To create a model:

from adarope import AdaRoPEConfig, AdaRoPELlamaForCausalLM

config = AdaRoPEConfig(
    vocab_size=32000,
    hidden_size=512,
    intermediate_size=1376,
    num_hidden_layers=8,
    num_attention_heads=8,
    num_key_value_heads=4,
    max_position_embeddings=4096,
    adascale_ref_len=64,
)
model = AdaRoPELlamaForCausalLM(config)

adafreq and adascale enable the two components. adascale_ref_len sets the reference length; by default it uses the original context length. rope_theta and rope_scaling control initialization, including YaRN. Pass attention_mask for padded batches.

import adarope
from transformers import AutoModelForCausalLM

model.save_pretrained("runs/adarope-llama")
model = AutoModelForCausalLM.from_pretrained("runs/adarope-llama")

For an existing Llama checkpoint, convert its configuration with AdaRoPEConfig.from_llama_config and load it with AdaRoPELlamaForCausalLM.from_pretrained. Loading with dtype=torch.bfloat16 keeps the AdaRoPE parameters in fp32.

adarope-synthetic --task window --positional rope --seq_len 256 --center 128 --width 16 \
  --d_model 64 --n_heads 1 --n_layers 1 --steps 4000 --batch_size 256 --save_dir runs/window_rope
adarope-synthetic --task window --positional adafreq --seq_len 256 --center 128 --width 16 \
  --d_model 64 --n_heads 1 --n_layers 1 --steps 4000 --batch_size 256 --save_dir runs/window_adafreq
adarope-plot --checkpoints runs/window_rope runs/window_adafreq --delta_step 1 --save_dir runs/figures

Use --adascale --adascale_ref_len 64 to enable AdaScale, --device cpu to run on a CPU, and --help to see the available tasks and options.

python -m pytest

@inproceedings{wang2026adarope,
  title={AdaRoPE: Not All Attention Heads Should Rotate and Scale Equally},
  author={Wang, Shaowen and Zheng, Yuke and Zhu, Tansheng and Chen, Shuang and Liu, Shaofan and Zheng, Suncong and Li, Jian},
  booktitle={Proceedings of the 43rd International Conference on Machine Learning},
  year={2026}
}

arXiv: https://arxiv.org/abs/2607.19363

Apache 2.0. src/adarope/modeling_adarope.py is adapted from the Llama implementation in Hugging Face Transformers; see NOTICE.