BetaQubit: Distilling Transformer Decoders for Low-Latency Surface Code Decoding

Saeng-nil Natthakan | Nanyang Technological University | College of Computing and Data Science | 2025

Supervisor: Asst Prof Tan Yong Kiam | Mentor: Dr. Derek Khu

Final Year Project (CCDS25-0718) submitted in partial fulfilment of the requirements for the Degree of Bachelor of Engineering in Computer Engineering. This work has been submitted to the International Quantum Error Correction Conference (QEC 2026) for the Poster track.

Abstract

Transformer-based decoders such as AlphaQubit improve surface code decoding accuracy by learning from experimental data, but their inference latency (~24 us/round) is too high for practical real-time QEC cycles. BetaQubit investigates knowledge distillation to compress an AlphaQubit-style Teacher decoder (1.1M parameters) into a compact Student model (180K parameters, 84% fewer), improving the accuracy-latency trade-off.

The Teacher is pretrained on 2 billion device-independent SI1000 samples, then finetuned on experimental data from Google's Sycamore (72 qubits) and Willow (105 qubits) processors. The Student is trained via a composite distillation loss combining hidden-state alignment, soft-label transfer, and hard-label supervision.

Key results (distance-3, Z-basis):

Dataset	Teacher LER	Distilled Student LER	Speedup vs Teacher	Speedup vs AlphaQubit
Sycamore	3.889 x 10^-2	3.762 x 10^-2	~2.5x	~4.5x
Willow	5.383 x 10^-3	5.747 x 10^-3	~2.5x	~4.5x

On Sycamore, the distilled Student actually outperforms the Teacher, likely due to a regularization effect where reduced capacity suppresses overfitting to noise fluctuations in the limited experimental data.

Model Architectures

The architecture is a recurrent Syndrome Transformer that processes stabilizer measurements round-by-round. Each layer consists of multi-head self-attention, a GELU-gated dense block, and a scattering residual convolution block that maps stabilizers to/from a 2D grid matching the physical qubit layout.

Model	Role	d_model	Layers	FFN dim	Conv dim	Params (d=3)
Original AlphaQubit	Reference	320	3	1600 (5x)	160	~8.3M
CompactGoogleDecoder	Teacher	128	2	384 (3x)	96	~1.1M
MicroGoogleDecoder	Student	64	1	192 (3x)	48	~180K
NanoGoogleDecoder	(experimental)	64	1	192 (3x)	48	~100K

Decoding Latency

Measured on NVIDIA A100 80GB PCIe, PyTorch float32, batch size 1, 25 rounds/sample, with torch.compile (reduce-overhead mode):

Model	d=3	d=7	d=13	d=25
Original AlphaQubit (8.3M)	505 us	527 us	639 us	964 us
Teacher (1.1M)	272 us	356 us	387 us	453 us
Student (180K)	127 us	141 us	159 us	212 us

Repository Layout

src/
  models/
    google_decoder.py             # Full decoder (RNN-style transformer + spatial convolutions)
    compact_google_decoder.py     # ~4x smaller student for distillation
    micro_google_decoder.py       # Micro (~180K) and Nano (~100K) variants
  data/
    google_data_utils.py          # STABILIZER_LOCATIONS per distance
    layout.py                     # Stabilizer grid utilities
    si1000_dem.py                 # SI1000 DEM-sampled dataset
    sycamore.py                   # Sycamore experimental dataset loaders
    pij_dem.py                    # Sycamore pij-derived DEM dataset
    willow.py                     # Willow experimental dataset loaders
  generators/
    generate_si1000_circuit.py    # Stim circuit generation for SI1000
  training/
    train_si1000_dem.py           # Pretrain on SI1000 DEM data
    finetune_pij.py               # Finetune on Sycamore pij experiments
    finetune_willow.py            # Finetune on Willow experiments
    distill_pretrain.py           # Distillation during pretraining
    distill_finetune.py           # Distillation during finetuning
    distill_pij_dem.py            # Distillation on pij DEM data
    distill_pij_finetune.py       # Distillation on Sycamore finetuning
    eval_checkpoints.py           # Batch checkpoint evaluation
    common.py                     # Noise curriculum, LR schedule, constants
    distributed.py                # DDP helpers
    ema.py                        # Exponential moving average
    checkpoint.py                 # Checkpoint save/load
    ler.py                        # Logical error rate fitting
    seed.py                       # Reproducibility
  benchmarks/
    benchmark_latency.py          # Inference latency profiling
    benchmark_compare_models.py   # Compare architectures (LER + throughput)
    benchmark_compare_general.py  # General model comparison
    benchmark_finetune_pij.py     # Evaluate finetuned Sycamore models
    benchmark_finetune_willow.py  # Evaluate finetuned Willow models
    plot_decoding_speed.py        # Decoding speed plots
    plot_compare_models_from_results.py  # Plot from saved results
  utils/
    model_param_counts.py         # Parameter counting utility

scripts/                          # Shell scripts for running experiments
  si1000/                         #   SI1000 pretraining & distillation
  pij_sycamore/                   #   Sycamore finetuning, distillation & benchmarks
  willow/                         #   Willow finetuning, distillation & benchmarks

jobs/                             # PBS cluster job scripts

checkpoints/                      # Pre-trained model weights (see below)

pretrain_si1000/                  # Pre-generated SI1000 DEM files

Setup

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Dependencies: PyTorch, Stim 1.15, NumPy, TensorBoard, TensorFlow (TFRecord parsing), pytorch-lamb, scipy, matplotlib, pymatching, PyYAML.

