---
license: mit
datasets:
- iohadrubin/wikitext-103-raw-v1
language:
- en
metrics:
- perplexity
base_model:
- openai-community/gpt2
pipeline_tag: text-generation
tags:
- length-extrapolation
- context-aware
- positional-encoding
- Cable
library_name: transformers
---
## Context-aware Biases for Length Extrapolation 

The source code of [(Context-aware Biases for Length Extrapolation)](https://arxiv.org/abs/2503.08067)

<p align="center">
 <img src="https://raw.githubusercontent.com/axiomlab/Cable/refs/heads/main/Figures/pull_figure.png"  width="400" height="300"/>
</p>

### 🚀 News
- [2025.02.3] Code release

#### Upcoming
- [ ] Cleaning codebase
- [ ] Adding scripts for training ALiBi, RoPE, T5-bias

### Datasets and Models
- Fineweb-Edu [🔗](https://arxiv.org/abs/2406.17557) [🤗](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)
- Fineweb [🔗](https://arxiv.org/abs/2406.17557) [🤗](https://huggingface.co/datasets/HuggingFaceFW/fineweb)
- WikiText-103 [🔗](https://arxiv.org/abs/1609.07843) [🤗](https://huggingface.co/datasets/iohadrubin/wikitext-103-raw-v1)
- WikiText-2 [🔗](https://arxiv.org/abs/1609.07843) [🤗](https://huggingface.co/datasets/mindchain/wikitext2)

Download the datasets from HuggingFace and use use ```dataset_preparation.py``` for saving tokenized dataset.


Some of trained models:
| Dataset | Model | Parameters |Sequence Length | Checkpoint |
| -------- | :-------: | :-------: | :-------: | :-------: |
| Fineweb-Edu(10B) | GPT-Medium | 334M | 1024 | [Link](https://huggingface.co/axiomlaborg/cable-edufineweb-md-1024) |
| Fineweb-Edu(10B) | GPT-Medium | 334M | 512 | [Link](https://huggingface.co/axiomlaborg/cable-edufineweb-md-512) |
| WikiText-103 | GPT-Tiny | 44M | 1024 | [Link](https://huggingface.co/axiomlaborg/cable-wiki-tiny-1024) |
| WikiText-103 | GPT-Tiny | 44M | 512 | [Link](https://huggingface.co/axiomlaborg/cable-wiki-tiny-512) |


### Training
- Single GPU
  ```shell
  python Cable.py --dataset-dir "path to dataset" --model "medium or small or tiny" --save-dir "dir for logs"
  ```

- Multiple GPUs
  ```shell
  torchrun --standalone --nproc_per_node=2 Cable.py
  ```

For Hellaswag benchmark and evaluating extrapolation please use ```evaluation.ipynb``` notebook.


### Length Extrapolation

<p align="center">
 <img src="https://raw.githubusercontent.com/axiomlab/Cable/refs/heads/main/Figures/ablation.png"  width="400" height="300"/>
</p>

A Cable model trained on T=1024 can extrapolate on T=8192, achieving a better performance (PPL=22.22) compared to the sinusoidal model (PPL=22.81) trained on T=8192.

### Runtime and Memory Overhead
<p align="center">
 <img src="https://raw.githubusercontent.com/axiomlab/Cable/refs/heads/main/Figures/time.png"  width="800" height="300"/>
</p>

Cable improves the model's extrapolation ability significantly with a negligible burden in time and memory compared to the vanilla transformer. Furthermore, compared to existing RPE methods, our approach maintains nearly identical training time and GPU memory usage, while its inference overhead remains either negligible or comparable, depending on the sequence length.
 
 ## Citation
If you use this repository for your research or wish to refer to our positional encoding method, please use the following BibTeX entry:
```bibtex
@article{veisi2025context,
  title={Context-aware Biases for Length Extrapolation},
  author={Ali Veisi and Amir Mansourian},
  journal={arXiv preprint arXiv:2503.08067},
  year={2025}
}
```

### Acknowledgement
This repo is based on [Karpathy/Build-NanoGPT](https://github.com/karpathy/build-nanogpt). Thanks for their excellent work.