Skip to content

Repository files navigation

BiKAN — Restoring the Collapsed Basis of Binary Kolmogorov–Arnold Networks with Walsh Parities

Paper: https://arxiv.org/abs/2608.01490

Everything needed to reproduce every number in the main paper and the technical appendix is in this repository.

Read this file to run one experiment. Read REPRODUCE.md to reproduce the whole paper end-to-end.


Contents

  1. What this code does
  2. Installation
  3. Datasets
  4. Five-minute smoke test
  5. Running a single experiment
  6. Every table and figure → exact command
  7. Reading the output artifacts
  8. Training configuration reference
  9. Repository layout
  10. Tests
  11. Troubleshooting

1. What this code does

BiKAN is a width-free recipe for binarizing Kolmogorov–Arnold Networks at W1A1 (1-bit weights, 1-bit activations). The central claim is that the standard "binarize then widen 4×" fix is unnecessary if you instead restore the basis capacity that binarization destroys, using a degree-2 Walsh parity path that costs zero MACs and zero hidden channels.

The student stays at the FP32 teacher's widths (64, 128, 256). Five mechanisms combine (bkan/layers/bireal.py):

Mechanism Switch Hardware cost at W1A1
Degree-2 Walsh parity path (ours) parity_rolls=(1,3) XNOR gates + fixed channel rotation (wiring)
Real-valued Bi-Real shortcut shortcut=True integer add + per-channel scale
RPReLU (ReActNet) rprelu=True 2 adds + 1 per-channel multiply (po2-snappable to a shift)
Libra-PB + EDE weight binarization (IR-Net) estimator='ede' training-time only
ApproxSign activations with trainable shifts estimator='ede' training-time only

Setting all of them off (estimator='analytic', parity_rolls=(), shortcut=False, rprelu=False) recovers the bare baseline — the conventional binary KAN layer that motivates widening.

Everything that reaches hardware is unchanged XNOR-popcount arithmetic. The recovery machinery is training-time only.


2. Installation

Python ≥ 3.9, PyTorch ≥ 2.0. A single modern GPU is enough for any individual experiment; see REPRODUCE.md for the full-campaign budget.

# from the repository root
python -m venv .venv && source .venv/bin/activate     # Windows: .venv\Scripts\activate
pip install --upgrade pip
pip install -e ".[all,dev]"

[all] pulls in the extras used by the appendix experiments (scikit-learn, pandas, ucimlrepo, openpyxl, sympy, pyyaml, matplotlib, tqdm). [dev] adds pytest.

Minimal install — if you only want the main-paper CIFAR/MNIST results, core dependencies suffice:

pip install -e .            # torch, torchvision, einops
pip install scipy           # optional: exact p-values in analyze_runs.py

Optional extras by experiment block:

You want to run Also install
Tabular / MLP suite (App. I) pip install -e ".[data]"
PyKAN cross-family row (App. J) pip install pykan (a vendored copy also ships in pykan/)
Exact paired-test p-values pip install scipy
The quickstart notebook pip install -e ".[notebooks]"

Verify the install:

python -c "import torch, bkan; print(torch.__version__, torch.cuda.is_available())"
python main.py --list          # prints every registered experiment + its defaults

3. Datasets

All datasets download automatically to ./data/ on first use. No manual preparation is needed.

Dataset Size Source Used for
MNIST 12 MB torchvision Main results, mechanism ablation, FPGA eval
CIFAR-10 170 MB torchvision All primary mechanism experiments
CIFAR-100 170 MB torchvision Main results
Tiny ImageNet ≈ 240 MB direct download (bkan/data/extra.py) Stretch row
JSC (hls4ml jet substructure) ≈ 3 MB OpenML data_id=42468 Tabular suite
Wine, Dry Bean < 1 MB scikit-learn / UCI Tabular suite
Traffic (California) ≈ 10 MB UCI Tabular regression row

Set a different root by editing bkan/data/loaders.py ('./data'), or symlink ./data at an existing dataset cache.

Dataloader workers default to 2. On a many-core node, raise them once per machine and keep the value fixed for the whole campaign:

export BKAN_NUM_WORKERS=8

4. Five-minute smoke test

Run this first. It exercises seeding, dataloading, the validation split, every mechanism arm, checkpointing, and the artifact writer at 1 epoch each.

bash experiments/run_paper_suite.sh smoke

Or the minimum end-to-end check (≈ 2 minutes on a GPU, ≈ 10 on CPU):

python run_seeded.py -e bireal_ablation_cifar10 --seed 0 --name smoke \
    --set "arms=['full','bare']" teacher_epochs=1 epochs=1 val_fraction=0.1