Data

Synthetic (SI1000)

Pre-generated DEM files are included under pretrain_si1000/. No download needed. The SI1000 noise model uses non-uniform Pauli error channels approximating superconducting hardware: measurement errors at 5p, reset at 2p, two-qubit gates at p, single-qubit/idle at p/10.

Sycamore Experiments

From Google's 72-qubit Sycamore processor. Four distance-3 regions + one distance-5 region, X/Z bases, 1--25 rounds, 50K shots each (6.5M total shots). Download from Zenodo. Place experiment folders under sycamore/, e.g. sycamore/surface_code_bZ_d3_r25_center_{row}_{col}. Each folder should contain:

• detection_events.b8
• measurements.b8
• obs_flips_actual.01
• properties.yml

Willow Experiments

From Google's 105-qubit Willow processor (higher fidelity than Sycamore). Five distance-3, three distance-5, and one distance-7 region, 1--250 rounds. Download from Zenodo. Place under willow/, same format as Sycamore.

Training Pipeline

1. Pretrain on SI1000 DEM Data

python -m src.training.train_si1000_dem --model_arch compact --distance 3 --basis Z

Pretrained on ~500M synthetic samples with noise curriculum (Gaussian-weighted sampling over scale factors 0.5--1.0), LAMB optimizer, EMA (decay 0.999), and step-decay LR schedule (factor 0.7).

2. Finetune on Experimental Data

# Sycamore (240 epochs, batch 1024)
python -m src.training.finetune_pij --model_arch compact \
  --pretrained_checkpoint path/to/pretrained.pth

# Willow (120 epochs, batch 256)
python -m src.training.finetune_willow --model_arch compact \
  --pretrained_checkpoint path/to/pretrained.pth

Uses relative weight decay (0.08) penalizing deviation from pretrained weights. Training split: 19,880 train / 5,120 dev, evaluated on 25,000 held-out shots.

3. Knowledge Distillation

# Pretrain distillation
python -m src.training.distill_pretrain \
  --teacher_arch compact --student_arch micro --dataset si1000

# Finetune distillation
python -m src.training.distill_finetune \
  --teacher_arch compact --student_arch micro --dataset willow

The distillation loss combines three terms:

L = a * MSE(f(h_S), h_T) + b * T^2 * BCE(z_S/T, sigma(z_T/T)) + g * BCE(z_S, y)

Stage	a (hidden-state)	b (soft-label)	g (hard-label)	Temperature
Pretraining	0.5	0.3	0.2	4
Finetuning	0.3	0.5	0.2	4

The shift toward higher soft-label weight during finetuning reflects the increased importance of transferring the Teacher's hardware-specific predictions.

See scripts/ for full experiment configurations.

Pre-trained Checkpoints

The checkpoints/ directory contains trained weights for distance-3, basis-Z models:

Directory	Description
pretrained_compact_si1000/	Compact (Teacher) pretrained on SI1000
pretrained_micro_si1000/	Micro (Student) pretrained on SI1000
finetuned_compact_pij/	Teacher finetuned on Sycamore
finetuned_compact_willow/	Teacher finetuned on Willow
finetuned_micro_pij/	Student finetuned on Sycamore
finetuned_micro_willow/	Student finetuned on Willow
distilled_micro_si1000/	Student distilled on SI1000
distilled_micro_sycamore/	Student distilled on Sycamore
distilled_micro_willow/	Student distilled on Willow

Evaluation

Performance is evaluated using the logical error rate per round (LER), estimated by fitting the decoder fidelity to the exponential decay model F(n) = (1 - 2e)^n. Lower LER indicates better decoding.

Full Results (distance-3, Z-basis)

Dataset	Pretrained SI1000	Teacher (finetuned)	Compact baseline	Compact distilled
Sycamore	8.983 x 10^-2	3.889 x 10^-2	4.128 x 10^-2	3.762 x 10^-2
Willow	5.600 x 10^-3	5.383 x 10^-3	5.974 x 10^-3	5.747 x 10^-3

Benchmarking

# Compare models on Sycamore
bash scripts/pij_sycamore/benchmark_sycamore.sh

# Compare models on Willow
bash scripts/willow/benchmark_willow.sh

# Inference latency
python -m src.benchmarks.benchmark_latency --checkpoint path/to/model.pth --device cuda

Notes

Stabilizer Geometry

Stabilizers are placed on a (2d+1) x (2d+1) grid using hard-coded STABILIZER_LOCATIONS from src/data/google_data_utils.py. All inputs (measurements, detection events, stabilizer indices) are aligned to flattened grid indices (y * grid_dim + x) so the decoder's spatial convolutions match the physical qubit layout.

Differences from AlphaQubit

• Reduced Teacher capacity: d_model 128 vs 320, 2 layers vs 3 (~1.1M vs ~8.3M params) due to compute constraints.
• Omitted spatial attention bias: found to contribute negligibly to decoding accuracy.
• The LER results here are not directly comparable to the original AlphaQubit LER of 2.9 x 10^-2 at d=3 due to the reduced Teacher size.

TensorBoard

Training logs are written to runs/:

tensorboard --logdir runs

Acknowledgements

Supervised by Asst Prof Tan Yong Kiam with mentorship from Dr. Derek Khu.

Computing resources provided by A*STAR (NVIDIA A100 GPU) and the National Supercomputing Centre (NSCC) Singapore (ASPIRE 2A cluster).

Based on the architecture described in:

Bausch et al., "Learning high-accuracy error decoding for quantum processors," Nature 635, 834--840 (2024).