SCOPE: Scoped Position Encoding for Efficient LLMs

Official Implementation of "SCOPE: Boosting LLM Efficiency with Scoped Position Encoding"

This repository contains the official implementation of SCOPE (Scoped Position Encoding), a novel positional encoding framework that reimagines structured sparsity as an intrinsic position encoding mechanism for Transformer-based language models.

Built on top of torchtitan v0.2.0, this codebase enables large-scale pre-training and evaluation of LLaMA, Qwen, and DeepSeek architectures with ScoPE and various baseline positional encoding schemes.

---

Overview

Motivation: Positional encodings are essential to Transformers, yet explicit methods rely on additional arithmetic operations and position-related parameters. Meanwhile, heuristic sparse attention patterns reduce FLOPs but may trade modeling fidelity for computational savings. We ask: can structured sparsity itself function as a position encoding mechanism?

Our Solution — ScoPE: ScoPE assigns exponentially distributed look-back scopes to attention heads, enabling the model to capture dependencies at diverse granularities without explicit arithmetic position signals. This design offers a different trade-off from standard explicit encodings:

• No position parameters: order awareness emerges from the attention topology rather than learned embeddings or rotary matrices
• Compatible efficiency: shorter scopes reduce active computation for most heads, while at least one head retains full context coverage

Note on Complexity: The asymptotic complexity remains $O(T^2)$, as the largest scope covers the full context history. The efficiency gains come from reducing the computational coefficient by sparsifying most heads, rather than changing the asymptotic order.

Key Results

Capability	Result
Native Extrapolation	Stable up to $4\times$ training context (e.g., 4k $\to$ 16k) without any adaptation
Long-Context Retrieval	>90% accuracy on Needle-in-a-Haystack at 128k context
Inference Speedup	$2.1\times$ prefill speedup at 256k tokens vs. RoPE
Training Efficiency	Consistently lower training loss than RoPE baseline; up to $7.3\times$ theoretical FLOP reduction
NLU Benchmarks	Competitive or superior on 8/10 tasks (ARC-C, HellaSwag, WinoGrande, etc.)

---

Installation

This repository is based on torchtitan v0.2.0. The installation requirements are the same as the upstream project. Please refer to the torchtitan installation guide for the recommended PyTorch version and environment setup.

git clone https://github.com/oncemoe/ScoPE.git
cd ScoPE
pip install -r requirements.txt

For evaluation scripts, you may also need:

pip install transformers accelerate lm-eval pyyaml

Note: FlexAttention is used for efficient training with per-head scope masks. Ensure your PyTorch version supports torch.nn.attention.flex_attention.

Context Parallel (CP) for Long-Context Training: To scale to 32k/128k context lengths, we enable Context Parallel together with FlexAttention. This requires a sufficiently recent PyTorch version (nightly >= 2.5 recommended) that supports torch.distributed.tensor.experimental.context_parallel. In our codebase, the original torchtitan restriction on CP + FlexAttention has been relaxed to allow this combination.

---

Quick Start

1. Training

We provide pre-configured TOML files under configs/ for small scale experiments. Below are concrete commands for training Llama3 models with ScoPE.

Small-scale pre-training (Llama3 nano):

CONFIG_FILE=configs/llama3-nano.toml \
  NGPU=8 bash run.sh \
  --model.pe_class_name ScoPE \
  --model.flavor nano \
  --training.global_batch_size 96 \
  --training.seq_len 1024 \
  --optimizer.lr 3e-4

Llama3 8B pre-training:

CONFIG_FILE=configs/llama3-nano.toml \
  NGPU=8 bash run.sh \
  --model.pe_class_name ScoPE \
  --model.flavor 8B \
  --training.global_batch_size 96 \
  --training.seq_len 4096 \
  --optimizer.lr 3e-4

For Qwen3 small-scale experiments, use CONFIG_FILE=configs/qwen3-0.6B.toml and --model.flavor 0.6B instead.

Supported pe_class_name values:

• RoPE — Standard Rotary Position Embedding (baseline)
• ScoPE — Scoped Position Encoding (our method)
• ALiBi — Attention with Linear Biases
• NoPE — No positional encoding (causal mask only)

2. Evaluation

We provide evaluation scripts for reproducing the main results in the paper:

Needle-in-a-Haystack (NIAH)

cd scripts/niah
python run.py --config configs/your_config.yaml

Perplexity on Long Documents

cd scripts/perplexity
# First, prepare long-text data
python select_text.py --output_dir ./proof-pile-10k
# Then run evaluation
python run.py --config configs/your_config.yaml

Standard NLU Benchmarks (lm-eval-harness)

cd scripts/lm_eval
python run.py --config configs/your_config.yaml

Linear Probing for Position Reconstruction

cd scripts/probe
python run.py --config configs/your_config.yaml

Each evaluation script supports batch-config execution via YAML/JSON files. See the respective configs/ subdirectories for examples.