Expected on success:

  • runs/smoke__s0/seed_meta.json — seed + overrides + timestamps
  • runs/smoke__s0/bireal_teacher_cifar10.pth — the cached FP32 teacher
  • runs/smoke__s0/bireal_ablation_cifar10.jsonfull and bare rows with acc, acc_bn_recal, params_M
  • runs/smoke__s0/bireal_ablation_cifar10_log_epochs.csv — per-epoch table

Accuracies after 1 epoch are meaningless (~30–45 %); what matters is that all four artifacts exist and that full.params_M == 11.94 and bare.params_M == 5.97.

Then confirm the analysis path works:

python analyze_runs.py --root runs --tag smoke

5. Running a single experiment

There are three entry points. They are equivalent; pick by what you need.

A. run_seeded.pyuse this for anything reported in the paper

It seeds the process before any model is built, writes to a collision-free run directory runs/<name>__s<seed>/, and can reuse a cached teacher so that compared arms share one teacher (the paired-seed protocol).

python run_seeded.py -e <experiment> --seed <n> --name <tag> \
    [--set KEY=VALUE ...] [--teacher <path/to/teacher.pth>]
# Train the seed-0 teacher implicitly, with the Table-2 'full' arm:
python run_seeded.py -e bireal_ablation_cifar10 --seed 0 --name abl_full \
    --set "arms=['full']" val_fraction=0.1

# Every later arm of seed 0 reuses that teacher — no retraining, exact pairing:
python run_seeded.py -e bireal_ablation_cifar10 --seed 0 --name abl_no_parity \
    --set "arms=['no_parity']" val_fraction=0.1 \
    --teacher runs/abl_full__s0/bireal_teacher_cifar10.pth

--set values are parsed with ast.literal_eval, so quote anything with brackets or parentheses: --set "arms=['full']", --set "parity_rolls=(1,3)".

B. main.py — quick single runs with the published defaults

No seeding, artifacts go to runs/ under the canonical experiment name. Fine for exploration, not for numbers you intend to report.

python main.py --list
python main.py --experiment bireal_progressive_cifar10
python main.py --experiment bireal_progressive_cifar10 --set lambda_at=0

C. From Python

import bkan
from bkan.pipelines.bireal_w1a1 import train_bireal_progressive

bkan.set_seed(0)
results = train_bireal_progressive(dataset="cifar10", val_fraction=0.1,
                                   save_dir="runs/my_run__s0")
print(results)   # {'teacher_fp32': ..., 'w1a8': ..., 'w1a1_final_bn_recal': ...}

Build a model directly:

from bkan.models import build_model
student = build_model("bireal_kagn", a_bits=1, groups=8,
                      parity_rolls=(1, 3), shortcut=True,
                      estimator="ede", rprelu=True,
                      in_channels=3, num_classes=10)

The experiments that matter for this paper

Experiment name What it trains
bireal_progressive_cifar10 Teacher → A8 → A4 → A2 → A1 at teacher width (the headline recipe)
bireal_progressive_mnist Same, MNIST
bireal_progressive_cifar100 Same, CIFAR-100
bireal_ablation_cifar10 Direct-to-W1A1 mechanism arms: full, no_shortcut, no_parity, no_rprelu, frozen_ste, bare
bireal_ablation_mnist Same arms, MNIST
bireal_width_sweep_cifar10 Full stack at 1× / 2× / 4× widths, plus bare at 4×
bireal_family_{efficientkan,fastkan,pykan} Width-free dense student per KAN family (MNIST)
bireal_tabular_{jsc,traffic,suite} Tabular / MLP rows
fastkan_phase6c_pure_bikan_v2 Trains + exports bit-packed HLS test vectors
fastkan_phase7b_hls_export Generates, compiles and runs the XNOR-popcount HLS testbench

python main.py --list shows all of them plus the legacy four-family phase pipelines, which are retained because the cross-family and hardware sections build on them.


6. Every table and figure → exact command

All commands run from the repository root. <S> is the seed. Primary CIFAR-10 comparisons use seeds 0–4; secondary arms and CIFAR-100 use 0–2.

Everything in this section shares one teacher per (dataset, seed). Train it once — the first command below — and pass it to every subsequent arm with --teacher. This is not an optimization: paired statistics require it.

# ---- STEP 0: per-seed CIFAR-10 teacher (also the Table-2 'full' cell) ----
for S in 0 1 2 3 4; do
  python run_seeded.py -e bireal_ablation_cifar10 --seed $S --name abl_full \
      --set "arms=['full']" val_fraction=0.1
