code/README.md

# Minimal PoPE code repository

This repository contains the minimal nanoGPT-based code used for the experiments in
["Decoupling the What and Where With Polar Coordinate Positional Embeddings"](https://arxiv.org/pdf/2509.10534).

![RoPE vs PoPE](assets/rope-vs-pope.png)

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](https://github.com/karpathy/nanoGPT). Thanks to the minimal and hackable nanoGPT repository by Andrej Karpathy, which this codebase reuses.

## Setup

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

```sh
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:

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

Common overrides:

```sh
--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:

```sh
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`.

## Main configs

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

| Experiment | Config |
| --- | --- |
| Indirect indexing | `config/train_indirect_idx.py` |
| OpenWebText 124M | `config/train_gpt124m.py` |
| OpenWebText 253M | `config/train_gpt253m.py` |
| OpenWebText 774M | `config/train_gpt774m.py` |
| JSB Chorales | `config/train_jsb.py` |
| MAESTRO | `config/train_maestro.py` |
| Human Reference Genome | `config/train_hrg205m.py` |
| OpenWebText length fine-tuning | `config/finetune_gpt124m.py`, `config/finetune_gpt253m.py` |

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

```sh
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.

## Data preparation

### Indirect indexing

Generate the synthetic dataset:

```sh
python data/indirect_idx/generate.py
```

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

### OpenWebText

Prepare GPT-2-tokenized OpenWebText:

```sh
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.

### JSB Chorales

Place `Jsb16thSeparated.json` in `data/jsb/`. The loader expects the JSON format with
`train`, `valid`, and `test` splits. Download the dataset files from [here](https://github.com/czhuang/JSB-Chorales-dataset)

### MAESTRO

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.

### Human Reference Genome

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.

### PG-19 length evaluation

`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/`.

## Representative runs

### Indirect indexing

```sh
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
```

### OpenWebText language modeling

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

```sh
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.

### JSB, MAESTRO, and HRG

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

```sh
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
```

### Length extrapolation

Fine-tuning configs are provided for OpenWebText context extension:

```sh
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`:

```sh
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.

## Citation

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

```bibtex
@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}
}
```