README.md

---
language: en
license: mit
tags:
- qec
- surface-code
- quantum
- pytorch
- quantum-error-correction
- neural-decoder
- bivariate-bicycle
- ldpc
- transfer-learning
- foundation-decoder
pipeline_tag: other
---

# NTU Neural Decoder Checkpoints

Pre-trained model weights for the neural decoders introduced in **Neural
Transfer Unification (NTU)**, an architecture-agnostic transfer-learning
framework for scalable quantum error correction.

📄 **Paper**: *Efficient Foundation Decoders for Fault-Tolerant Quantum Computing*

🌐 **Project page**: [https://grahamyan.github.io/ntu-decoder/](https://grahamyan.github.io/ntu-decoder/)

💻 **Source code**: [https://github.com/GrahamYan/ntu-decoder](https://github.com/GrahamYan/ntu-decoder)

---

## Overview

NTU exploits the algebraic scale invariance of structured QEC code families to
transfer error knowledge from small codes to large-scale fault-tolerant regimes,
eliminating the cold-start optimization barrier.  The framework is instantiated
with two backbone architectures:

| Backbone | Description | Code families |
|---|---|---|
| **NTU-Transformer** | Interleaved RNN-Transformer with 2D RoPE and cross-attention readout | Surface, BB |
| **NTU-Neural-BP** | Graph-neural belief propagation on the code Tanner graph | BB |

For planar surface codes under circuit-level depolarizing noise, NTU-Transformer
surpasses standard PyMatching at *d* = 25 within a ~10³ GPU-hour training budget.
For the [[72, 12, 6]] bivariate-bicycle (BB) code, it outperforms BP+OSD across
all tested physical error rates and is competitive with multi-stage Relay BP.
Transfer from [[72, 12, 6]] to [[144, 12, 12]] reaches 93.1% block accuracy
within 2,500 steps (NTU-Transformer) and 95.3% within 500 steps (NTU-Neural-BP).

---

## Repository Structure

```
ntu-surface-code-decoder/
├── README.md
├── surface/                         ← Surface code (NTU-Transformer)
│   ├── d7.pth                       (121 MB, trained from scratch)
│   ├── d11.pth                      (121 MB, transferred from d=7)
│   ├── d15.pth                      (121 MB, transferred from d=11)
│   ├── d19.pth                      (121 MB, transferred from d=15)
│   ├── d23.pth                      (121 MB, transferred from d=19)
│   └── d25.pth                      (122 MB, transferred from d=23)
└── bb/                              ← BB code checkpoints
    ├── bb72_transformer.pt          (138 MB, NTU-Transformer, [[72,12,6]])
    └── neural_bp_bb72.pt            (1.2 MB, NTU-Neural-BP,  [[72,12,6]])
```

Each surface code checkpoint contains `model_state` (OrderedDict of weights),
`d` (code distance), `rounds` (syndrome extraction rounds), and `step`
(training step).  BB Transformer checkpoints additionally include `block_acc`
and `output_convention` metadata.  NTU-Neural-BP checkpoints store the raw
`state_dict` directly.

---

## Usage

### NTU-Transformer — Surface code

```python
import torch
from huggingface_hub import hf_hub_download

ckpt_path = hf_hub_download(
    repo_id="Dreamworldsmile/ntu-surface-code-decoder",
    filename="surface/d7.pth",
)

ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model.load_state_dict(
    {k.replace("_orig_mod.", "").replace("module.", ""): v
     for k, v in ckpt["model_state"].items()},
    strict=False,
)
```

### NTU-Transformer — BB code

```python
ckpt_path = hf_hub_download(
    repo_id="Dreamworldsmile/ntu-surface-code-decoder",
    filename="bb/bb72_transformer.pt",
)

ckpt = torch.load(ckpt_path, map_location="cpu")
state_dict = {k.replace("_orig_mod.", "").replace("module.", ""): v
              for k, v in ckpt["model_state"].items()}
# Filter to keys present in the target model (skip logical_readout_bias).
model_sd = model.state_dict()
filtered = {k: v for k, v in state_dict.items()
            if k in model_sd and model_sd[k].shape == v.shape
            and k != "logical_readout_bias"}
model.load_state_dict(filtered, strict=False)
```

### NTU-Neural-BP — BB code

```python
ckpt_path = hf_hub_download(
    repo_id="Dreamworldsmile/ntu-surface-code-decoder",
    filename="bb/neural_bp_bb72.pt",
)

ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
state_dict = {k.replace("module.", ""): v for k, v in ckpt.items()}
model.load_state_dict(state_dict, strict=True)
```

### Inference with the official code

```bash
git clone https://github.com/GrahamYan/ntu-decoder.git
cd ntu-decoder

# Surface code.
bash inference.sh --code surface --d 7 \
    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000

# BB code — NTU-Transformer.
bash inference.sh --code bb --model transformer --block_size 72 \
    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000 --p 0.005

# BB code — NTU-Neural-BP.
bash inference.sh --code bb --model neural_bp --block_size 72 \
    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000 --p 0.005
```

---

## Model Architecture

### NTU-Transformer

The Transformer-based decoder combines standard multi-head self-attention
blocks with two QEC-specific components:

- **Scalable STEM embedding** — Encodes syndrome data from variable-size
  lattices into a shared latent representation, absorbing the distance
  dependence into the input encoding (Eq. 2 in the paper).
- **QEC-aware 2D RoPE** — Applies rotary position embeddings to relative
  algebraic displacements defined by the code's shift set *M*(*x*, *y*, *t*; *d*),
  preserving detector geometry across code distances.
- **Interleaved RNN-Transformer backbone** — 5 GRU-based recurrent blocks
  alternating with 6 spatial self-attention blocks.
- **Cross-attention logical readout** — Learnable logical query tokens attend
  over the encoded detector representations.

### NTU-Neural-BP

A graph-neural-network decoder operating on the bipartite Tanner graph:

- **Message passing** between variable and check nodes with gated recurrent
  units (GRU) for message updates.
- **Syndrome-aware encoding** of check node states and prior LLRs.
- **Focal loss** with syndrome consistency regularization.
- Compact model (~300K parameters for the [[72, 12, 6]] code).

---

## Authors

[Ge Yan](https://grahamyan.github.io)<sup>1</sup>,
Shanchuan Li<sup>1, 2</sup>,
Shiyi Xiao<sup>1, 3</sup>,
Pengyue Ma<sup>1</sup>,
Hanyan Cao<sup>4</sup>,
[Feng Pan](https://scholar.google.com/citations?user=Vp6hFhUAAAAJ)<sup>4,\*</sup>,
[Yuxuan Du](https://yuxuan-du.github.io)<sup>1,\*</sup>

<sup>1</sup> College of Computing and Data Science, Nanyang Technological University, Singapore<br>
<sup>2</sup> Department of Electrical Engineering and Computer Science, Tokyo University of Agriculture and Technology, Japan<br>
<sup>3</sup> School of Artificial Intelligence, Shanghai Jiao Tong University, China<br>
<sup>4</sup> Science, Mathematics and Technology Cluster, Singapore University of Technology and Design, Singapore

<small><sup>\*</sup> Corresponding authors</small>

---

## Citation

```bibtex
@article{ntu2026,
  title={Efficient Foundation Decoders for Fault-Tolerant Quantum Computing},
  author={Yan, Ge and Li, Shanchuan and Xiao, Shiyi and Ma, Pengyue and
          Cao, Hanyan and Pan, Feng and Du, Yuxuan},
  year={2026},
}
```

---

## License

This repository is released under the [MIT License](https://opensource.org/licenses/MIT).