Document BB72 checkpoints for ultra-low p and high p

Adds the new low-p fine-tuned BB72 checkpoint to the repository structure and usage example, while documenting that the original checkpoint remains recommended for high-p evaluation and transfer learning.

README.md

@@ -17,9 +17,7 @@ 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.
+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*
 
@@ -31,47 +29,46 @@ framework for scalable quantum error correction.
 
 ## 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:
+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).
+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
 
-```
+```text
 ntu-surface-code-decoder/
 ├── README.md
-├── surface/                         ← Surface code (NTU-Transformer)
+├── 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]])
+└── bb/                               ← BB code checkpoints
+    ├── bb72_transformer.pt           (144 MB, high-p evaluation and transfer learning)
+    ├── bb72_transformer_lowp_finetuned.pt
+    │                                 (144 MB, fine-tuned for ultra-low physical error rates)
+    └── 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.
+### BB72 Transformer checkpoint selection
+
+| Checkpoint | Recommended use |
+|---|---|
+| `bb/bb72_transformer.pt` | **High physical error rates (high p)** and **transfer learning** |
+| `bb/bb72_transformer_lowp_finetuned.pt` | **Ultra-low physical error rates (ultra-low p)** |
+
+The original `bb72_transformer.pt` checkpoint is retained unchanged. Use the low-p fine-tuned checkpoint only for evaluation or inference in the ultra-low-p regime.
+
+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.
 
 ---
 
@@ -87,7 +84,6 @@ 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
@@ -98,20 +94,31 @@ model.load_state_dict(
 
 ### NTU-Transformer — BB code
 
+Select the checkpoint according to the target physical error-rate regime:
+
 ```python
+# Use bb72_transformer.pt for high p and transfer learning.
+# Use bb72_transformer_lowp_finetuned.pt for ultra-low p.
+checkpoint_name = "bb72_transformer_lowp_finetuned.pt"
+
 ckpt_path = hf_hub_download(
     repo_id="Dreamworldsmile/ntu-surface-code-decoder",
-    filename="bb/bb72_transformer.pt",
+    filename=f"bb/{checkpoint_name}",
 )
-
 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()}
+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"}
+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)
 ```
 
@@ -122,7 +129,6 @@ 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)
@@ -136,15 +142,15 @@ cd ntu-decoder
 
 # Surface code.
 bash inference.sh --code surface --d 7 \
-    --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000
+  --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
+  --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
+  --hf_repo Dreamworldsmile/ntu-surface-code-decoder --shots 100000 --p 0.005
 ```
 
 ---
@@ -153,48 +159,33 @@ bash inference.sh --code bb --model neural_bp --block_size 72 \
 
 ### NTU-Transformer
 
-The Transformer-based decoder combines standard multi-head self-attention
-blocks with two QEC-specific components:
+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.
+- **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.
+- **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).
+- 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
+[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>
 
-<small><sup>\*</sup> Corresponding authors</small>
+<sup>1</sup> College of Computing and Data Science, Nanyang Technological University, Singapore  
+<sup>2</sup> Department of Electrical Engineering and Computer Science, Tokyo University of Agriculture and Technology, Japan  
+<sup>3</sup> School of Artificial Intelligence, Shanghai Jiao Tong University, China  
+<sup>4</sup> Science, Mathematics and Technology Cluster, Singapore University of Technology and Design, Singapore  
+<sup>\*</sup> Corresponding authors
 
 ---
 
@@ -203,8 +194,7 @@ Hanyan Cao<sup>4</sup>,
 ```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},
+  author={Yan, Ge and Li, Shanchuan and Xiao, Shiyi and Ma, Pengyue and Cao, Hanyan and Pan, Feng and Du, Yuxuan},
   year={2026},
 }
 ```