Upload README.md with huggingface_hub

README.md

@@ -2,66 +2,104 @@
 language: en
 license: mit
 tags:
-  - qec
-  - surface-code
-  - quantum
-  - pytorch
-  - quantum-error-correction
-  - neural-decoder
+- qec
+- surface-code
+- quantum
+- pytorch
+- quantum-error-correction
+- neural-decoder
+- bivariate-bicycle
+- ldpc
 pipeline_tag: other
 ---
 
-# NTU Surface Code Decoder (AlphaQubit V2)
+# NTU Neural Decoder Checkpoints
 
-Pre-trained neural decoder checkpoints for rotated surface codes, based on the
-**Neural Transfer Unification (NTU)** framework.
+Pre-trained neural decoder model weights for quantum error correction (QEC)
+codes, based on the **Neural Transfer Unification (NTU)** framework introduced
+in the accompanying paper.
 
 📄 **Paper**: *Transfer Learning is All You Need for Scalable Neural Decoder*
+🌐 **Project page**: [https://grahamyan.github.io/ntu-decoder/](https://grahamyan.github.io/ntu-decoder/)
 
-## Model Architecture
+---
 
-**AlphaQubit V2** — A high-capacity neural decoder (~58M parameters) featuring:
+## Overview
 
-- **Interleaved RNN-Transformer backbone** (5 GRU + 6 self-attention layers)
-- **2D Rotary Position Embedding (RoPE)** based on physical detector coordinates
-- **Joint X+Z stabilizer processing** with spatial hint connections
-- **Cross-attention readout** with learnable logical query tokens
-- Trained with **progressive knowledge distillation** from MWPM pseudo-labels
+This repository hosts the official model checkpoints for two families of QEC
+codes:
+
+| Code family | Architecture | Decoder |
+|---|---|---|
+| Rotated surface code | AlphaQubit V2 (~58M parameters) | Transformer-based |
+| Bivariate-bicycle (BB) code | AlphaQubitV2_BB (~XXM parameters) | Transformer-based |
+| Bivariate-bicycle (BB) code | Neural Belief Propagation | GNN-based message passing |
+
+All models are implemented in PyTorch and trained with distributed data-parallel
+(DDP) across 8 GPUs. The surface code decoder uses progressive knowledge
+distillation from minimum-weight perfect matching (MWPM) pseudo-labels;
+the BB decoder is trained end-to-end on sampled syndromes.
+
+---
 
 ## Repository Structure
 
 ```
 ntu-surface-code-decoder/
 ├── README.md
-├── surface/           ← Surface code checkpoints (AlphaQubit V2)
-│   ├── d7.pth         (~121 MB, scratch)
-│   ├── d11.pth        (~121 MB, transfer from d7)
-│   ├── d15.pth        (~121 MB, transfer from d11)
-│   ├── d19.pth        (~121 MB, transfer from d15)
-│   ├── d23.pth        (~121 MB, transfer from d19)
-│   └── d25.pth        (~122 MB, transfer from d23)
-└── bb/                ← BB code checkpoints (coming soon)
+├── surface/                      ← Surface code checkpoints (AlphaQubit V2)
+│   ├── d7.pth                    (121 MB, trained from scratch)
+│   ├── d11.pth                   (121 MB, transfer learning from d=7)
+│   ├── d15.pth                   (121 MB, transfer learning from d=11)
+│   ├── d19.pth                   (121 MB, transfer learning from d=15)
+│   ├── d23.pth                   (121 MB, transfer learning from d=19)
+│   └── d25.pth                   (122 MB, transfer learning from d=23)
+└── bb/                           ← BB code checkpoints
+    ├── bb72_transformer.pt       (138 MB, AlphaQubitV2_BB, [[72,12,6]] code)
+    └── neural_bp_bb72.pt         (1.2 MB, Neural-BP, [[72,12,6]] code)
 ```
 
-Each checkpoint contains:
-- `model_state` — OrderedDict of model weights
-- `d` — code distance (int)
-- `rounds` — decoding rounds (int)
-- `step` — training step (int)
+### Checkpoint format
+
+**Surface code checkpoints** (`surface/*.pth`):
+| Key | Type | Description |
+|---|---|---|
+| `model_state` | `OrderedDict` | Model weights (strip `_orig_mod.` and `module.` prefixes before loading) |
+| `d` | `int` | Code distance |
+| `rounds` | `int` | Syndrome extraction rounds |
+| `step` | `int` | Training step at which the checkpoint was saved |
+
+**BB Transformer checkpoints** (`bb/bb*_transformer.pt`):
+| Key | Type | Description |
+|---|---|---|
+| `model_state` | `OrderedDict` | Model weights |
+| `step` | `int` | Training step |
+| `block_acc` | `float` | Block accuracy at save time |
+| `per_log_mean` | `float` | Per-logical average accuracy |
+| `output_convention` | `dict` | Logical observable convention metadata |
+
+**Neural-BP checkpoints** (`bb/neural_bp_*.pt`):
+| Key | Type | Description |
+|---|---|---|
+| (raw `state_dict`) | `OrderedDict` | Model weights (strip `module.` prefix before loading) |
+
+---
 
 ## Usage
 
+### Surface code — AlphaQubit V2
+
 ```python
 import torch
 from huggingface_hub import hf_hub_download
 
-# Download a surface code checkpoint
+# Download a surface code checkpoint.
 ckpt_path = hf_hub_download(
     repo_id="Dreamworldsmile/ntu-surface-code-decoder",
     filename="surface/d7.pth",
 )
 
-# Load into AlphaQubit V2
+# Load into an AlphaQubit V2 model instance.
 ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
 model.load_state_dict(
     {k.replace("_orig_mod.", "").replace("module.", ""): v
@@ -70,30 +108,123 @@ model.load_state_dict(
 )
 ```
 
-### With the official code
+### BB code — AlphaQubitV2_BB (Transformer)
+
+```python
+import torch
+from huggingface_hub import hf_hub_download
+
+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 = ckpt["model_state"]
+state_dict = {k.replace("_orig_mod.", "").replace("module.", ""): v
+              for k, v in state_dict.items()}
+# Filter to keys present in the 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)
+```
+
+### BB code — Neural Belief Propagation
+
+```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
+
+The [official implementation](https://github.com/GrahamYan/ntu-decoder) provides a
+unified inference launcher that automatically downloads the required checkpoint:
 
 ```bash
-# Inference — auto-downloads surface/d{d}.pth
-python inference.py --hf_repo Dreamworldsmile/ntu-surface-code-decoder --d 7 --shots 100000
+# Surface code inference.
+bash inference.sh --code surface --d 7 \
+    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000
+
+# BB Transformer inference.
+bash inference.sh --code bb --model transformer --block_size 72 \
+    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000 --p 0.005
 
-# Transfer learning — specify full path within the repo
-bash train.sh --mode transfer \
-    --hf_ckpt Dreamworldsmile/ntu-surface-code-decoder/surface/d7.pth --d 11 ...
+# BB Neural-BP inference.
+bash inference.sh --code bb --model neural_bp --block_size 72 \
+    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000 --p 0.005
 ```
 
+For training and baseline evaluations, please refer to the shell scripts under
+`codes/Surface/` and `codes/BB/` in the source repository.
+
+---
+
+## Model Architecture
+
+### AlphaQubit V2 / AlphaQubitV2_BB
+
+A high-capacity neural decoder featuring:
+
+- **Interleaved RNN-Transformer backbone** (5 GRU + 6 self-attention layers)
+- **2D Rotary Position Embedding (RoPE)** based on physical detector coordinates
+- **Joint X+Z stabilizer processing** with spatial hint connections between
+  same-type and cross-type stabilizers
+- **Cross-attention readout** with learnable logical query tokens
+- Trained with **progressive knowledge distillation** from MWPM pseudo-labels
+  (surface code) or end-to-end on sampled syndromes (BB code)
+
+### Neural Belief Propagation
+
+A graph-neural-network decoder operating on the Tanner graph of the code:
+
+- **Bipartite message passing** between variable and check nodes
+- **Gated recurrent units (GRU)** for message updates
+- **Focal loss** with syndrome consistency regularization
+- Compact model size (~300K parameters for BB72)
+
+---
+
 ## Authors
 
-Ge Yan, Shanchuan Li, **Shiyi Xiao**, Pengyue Ma, Hanyan Cao, Feng Pan, Yuxuan Du
+Ge Yan<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<sup>4,\*</sup>,
+Yuxuan Du<sup>1,\*</sup>
+
+<sup>1</sup> Nanyang Technological University &nbsp;
+<sup>2</sup> Tokyo University of Agriculture and Technology &nbsp;
+<sup>3</sup> Shanghai Jiao Tong University &nbsp;
+<sup>4</sup> Singapore University of Technology and Design
 
-*Nanyang Technological University · TUAT · Shanghai Jiao Tong University · SUTD*
+<small><sup>\*</sup> Corresponding authors</small>
+
+---
 
 ## Citation
 
+If you use these model weights or the NTU framework in your research, please
+cite the accompanying paper:
+
 ```bibtex
 @article{ntu2026,
   title={Transfer Learning is All You Need for Scalable Neural Decoder},
   author={Yan, Ge and Li, Shanchuan and Xiao, Shiyi and Ma, Pengyue and
           Cao, Hanyan and Pan, Feng and Du, Yuxuan},
+  journal={arXiv preprint},
   year={2026},
 }
 ```
+
+---
+
+## License
+
+This repository is released under the [MIT License](https://opensource.org/licenses/MIT).