done
TEACHER=runs/abl_full__s0/bireal_teacher_cifar10.pth   # per seed: abl_full__s$S/...

Main results — accuracy at W1A1

# ours, teacher width, progressive descent
for S in 0 1 2 3 4; do
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name prog_full \
      --set val_fraction=0.1 --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done

for S in 0 1 2; do
  python run_seeded.py -e bireal_progressive_mnist   --seed $S --name prog_mnist --set val_fraction=0.1
  python run_seeded.py -e bireal_progressive_cifar100 --seed $S --name prog_c100  --set val_fraction=0.1
done

# Tiny ImageNet stretch row (2 seeds, val_fraction=0 fallback — disclosed)
for S in 0 1; do
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name prog_tinyin \
      --set dataset=tiny_imagenet teacher_epochs=60 epochs_per_stage=50 val_fraction=0
done

Expected (w1a1_final_bn_recal, mean ± std): CIFAR-10 84.38 ± 0.14 (n=5), MNIST 99.48 ± 0.05 (n=3), CIFAR-100 55.81 ± 1.86 (n=3), Tiny ImageNet 37.15 ± 0.68 (n=2).

Mechanism ablation

for S in 0 1 2 3 4; do
  for ARM in no_parity bare; do
    python run_seeded.py -e bireal_ablation_cifar10 --seed $S --name abl_$ARM \
        --set "arms=['$ARM']" val_fraction=0.1 \
        --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  done
done
for S in 0 1 2; do
  for ARM in no_shortcut no_rprelu frozen_ste; do
    python run_seeded.py -e bireal_ablation_cifar10 --seed $S --name abl_$ARM \
        --set "arms=['$ARM']" val_fraction=0.1 \
        --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  done
done

Expected (<arm>.acc_bn_recal): full 83.00 ± 0.31, no_shortcut 83.08 ± 0.28, frozen_ste 83.05 ± 0.29, no_parity 81.77 ± 0.38, no_rprelu 78.93 ± 0.85, bare 78.19 ± 0.33.

Mechanism vs. width (H3) and equal-parameter control (H3′)

# the fair widened baseline: 'bare' at 4x width, our identical budget + KD
for S in 0 1 2 3 4; do
  python run_seeded.py -e bireal_width_sweep_cifar10 --seed $S --name widbare4 \
      --set "width_mults=()" include_bare_wide=True val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done

# width frontier: full stack at 2x and 4x
for S in 0 1 2; do
  python run_seeded.py -e bireal_width_sweep_cifar10 --seed $S --name widfull_x2 \
      --set "width_mults=(2.0,)" include_bare_wide=False val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done
python run_seeded.py -e bireal_width_sweep_cifar10 --seed 0 --name widfull_x4 \
    --set "width_mults=(4.0,)" include_bare_wide=False val_fraction=0.1 \
    --teacher runs/abl_full__s0/bireal_teacher_cifar10.pth

# equal-parameter control: widen 'bare' until it matches BiKAN's parameter count
for S in 0 1 2 3 4; do
  WB=$(python hpc/solve_bare_widths.py)
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name barematch \
      --set "a_bits_schedule=(1,)" lambda_ema=0 "widths=$WB" "parity_rolls=()" \
            shortcut=False rprelu=False estimator=analytic val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done

hpc/solve_bare_widths.py prints the widths tuple that makes the bare student parameter-matched to the parity student; it is called inline so the value is never hand-copied.

Expected: bare_x4.acc_bn_recal 82.68 ± 0.21 at ≈ 8× the parameters, barematch 79.87 ± 0.27 at matched parameters, full_x2 84.60 ± 0.27, full_x4 85.77 (n=1).

Parity dose–response and pairing controls

Mode-B isolation: direct-to-W1A1, EMA off, only |R| varies.

declare -A R=( [par_r0]="()" [par_r1]="(1,)" [par_r2]="(1,3)" \
               [par_r4]="(1,2,3,5)" [par_r8]="(1,2,3,4,5,6,7,8)" )
for TAG in par_r0 par_r1 par_r2 par_r4 par_r8; do
  for S in 0 1 2; do
    python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name $TAG \
        --set "a_bits_schedule=(1,)" lambda_ema=0 "parity_rolls=${R[$TAG]}" \
              val_fraction=0.1 \
        --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  done
done
# par_r2 additionally at seeds 3 and 4 (it is the H3' reference cell)

