Selective Rotary Position Embeddings

This repository contains the official implementation of the paper Selective Rotary Position Embedding (ICLR 2026).

We introduce Selective RoPE, an input-dependent rotary embedding that enhances gated linear transformers. We demonstrate that softmax attention performs implicit rotations on query-key pairs and show how real and imaginary components in state-space models manage forgetting and positional encoding respectively. Our method improves performance on language modeling and sequence tasks including copying, state tracking, and retrieval.

Installation

Make sure you have uv installed:

curl -LsSf https://astral.sh/uv/install.sh | sh

Clone the repository and sync the environment:

git clone https://github.com/<org>/selective-rope.git
cd selective-rope
uv sync -p 3.12

Usage

All commands should be run through uv run:

uv run --frozen --no-sync python <script.py>

Before scheduling experiments

Before running any scheduling script (schedule_lm.py, schedule_mad.py, etc.), fill in the following placeholders that were left null for the public release:

File	Field	Description
configs/language_modeling/cluster/capella.yaml	data_home	Path to tokenized training data
configs/language_modeling/language_modeling.yaml	logger.wandb_project	WandB project name
configs/*/logger/logger.yaml	wandb_entity	WandB entity (user or team)

For the evaluation script (configs/language_modeling/scripts/capella_eval.sh), set HF_HOME and HF_DATASETS_CACHE if you want HuggingFace data cached outside the default ~/.cache/huggingface.

Reproducing Experiments

Language Modeling

Training is launched via SLURM using the scheduling script, which creates isolated git worktree snapshots of the current code state:

uv run --frozen --no-sync python schedule_lm.py

Configuration is Hydra-based. Position embedding variants are swapped via config groups:

• model/position_embedding=nope — no positional embedding (baseline)
• model/position_embedding=rope — standard RoPE
• model/position_embedding=selective_rope — Selective RoPE (ours)

See configs/language_modeling/ for the full configuration structure.

MAD Benchmark

First, pre-generate the data:

uv run --frozen --no-sync python synthetic_tasks/mad/generate_data.py \
    --data-path ./data/mad --num-workers 8

Then schedule training:

uv run --frozen --no-sync python synthetic_tasks/mad/schedule_mad.py

Copying & State Tracking

These follow the same Hydra + SLURM pattern. See configs/copying/ and configs/state_tracking/ for configurations.

Citation

@inproceedings{movahedi2026selective,
  title={Selective Rotary Position Embedding},
  author={Movahedi, Sajad and Carstensen, Timur and Afzal, Arshia and Hutter, Frank and Orvieto, Antonio and Cevher, Volkan},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=AQo1SEElNb}
}

Acknowledgments

The synthetic evaluation tasks build on MAD-Lab, Zoology, and DeltaProduct.

License

This project is licensed under the Apache License 2.0 — see LICENSE for details.