This repository contains the minimal nanoGPT-based code used for the experiments in "Decoupling the What and Where With Polar Coordinate Positional Embeddings".

The training configs, dataset and preprocessing details are reported in the paper. This README is intended to make the code path easy to follow and to provide representative commands for reproducing the main results.

The code is adapted from Andrej Karpathy's nanoGPT. Thanks to the minimal and hackable nanoGPT repository by Andrej Karpathy, which this codebase reuses.

Create an environment with a recent PyTorch install appropriate for your hardware, then install the Python dependencies:

pip install -r requirements.txt

The optional custom complex FlashAttention path for PoPE uses Triton and is controlled with --complex_flash=True.

Most scripts accept config-file and command-line overrides through configurator.py. For example:

python train.py config/train_indirect_idx.py --pos_type=pope --wandb_log=True

Common overrides:

--pos_type=rope          # RoPE baseline
--pos_type=pope          # PoPE
--base_dir=/path/to/root # train.py reads datasets from $base_dir/data/<dataset>
--compile=False          # useful for CPU/debug runs
--wandb_log=True         # enable Weights & Biases logging

For multi-GPU training, use PyTorch DDP:

torchrun --standalone --nproc_per_node=8 train.py config/train_gpt124m.py --pos_type=pope

For multi-node training, use the usual torchrun --nnodes, --node_rank, --master_addr, and --master_port arguments. If your cluster has no InfiniBand, you may need to prepend NCCL_IB_DISABLE=1.

The following configs correspond to the paper's main training/fine-tuning runs:

Each training config can be run as RoPE or PoPE by overriding pos_type:

python train.py <config.py> --pos_type=rope
python train.py <config.py> --pos_type=pope

Some configs include machine-specific defaults such as base_dir. Override those on the command line rather than editing the config when running on a new system.

Generate the synthetic dataset:

python data/indirect_idx/generate.py

This writes the dataset file to data/indirect_idx/ds_minl20_maxl40_shift_15.txt.

Prepare GPT-2-tokenized OpenWebText:

python data/openwebtext/prepare.py

The training script expects train.bin and val.bin under $base_dir/data/openwebtext. With the default --base_dir='', that is data/openwebtext. The prep script writes there by default; set POPE_DATA_DIR=/path/to/data if you want the generated files and Hugging Face cache under a larger storage volume.

Place Jsb16thSeparated.json in data/jsb/. The loader expects the JSON format with train, valid, and test splits. Download the dataset files from here

Download MAESTRO v3.0.0 MIDI files and place/extract them under data/maestro/, or set --base_dir so that the training path resolves to $base_dir/data/maestro. The loader tokenizes MIDI files with REMI and creates local chunk directories.

The HRG loader uses the Hugging Face dataset InstaDeepAI/human_reference_genome. It downloads through datasets; set your normal Hugging Face cache variables or POPE_HF_CACHE_DIR if needed.

length_gen.py uses the Hugging Face dataset emozilla/pg19-test for length extrapolation evaluation. It loads a saved checkpoint and writes loss-vs-length output under length_gen/.

python data/indirect_idx/generate.py

python train.py config/train_indirect_idx.py \
  --pos_type=rope \
  --wandb_run_name=indirect-rope

python train.py config/train_indirect_idx.py \
  --pos_type=pope \
  --wandb_run_name=indirect-pope

Prepare the data, then launch RoPE and PoPE runs with the same config:

python data/openwebtext/prepare.py

torchrun --standalone --nproc_per_node=8 train.py config/train_gpt124m.py \
  --pos_type=rope \
  --wandb_run_name=owt-124m-rope

torchrun --standalone --nproc_per_node=8 train.py config/train_gpt124m.py \
  --pos_type=pope \
  --wandb_run_name=owt-124m-pope

Use config/train_gpt253m.py and config/train_gpt774m.py for the larger model sizes.

After preparing the relevant dataset, run paired RoPE/PoPE jobs in the same pattern:

python train.py config/train_jsb.py --pos_type=rope --wandb_run_name=jsb-rope
python train.py config/train_jsb.py --pos_type=pope --wandb_run_name=jsb-pope

python train.py config/train_maestro.py --pos_type=rope --wandb_run_name=maestro-rope
python train.py config/train_maestro.py --pos_type=pope --wandb_run_name=maestro-pope

torchrun --standalone --nproc_per_node=8 train.py config/train_hrg205m.py \
  --pos_type=rope \
  --wandb_run_name=hrg-rope

torchrun --standalone --nproc_per_node=8 train.py config/train_hrg205m.py \
  --pos_type=pope \
  --wandb_run_name=hrg-pope

Fine-tuning configs are provided for OpenWebText context extension:

torchrun --standalone --nproc_per_node=4 train.py config/finetune_gpt124m.py \
  --pos_type=rope \
  --complex_flash=True

torchrun --standalone --nproc_per_node=4 train.py config/finetune_gpt124m.py \
  --pos_type=pope \
  --complex_flash=True

Evaluate saved checkpoints on PG-19 with length_gen.py:

python length_gen.py \
  --base_dir=/path/to/root \
  --ckpt_dir=final-owt-ckpts \
  --ckpt_fname=gpt2-124M-pope-ckpt.pt \
  --pos_type=pope \
  --compile=False

The script evaluates sequence lengths from the training context length up to 10x that length.

If you use PoPE, the Indirect Indexing task, or build on the experimental results or what-where decoupling analysis, please cite:

@inproceedings{gopalakrishnan2026decoupling,
  title={Decoupling The ''What'' and ''Where'' With Polar Coordinate Positional Embedding},
  author={Gopalakrishnan, Anand and Csord{\'a}s, Robert and Schmidhuber, J{\"u}rgen and Mozer, Michael Curtis},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=I3Z9za1EkO}
}