# controls: does the *choice* of offsets matter, or only their number?
for S in 0 1 2; do
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name par_off17 \
      --set "a_bits_schedule=(1,)" lambda_ema=0 "parity_rolls=(1,7)" val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name par_off211 \
      --set "a_bits_schedule=(1,)" lambda_ema=0 "parity_rolls=(2,11)" val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name par_rand \
      --set "a_bits_schedule=(1,)" lambda_ema=0 parity_pairing=random val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done

Expected (w1a1_final_bn_recal): |R| = 0 → 81.74 ± 0.53, 1 → 82.67 ± 0.10, 2 → 82.96 ± 0.12, 4 → 83.47 ± 0.37, 8 → 84.00 ± 0.08 — monotone, not yet saturated. Controls: par_off17 83.08 ± 0.38, par_off211 82.76 ± 0.08, par_rand 82.94 ± 0.10 (random pairing matches circulant rolls, but would need a stored routing table on hardware — it is a science control, not a deployment path).

Capacity starvation (H5)

for S in 0 1 2; do
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name w025_par \
      --set "a_bits_schedule=(1,)" lambda_ema=0 "widths=(16,32,64)" val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name w025_nopar \
      --set "a_bits_schedule=(1,)" lambda_ema=0 "widths=(16,32,64)" "parity_rolls=()" \
            val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done
# repeat with widths=(32,64,128) for the x0.5 pair (tags w05_par / w05_nopar)

Expected paired parity gap: +2.87 ± 0.67 at ×0.25, +1.66 ± 0.40 at ×0.5, +1.23 ± 0.43 at ×1 — the gap grows as capacity shrinks.

Progressive vs. direct (H4)

for S in 0 1 2 3 4; do   # equal final-stage budget (40 epochs at A1)
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name direct40 \
      --set "a_bits_schedule=(1,)" val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done
for S in 0 1 2; do       # total-compute matched (160 epochs at A1)
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name direct160 \
      --set "a_bits_schedule=(1,)" epochs_per_stage=160 val_fraction=0.1 \
      --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done

Expected: direct40 82.51 ± 0.25 (progressive wins by +1.87, p < 0.001), direct160 84.43 ± 0.29 (progressive advantage disappears, −0.08, p = 0.77). Both are reported in the paper.

Recipe ablation

One override per arm, on the full progressive recipe.

for S in 0 1 2; do
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name rec_noat   --set lambda_at=0        val_fraction=0.1 --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name rec_noema  --set lambda_ema=0       val_fraction=0.1 --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name rec_nodiv  --set lambda_div=0       val_fraction=0.1 --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name rec_nolat  --set latent_lr_mult=1.0 val_fraction=0.1 --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
  python run_seeded.py -e bireal_progressive_cifar10 --seed $S --name rec_noede  --set estimator=analytic val_fraction=0.1 --teacher runs/abl_full__s$S/bireal_teacher_cifar10.pth
done

Expected: all five land within ±0.4 of prog_full (84.38) — and all five are at or slightly above it. At this budget none of attention transfer, the EMA self-teacher, shift diversity, the latent-LR boost or EDE measurably improves the W1A1 result; they are second-order next to parity and progressive descent. Report them as such.

Synthetic degree-2 Walsh tasks (Proposition 1)

CPU-only, a few minutes total.

python synthetic_parity_task.py --regime covered   --rolls 1,3 --seeds 0 1 2 --out runs/synth_covered.json
python synthetic_parity_task.py --regime uncovered --rolls 1,3 --seeds 0 1 2 --out runs/synth_uncovered.json
python synthetic_parity_task.py --regime covered   --rolls 1,3 --binary-head --seeds 0 1 2 --out runs/synth_covered_binhead.json

Expected: parity reaches 100 % with 193 parameters when the target's pair distances lie inside the roll orbit, and sits at chance (≈ 50 %) when they do not; the FP linear baseline is at chance in both; two-layer sign-MLPs need ≈ 1.1 M parameters to reach ≈ 87 %. Reference outputs ship in results/synth_*.json.

Cross-family transfer

python main.py --experiment bireal_family_efficientkan     # student (784, 64, 10)
python main.py --experiment bireal_family_fastkan          # student (784, 128, 10)
python main.py --experiment bireal_family_pykan            # requires: pip install pykan
# widened frozen references from the same families:
python main.py --experiment pykan_phase6_w1a1_mnist
python main.py --experiment fastkan_phase6c_pure_bikan
python main.py --experiment efficientkan_train_teacher && \
python main.py --experiment efficientkan_phase6_8_w1a1_frontier

These rows were collected under the earlier single-seed protocol and are reported in the appendix as exploratory, not as primary statistical evidence.

Tabular and MLP suite