---

Repository Structure

This repository is based on torchtitan v0.2.0. Below we highlight the key modifications and additions:

ScoPE/
├── torchtitan/
│   ├── models/
│   │   ├── mask.py                    # NEW: Core ScoPE implementation
│   │   ├── llama3/model/model.py      # MODIFIED: Integrated ScoPE + YaRN + Partial RoPE
│   │   ├── qwen3/model/model.py       # MODIFIED: Integrated ScoPE + YaRN + ALiBi
│   │   ├── deepseek_v3/model/model.py # MODIFIED: Integrated ScoPE
│   │   ├── llama3/__init__.py         # MODIFIED: Added nano, 2B, 8BP, 8BLL flavors
│   │   ├── qwen3/__init__.py          # MODIFIED: Added ScoPE model flavors
│   │   └── ...
│   ├── inference/                     # NEW: Efficient inference kernels
│   │   ├── flash_mask_attention.py    #   Triton-based FlashAttention with per-head masks
│   │   ├── sdpa_mask_attention.py     #   SDPA wrapper with scope masking
│   │   ├── modeling.py                #   HuggingFace-compatible inference model
│   │   └── ...
│   ├── datasets/hf_datasets.py        # MODIFIED: Added RefinedWeb, long-context mix
│   ├── config/job_config.py           # MODIFIED: Added PE config fields
│   └── train.py                       # MODIFIED: Minor adaptions for tokenizer passing
│
├── scripts/
│   ├── niah/                          # Needle-in-a-Haystack evaluation
│   ├── lm_eval/                       # lm-evaluation-harness wrapper
│   ├── perplexity/                    # Long-document perplexity evaluation
│   └── probe/                         # Linear probing for position reconstruction
│
├── configs/                           # NEW: Training presets
│   ├── llama3-nano.toml
│   └── qwen3-0.6B.toml
│
├── run.sh                             # torchtitan launcher script
└── run_train.sh                       # Original torchtitan training entry

Core Implementation: torchtitan/models/mask.py

The heart of ScoPE lies in assigning a unique look-back scope to each attention head. Given $H$ heads and max sequence length $T$, scopes are assigned via:

# Exponential scope assignment (ScoPE)
scopes = [int(T ** (i / H)) for i in range(1, H + 1)]

This ensures heads collectively cover the context at diverse granularities, from local (small scope) to global (large scope). The scope mask is then applied via FlexAttention:

def get_scope_mask_mod(scopes):
    def scope_mask(b, h, q_idx, kv_idx):
        return (q_idx - kv_idx) <= scopes[h]
    return scope_mask

For inference, we also provide Triton-based FlashAttention kernels that support per-head window sizes (torchtitan/inference/flash_mask_attention.py), enabling efficient decoding with ScoPE.

---

Supported Architectures & Configurations

Architecture	Supported PEs	Notes
Llama 3	RoPE, ScoPE, MAPE, FAPE, Mistral, NoPE	Full training & inference support
Qwen 3	RoPE, ScoPE, MAPE, FAPE, ALiBi, R-ScoPE, NoPE	Includes MoE variants
DeepSeek V3	RoPE, ScoPE	Initial integration for MLA architecture

Additional model flavors added:

• nano / 0.6B — Small-scale models for rapid prototyping and probing experiments
• 2B — Mid-size Llama3 variant
• 8BP — Llama3-8B with Partial RoPE (reduced rotation dim)
• 8BLL / 8BLLL — Llama3-8B with YaRN scaling (4× / 8×)

---

Citation

If you find this work useful, please cite our paper:

@inproceedings{qi2026scope,
  title={SCOPE: Boosting LLM Efficiency with Scoped Position Encoding},
  author={Qi, Qingguo and Chen, Hongyang and Li, Zhao},
  booktitle={Proceedings of the Annual Meeting of the Association for Computational Linguistics},
  year={2026}
}

This codebase is built on torchtitan. Please also cite:

@inproceedings{liang2025torchtitan,
  title={TorchTitan: One-stop PyTorch native solution for production ready {LLM} pretraining},
  author={Liang, Wanchao and Liu, Tianyu and Wright, Less and Constable, Will and Gu, Andrew and Huang, Chien-Chin and Zhang, Iris and Feng, Wei and Huang, Howard and Wang, Junjie and Purandare, Sanket and Nadathur, Gokul and Idreos, Stratos},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2025},
  url={https://openreview.net/forum?id=SFN6Wm7YBI}
}

---

License

This project is licensed under the BSD 3-Clause License, consistent with the upstream torchtitan project.

The original torchtitan code is Copyright (c) Meta Platforms, Inc. and affiliates. Our modifications and additions are released under the same license.

---

Acknowledgments

We thank the PyTorch team for the excellent torchtitan framework, which provides a clean and scalable foundation for LLM training experiments.