python main.py --experiment bireal_tabular_jsc            # hls4ml jet substructure
python main.py --experiment bireal_tabular_traffic        # regression (RMSE, 72 -> 96)
python main.py --experiment bireal_tabular_suite          # wine + dry_bean + jsc + traffic
python main.py --experiment bireal_tabular_tinyimagenet_mlp   # heavy, 12288-dim input

# faster JSC iteration (full set is ~830k rows):
python main.py --experiment bireal_tabular_jsc --set "loader_kwargs={'subsample':100000}"

Four arms per dataset under identical budgets: FP32 KAN teacher, ours_w1a1 (teacher dims), wide_w1a1 (bare config, hidden dims × 4), nonkan_w1a1 (plain binary MLP, use_basis=False, groups=1). Traffic additionally trains mlp_traffic as an FP32 non-KAN reference. Also single-seed / exploratory.

FPGA resources and latency

The complete hardware track ships in this repository: the exporters (bkan/export/), the generated Vivado HLS kernels and testbenches (hls/), the build scripts, and the post-place-and-route reports the paper's hardware table is read from. See hls/README.md for the full flow — this is a summary.

Nine designs target Xilinx Zynq-7020 (xc7z020clg400-1) at 100 MHz, evaluated on the MNIST 10k test set: GRAM dense (FP32 / W1A1 / W1A1+Po2-QAT), KAGN conv (FP32 / W1A1 / naive Po2), EfficientKAN (FP32 / binary), and a tiny _mini_full smoke design.

Five stages take a checkpoint to post-route numbers:

python main.py --experiment fp32_teacher_gram     # 1. FP32 teacher
python main.py --experiment w1a1_student_gram     # 2. W1A1 student
python main.py --experiment export_gram           # 3. fold BN, pack bits -> hw_data/gram/
python main.py --experiment verify_gram           # 4. numpy replay == PyTorch
python main.py --experiment hls_gram_generate     # 5. emit hls/gram/{src,tb}
python main.py --experiment hls_gram_csim         #    g++ compile + run over all 10k

Then, with Vivado HLS 2019.1:

cd hls/gram
vivado_hls -f run_hls.tcl      # csim -> csynth -> cosim -> export IP
vivado_hls -f run_impl.tcl     # place & route -> real LUT/FF/DSP/BRAM + Fmax

Substitute kagn_conv, efficientkan, gram_po2, kagn_conv_po2, gram_fp32, efficientkan_fp32, conv_fp32. The Po2 designs insert po2qat_<variant> before the export.

Correctness is verifiable without any Xilinx tooling. The kernels compile with plain g++ (the headers fall back to uint64_t when the ap_int types are unavailable), and the shipped _mini_full design runs immediately because its test vectors are included:

cd hls/_mini_full
g++ -O2 -std=c++11 -Isrc src/gram_hls.cpp tb/gram_tb.cpp -o mini_test
./mini_test ../../hw_data/_mini_full
# -> accumulator mismatches: 0 | argmax matches: 50/50 | bit-exact with the Python export

The export contract is checked by three CPU-only regression tests:

python -m pytest tests/additions/test_packing.py \
                 tests/additions/test_export_replay.py \
                 tests/additions/test_export_replay_conv.py -q

Only the resource and latency numbers need Vivado. To keep the package small, the raw Vivado reports and proj_*/ build trees are not shipped — the .tcl scripts regenerate them — but the numbers we measured are preserved in hls/RESULTS.json (post-place-and-route LUT/FF/DSP/BRAM, achieved Fmax, and C/RTL-cosim cycle counts for all nine designs).

Each hw_data/<design>/manifest.json records the geometry, packing convention, fixed-point scales, source-checkpoint SHA-256, and the replay-vs-PyTorch agreement (e.g. kagn_conv: "replay_vs_torch_argmax": [10000, 10000], replay_acc 99.39 vs deployed_torch_acc 99.39).

The older single-layer FastKAN HLS harness is still present (fastkan_phase7b_hls_export, fastkan_phase7b_export_mnist) and is what docs/bireal_hw_notes.md documents; the hls/ designs above superseded it.


7. Reading the output artifacts

Every bireal_* run writes incrementally, so the JSON on disk is complete up to the last finished epoch even if the job dies.

runs/<tag>__s<seed>/
  seed_meta.json                          seed, overrides, argv, timestamps
  bireal_teacher_<dataset>.pth            cached FP32 teacher
  bireal_w1a1_<dataset>.pth               final student checkpoint
  bireal_progressive_<dataset>.json       config + env + per-stage + results
  bireal_progressive_<dataset>_epochs.csv per-epoch metrics, flattened
  bireal_ablation_<dataset>.json          per-arm acc, acc_bn_recal, params_M
  bireal_ablation_<dataset>_log.json      full per-epoch record, per stage

Re-running never clobbers a previous log — the old one moves to *.prev_<timestamp>.json.

Use the _log.json, not the CSV, for per-epoch curves. The CSV writer fixes its header from the first row it sees, which is the teacher stage (epoch, train_loss, test_acc, lr), and silently drops extra columns from later stages. The student stages log more than that — val_acc, best_acc, ede_t, ema_active — and those survive only in the JSON:

import json
log = json.load(open("runs/prog_full__s0/bireal_progressive_cifar10.json"))
for stage, rec in log["stages"].items():
    # the teacher stage has no val_acc, so use .get
    print(stage, [e.get("val_acc") for e in rec["epochs"]])

The primary metric is w1a1_final_bn_recal (progressive runs) or <arm>.acc_bn_recal (ablation runs): BN-recalibrated test accuracy of the validation-selected checkpoint. Best-on-test values (w1a1, <arm>.acc) are secondary and appear in the appendix only.

Aggregate and test:

python analyze_runs.py --root runs                       # mean ± std over seeds, every tag
python analyze_runs.py --root runs --tag par_ --metric w1a1_final_bn_recal
python analyze_runs.py --root runs --pair prog_full direct40 --metric w1a1_final_bn_recal
python analyze_runs.py --root runs --pair abl_full abl_no_parity \
    --metrics full.acc_bn_recal no_parity.acc_bn_recal

--pair matches seeds and runs a two-sided paired t-test. Install scipy for exact p-values; without it the t statistic and df are printed.

Lost console output but still have checkpoints:

python experiments/recover_results.py     # -> runs/recovered_results.json

It infers each checkpoint's architecture from its state-dict keys and re-evaluates on the matching test set. Final accuracies are fully recoverable this way; per-epoch curves and intermediate A8/A4/A2 stage accuracies are not.


8. Training configuration reference

Defaults live in bkan/pipelines/bireal_w1a1.py and bkan/pipelines/bireal_ablation.py. Nothing below needs to be passed explicitly — this is what runs when you pass nothing.

Architecture

Teacher ConvolutionalKAGN, degree-3 Gram polynomial, widths (64, 128, 256), 1.86 M params
Student bireal_kagn, same widths (64, 128, 256) — no widening, 11.94 M params with parity, 5.97 M without
Macro-structure 3 conv blocks → MaxPool ×2 → AdaptiveAvgPool → dense head; hardtanh between blocks; feature taps f1/f2/f3
Group replication groups=8 (input copies before shifted quantization)
Parity offsets parity_rolls=(1, 3), pairing roll (circulant)
Widened baseline HYPERWIDE_WIDTHS = (256, 512, 1024) = 4× teacher width, 94.67 M params

Optimization (identical across all arms)

Optimizer AdamW
Base LR 1e-3, cosine annealed to 0 over each stage (CosineAnnealingLR, T_max = epochs)
Latent binary weights weight_decay = 0, lr × latent_lr_mult = 2.0
All other parameters weight_decay = 1e-4, lr × 1
Batch size 128
Teacher AdamW, lr = 1e-3, weight_decay = 1e-4, cosine, cross-entropy

Zero decay on latent weights is deliberate: decay pulls them toward 0 and causes chronic sign-flip churn late in training, and only the sign is deployed.

Distillation objective

L = KD(logits; T=4.0, alpha=0.9)
  + lambda_at  * attention_transfer(f1,f2,f3)     lambda_at  = 1000.0
  + lambda_div * shift_diversity_penalty          lambda_div = 0.1
  + lambda_ema * KL(student || EMA_student; T=4) * 16    (final A1 stage only)

KD = KL(student/T ‖ teacher/T) · α·T² + CE(student, labels) · (1−α). Attention transfer is the scale-free Zagoruyko–Komodakis criterion — student and teacher share widths and spatial dims, so no learnable projectors are needed. The EMA self-teacher (ema_decay = 0.999) activates after ema_warmup_frac = 0.5 of the final stage, with lambda_ema = 0.3.

Progressive schedule and post-processing

Bit schedule a_bits_schedule = (8, 4, 2, 1), weights inherited between stages
EDE temperature log-annealed 0.1 → 10.0 within each stage (IR-Net)
Post-training RPReLU slopes snapped to signed powers of two (po2_slopes=True)
Post-training BatchNorm recalibration, 100 batches, frozen weights, cumulative averaging

Snapping RPReLU slopes to powers of two makes the negative-half multiply a barrel shift. BN recalibration matters more for binary nets than usual: BN effectively sets the sign thresholds, and folds into the integer comparison threshold on hardware.

Epoch budgets

Dataset Teacher epochs Epochs per stage Ablation epochs
MNIST 20 20 30
CIFAR-10 50 40 40
CIFAR-100 60 50
Tiny ImageNet 60 50

Progressive runs therefore total 160 student epochs on CIFAR-10 (4 stages × 40). direct160 exists precisely to match that total against a single A1 stage.

Data pipeline

Dataset Train augmentation Normalization
MNIST none (0.1307,) / (0.3081,)
CIFAR-10 RandomCrop(32, padding=4) + RandomHorizontalFlip (0.4914,0.4822,0.4465) / (0.2023,0.1994,0.2010)
CIFAR-100 RandomCrop(32, padding=4) + RandomHorizontalFlip (0.5071,0.4865,0.4409) / (0.2673,0.2564,0.2762)

Validation is carved from the train split with fixed split seed 20260716, independent of the experiment seed, so the split is identical across every arm and seed. Val and test use the eval transform (no augmentation). val_fraction=0.1 unless stated; Tiny ImageNet uses the val_fraction=0 fallback and selects the final checkpoint (disclosed in the appendix).

Seeding

run_seeded.py calls bkan.set_seed(seed) before the pipeline builds anything, so model init and dataloader shuffling are both seed-controlled. --deterministic additionally requests deterministic torch algorithms; it is slower and is not used for the reported numbers — the variance protocol is mean ± std over seeds, not bitwise determinism.


9. Repository layout

bkan/
  layers/bireal.py          BiRealKANConv2D / BiRealKANDense — the BiKAN layer
  layers/kagn.py            FP32 Gram-polynomial KAN layers (teacher)
  layers/{quantized,half_binary,boolean}.py   frozen W1A[K] / W1A32 / W1A1 baselines
  quantization/parity.py    degree-2 Walsh parity features (the mechanism)
  quantization/hadamard.py  Walsh–Hadamard basis projection
  quantization/estimators.py  ApproxSign, EDE/Libra-PB, shifted multi-bit quantizer
  models/backbones.py       TEACHER_WIDTHS / HYPERWIDE_WIDTHS backbones
  models/{bireal,bireal_mlp,kan_mlp}.py       BiKAN conv/dense students, KAN MLPs
  models/registry.py        build_model(name, **kw)
  engine/bnn.py             BNN param groups, EDE anneal, BN recalibration
  engine/runlog.py          crash-safe incremental JSON/CSV artifact writer
  distillation/losses.py    KD, attention transfer, EMA self-teacher
  data/{loaders,extra}.py   image / tabular / time-series datasets + val split
  pipelines/bireal_w1a1.py      progressive teacher -> A8 -> A4 -> A2 -> A1
  pipelines/bireal_ablation.py  mechanism arms + width sweep
  pipelines/bireal_families.py  cross-family width-free students
  pipelines/bireal_tabular.py   tabular / MLP suite
  pipelines/fp32_teachers.py    FP32 MNIST teachers for the hardware variants
  pipelines/w1a1_students.py    W1A1 MNIST students for the hardware variants
  quantization/po2_qat.py       power-of-two-aware QAT (the zero-DSP path)
  {pykan,fastkan,efficientkan}_variant/   other KAN families (cross-family + HLS export)

  export/                   HARDWARE EXPORT — Python model → FPGA artifacts
    packing.py                bit-packing / sign convention (single source of truth)
    replay*.py                numpy replay of the folded integer datapath
    verify.py                 replay from the exported artifacts alone
    gram|kagn_conv|efficientkan.py    per-variant weight/threshold export
    hls_*.py                  Vivado HLS C++ codegen + g++ C-sim runners
    DESIGN.md                 the progressive-folding export contract
  paper/gain_table.py       W1A1-vs-FP32 table from the post-P&R reports
  paper/collate.py          joins HLS reports + runs/ + manifests into a table

hls/                        VIVADO HLS DESIGNS — see hls/README.md
  <design>/src, tb          generated kernel + self-checking testbench
  <design>/run_*.tcl        csim / csynth / cosim / implementation scripts
  RESULTS.json              our measured post-route resources + cosim cycles
hw_data/<design>/manifest.json   export provenance + replay-vs-PyTorch agreement

main.py                     CLI: --list / --experiment / --set
run_seeded.py               seeded, collision-free launcher (use this for the paper)
analyze_runs.py             mean ± std + paired t-tests over runs/
synthetic_parity_task.py    Proposition 1 as a runnable experiment
hpc/                        queue definitions + multi-machine worker (see REPRODUCE.md)
experiments/                run_paper_suite.sh (smoke/full) + recover_results.py
tests/                      golden-master + unit + export-replay tests
results/                    the aggregate tables and paired tests behind the paper
docs/architecture.md        design notes
docs/bireal_hw_notes.md     FPGA deployment checklist (the earlier FastKAN harness)
PREREGISTRATION.md          the pre-registered analysis plan (thresholds fixed before the runs)

10. Tests

python -m pytest tests -q

Expect 118 passed in roughly one minute on CPU.

Heads-up: tests/additions/test_extra_datasets.py exercises the download loaders, so on a networked machine the first pytest tests fetches Tiny ImageNet and the UCI tabular sets (~770 MB into ./data). The tests skip gracefully when offline. To keep the run local and fast, either pre-populate ./data or use python -m pytest tests -q --ignore=tests/additions.

Three suites:

  • tests/golden/golden-master characterization tests. They seed the RNG, run every quantizer / estimator / layer / model on CPU with fixed inputs, and compare forward outputs and backward gradients bitwise against recorded references. Their job is to prove a refactor changed no numerics.
  • tests/additions/ — unit tests for the BiKAN layers, the KAN-MLP models, and the extra dataset loaders.
  • tests/engine/ — trainer / evaluate / checkpoint / BN-recalibration tests.

Note on golden references. The binary reference tensors (~1.5 GB) are not shipped in this anonymous release. On the first run the harness records them from the code as shipped and reports recorded; every subsequent run verifies against that snapshot bitwise and reports matched. So:

python -m pytest tests -q     # first run: records references (118 passed)
python -m pytest tests -q     # second run: verifies bitwise (118 passed)

A mismatch on the second run means something is non-deterministic in your environment. To regenerate references after an intended numeric change: UPDATE_GOLDEN=1 python -m pytest tests/golden.

These references are an environment-local baseline, not a cross-version guarantee. The golden suite compares bitwise by design, and binarized layers are unusually sensitive to floating-point reassociation: a change of one ULP in a pre-activation that happens to sit on a sign() or rounding boundary flips a discrete decision, so the output moves by O(1) rather than by an ULP. Comparing references recorded under two different PyTorch releases, 56 of the 67 cases agreed to within 1.5e-5, while 11 — all of them binarized or quantized paths (BooleanKANDense, QuantizedKANDense.a1, the hyper-wide backbones, the FastKAN binary models) — diverged visibly.

So: do not expect references recorded on one PyTorch version to validate on another, and do not read such a mismatch as a regression. Re-record after a toolchain change, or relax the comparison with GOLDEN_ATOL=1e-4 GOLDEN_RTOL=0 python -m pytest tests/golden (which absorbs the 56 well-behaved cases but not the 11 threshold-sensitive ones). None of this affects the reported accuracies, which are seed-averaged over full training runs rather than single-forward-pass bit patterns.


11. Troubleshooting

Unknown experiment '...' — run python main.py --list. Names are exact.

--set value parsed as a string — quote the whole KEY=VALUE when it contains brackets: --set "arms=['full']", not --set arms=['full'].

ImportError: the 'pykan' teacher is unavailablepip install pykan. A vendored copy also ships in pykan/; putting the repository root on PYTHONPATH makes it importable as kan.

Out of memory — lower --set batch_size=64. The width-sweep 4× arm (94.67 M parameters) is the heaviest run in the paper; it needs roughly 4× the memory of the teacher-width arms.

Runs shadowing each other — always use run_seeded.py with a distinct --name. main.py writes to runs/ under the canonical experiment name and two concurrent runs of the same experiment will collide.

A comparison shows no matched seedsanalyze_runs.py --pair needs both tags present at the same seed. Check ls runs/ and that both arms finished.

Parity roll rejected (Parity roll r is 0 (mod C)) — x * x == 1 identically at r ≡ 0, which carries no information. Use nonzero, distinct offsets.

pytest cache file lock on Windows or a synced directory — append -p no:cacheprovider (already the default in pyproject.toml).


Third-party code

bkan/fastkan_variant/fastkan.py (fast-kan, Apache 2.0, original header retained), bkan/efficientkan_variant/efficient_kan.py (efficient-kan, MIT), and pykan/ (pykan, MIT) are vendored unmodified. bkan/layers/kagn.py and bkan/regularization.py are derived from open-source implementations cited in their headers. See THIRD_PARTY_LICENSES.md. Everything else is released under the MIT License in LICENSE.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages