“SparseKAN: Hardware-Cost-Aware Hierarchical Sparsification and Quantization of Kolmogorov–Arnold Networks.”
Paper: https://arxiv.org/abs/2608.00859
This repository contains everything needed to reproduce every number, table,
and figure in the main paper and the technical appendix. All results come from
a single code path (train.py + sparsekan/core/); there are no
per-experiment forks.
Contents
- What the method does
- Installation
- Datasets
- Quick verification (5 minutes)
- Repository layout
- Anatomy of a run
- Running individual experiments
- Full reproduction of the paper
- Reading and aggregating results
- Experiment → paper-table map
- Compute budget
- Troubleshooting
Every KAN edge from input feature i to output o computes a learnable
univariate function split into a cheap base branch and an expensive
basis branch:
phi_oi(x_i) = w_base[o,i] * act(x_i) + sum_k c[o,i,k] * B_k(x_i)
SparseKAN inserts learnable gates at three levels of that computation:
phi_oi(x_i) = g_base[o,i] * w_base[o,i] * act(x_i)
+ g_branch[o,i] * sum_k g_term[o,i,k] * c[o,i,k] * B_k(x_i)
and trains them under a differentiable active-cost regularizer rather than a plain gate-mean penalty:
active_cost = base_cost * sum(g_base) + term_cost * sum(g_branch * g_term)
cost_ratio = active_cost / dense_cost # the headline metric
The same staged protocol runs for every basis family and every dataset:
| Stage | What happens |
|---|---|
| 1 | Sparse training: warmup_epochs at zero λ, then a linear λ ramp over ramp_epochs, then constant λ |
| 2 | Hard pruning: gates snapped to exactly 0/1 at --prune-threshold |
| 2b | (optional) Top-k term hardening: keep the k best terms per active edge |
| 2c | (optional) Structured hardening: shared_k / block / neuron patterns |
| 3− | (optional) QAT: fake-quantize surviving coefficients to B bits (STE) |
| 3 | Fine-tuning with the sparse structure frozen (λ = 0) |
| 4 | Plots, summary.json |
The single most important thing to understand when reading logs. During Stage 1 the reported
CostRatiostays around0.99. This is correct and by design. Soft training arranges importance — which gates sit near 0 and which near 1 — it does not remove computation. All visible compression happens at the explicit hardening stages (2, 2b, 2c) and in the post-hoc tools. One Stage-1 checkpoint then serves many operating points cheaply.Always read a result as the triplet pre-prune / post-prune / healed, and judge the healed number.
Masked models (gates multiplied by 0/1) reduce the analytic cost but not
wall-clock time on dense hardware. compact_pruned_model.py turns a masked
model into a physically smaller one by slicing tensors; that is the model
benchmark_cost.py times.
python -m venv .venv
source .venv/bin/activate # Windows: .venv\Scripts\activate
pip install -r requirements.txtRequirements (requirements.txt): torch>=2.0, torchvision>=0.15,
pandas, matplotlib, einops (GRAM/KAGN variants), numpy, pyyaml.
Tested with Python 3.10–3.12.
Optional extras, only needed for specific features:
| Package | Needed for |
|---|---|
tensorboard or wandb |
--logger tensorboard / --logger wandb |
scikit-learn |
the tabular datasets (wine, dry_bean, moons, …) |
sympy |
some PyKAN reference utilities in the QuantKAN baselines |
A CUDA build of PyTorch is required for --amp and for the CUDA latency
measurements; everything else runs on CPU. The GPU is used automatically when
available.
| Dataset | --dataset |
Source |
|---|---|---|
| MNIST, FashionMNIST, CIFAR-10, CIFAR-100, SVHN | mnist, fashion_mnist, cifar10, cifar100, svhn |
auto-downloaded by torchvision into --data-dir (default ./data) |
| Tiny ImageNet, Imagenette, Imagewoof, ImageNet | tiny_imagenet, imagenette, imagewoof, imagenet |
expects an extracted ImageFolder layout <data-dir>/<subdir>/{train,val} |
| Wine, Dry Bean, Mushroom, JSC (OpenML), Traffic (California) | wine, dry_bean, mushroom, jsc_openml, traffic_california |
fetched via scikit-learn / OpenML on first use, then cached |
| Synthetic | toy_regression, moons, circle_in_circle |
generated deterministically, no download |
On a machine without internet you can still exercise the full mechanics:
python tests/make_synthetic_mnist.py # writes a tiny fake MNIST into ./dataAll splits, normalizations, and augmentations are fixed in
sparsekan/core/data.py.
Checkpoint selection. --val-size N carves a seeded validation split off
the training set (evaluated without augmentation) and, when present, it is
what the best-epoch checkpoint of each stage is selected on. The campaign in
experiments_manifest.py does not request one, so with the default
--val-size 0 the best-epoch checkpoint is selected by test accuracy
(select_metric = (val_stats or test_stats)["metric"] in
sparsekan/core/engine.py). This applies identically to every arm of every
comparison and all comparisons are paired by seed, so reported differences
are unaffected; absolute numbers are best-epoch-on-test. Pass
--val-size 5000 (as configs/cifar100_eff_budget.yaml does) for a fully
held-out selection protocol.
Run these three commands from the repository root before anything else.
# 1. Layer math is identical to the reference implementations
python tests/test_layer_equivalence.py # -> ALL EQUIVALENCE TESTS PASSED
# 2. Every feature used by the paper works end to end
python tests/test_new_features.py # -> ALL NEW-FEATURE TESTS PASSED
# 3. All paper-model presets build and run
python tests/test_presets_integration.py # -> ALL PRESET TESTS PASSEDThen a real 2-minute training run:
python train.py --variant efficientkan --dataset mnist --hidden-dims 32 \
--epochs 3 --warmup-epochs 1 --ramp-epochs 1 --finetune-epochs 1 \
--batch-size 128 --run-name sanityExpected shape of the output:
SparseKAN unified training | variant=efficientkan | dataset=mnist
Model parameters: ...
[Stage 1] Sparse training
Epoch 001 | Train Accuracy ... | Test ... | CostRatio 1.0000 | 0.8s
...
[Stage 2] Hard pruning at threshold 0.5
[Stage 3] Fine-tuning for 1 epochs
Final summary
and on disk:
ls runs_unified/mnist_efficientkan/sanity/
# config.json metrics.csv summary.json *.pt *.pngsmoke_test.sh chains a broader set of tiny runs — the three test suites, a
dense/sparse pair, all seven basis families, both alternative gate
relaxations, budget mode, in-training top-k, all three structured modes, QAT,
both QAS levels, all four post-hoc ablation score modes, a preset, a
regression task, a YAML config, an executor dry run, physical compaction,
both benchmarks, and aggregation:
bash smoke_test.sh # ~40 min on a single GPU; longer on CPU
FAST=1 bash smoke_test.sh # skips the conv/compaction/latency partIt exits non-zero on the first failure and ends with SMOKE TEST PASSED.
Runs land in runs_smoke/, aggregated into runs_smoke/_summary/. These are
pipeline-validation numbers, not paper results.
SparseKAN/
├── README.md # this file
├── requirements.txt
├── smoke_test.sh # tiny end-to-end runs over the whole pipeline
│
├── train.py # THE trainer: any variant x any dataset, staged protocol
├── ablate.py # post-hoc ablations (importance_prune | topk_harden)
├── compact_pruned_model.py # masked model -> physically smaller model
├── benchmark_cost.py # analytic FLOPs + measured latency (gated / dense / compact)
├── latency_batch_sweep.py # batch-size crossover figure
├── plot_efficiency.py # accuracy-vs-cost Pareto figure + LaTeX rows
├── summarize.py # aggregate runs_unified/ -> CSV + LaTeX table + Pareto plot
├── eval_quantkan_baseline.py # as-is dense/PTQ KAN baselines
│
├── experiments_manifest.py # SINGLE SOURCE OF TRUTH: every training run in the paper
├── run_experiments.sh # sharded, parallel, resumable executor over the manifest
├── run_posthoc.py # idempotent driver for every post-hoc stage
│
├── exp_scripts/ # the campaign, in the order it was run (see its README)
│ ├── exp_all_tierA.sh # chains steps 1-3
│ ├── exp_1_anchors.sh # E1 dense anchors + E2a base sparse
│ ├── exp_2_structuring.sh # E4/E5a/E9a + e3/e11a post-hoc
│ ├── exp_3_composition.sh # E6a/E7a + e5c/e6b/e6c/e7b/e8 post-hoc
│ ├── exp_4_tierb_and_report.sh # Tier B, QAS, aggregation
│ ├── exp_scale_presets.sh # large presets + tabular rung
│ └── exp_baselines_ptq.sh # as-is baseline PTQ rows
│
├── configs/ # example YAML configs (single runs + a sweep)
├── tests/ # equivalence, feature, and preset test suites
│
├── hw/ # FPGA/HLS track — see hw/README.md
│ ├── export/ # checkpoint -> folded artifacts -> HLS inputs,
│ │ # NumPy golden reference, verification gates,
│ │ # HLS report parsing and table building
│ ├── hls/common/ # the HLS C++ sources (3 families x fp32/int)
│ ├── hls/<design>/ # per-design run_hls.tcl + generated params.h
│ ├── artifacts/<tag>/ # arch.json / manifest.json provenance records
│ └── reports/ # generated: per-design JSON + markdown tables
│
└── sparsekan/
├── core/ # THE implementation
│ ├── layers/
│ │ ├── base.py # SparseKANLayerBase: the three gate tensors, gated
│ │ │ # forward, mean/global/cost regularization, entropy,
│ │ │ # sparsity report, hard prune, top-k, occupancy, QAT
│ │ ├── efficientkan.py # B-spline basis
│ │ ├── fastkan.py # Gaussian RBF basis
│ │ ├── kagn.py # Gram polynomial basis
│ │ ├── pykan.py # original-KAN B-splines with adaptive grids
│ │ ├── chebykan.py # Chebyshev polynomial basis
│ │ ├── wavkan.py # Ricker wavelet basis
│ │ ├── relukan.py # squared-ReLU bump basis
│ │ ├── conv_base.py # sparse KAN convolution base
│ │ └── conv_variants.py # kagn_conv / efficientkan_conv
│ ├── gates.py # sigmoid | hard_concrete (L0) | gumbel_st relaxations
│ ├── models.py # SparseKANClassifier (flatten | convstem frontends)
│ ├── conv_model.py # conv_kan architecture
│ ├── quantkan_models.py # sparsifiable ports of published KAN architectures
│ ├── data.py # dataset registry (image / tabular / regression)
│ ├── engine.py # the staged protocol, budget mode, schedulers, AMP, QAS
│ ├── structured.py # shared_k / block / neuron hardening
│ ├── quant.py # fake-quantization (fixed-bit and learned-bit)
│ ├── posthoc.py # importance_prune / topk_harden kernels
│ ├── analysis.py # summarizer engine
│ ├── reporting.py # metrics.csv, plots, checkpoints, summary.json
│ └── utils.py # seeding, device, lambda schedule
├── presets.py # 25 named paper-model presets
├── quantkan_baselines.py # registry over the vendored baseline models
├── quantkan_vendor/ # published baseline models, vendored unmodified
├── variants/ # third-party reference sources reused verbatim
│ ├── fastkan/fastkan.py (RadialBasisFunction, SplineLinear)
│ └── pykan/pykan_spline.py (B_batch, coef2curve, curve2coef, extend_grid)
└── vendor/kan_datasets.py # tabular/timeseries dataset loaders
Every train.py invocation writes one folder:
runs_unified/<dataset>_<variant>/<run-name>/
├── config.json # the EXACT resolved arguments (reproducibility record)
├── metrics.csv # one row per epoch/stage, full per-layer report
├── summary.json # the final headline numbers <-- "run is COMPLETE" marker
├── best_before_prune.pt # best Stage-1 checkpoint
├── last_before_prune.pt
├── after_hard_prune.pt # immediately after Stage 2 (before 2b/2c)
├── best_after_prune.pt # best fine-tuned checkpoint <-- post-hoc tools load this
├── last_after_prune.pt
└── *.png # accuracy, cost-ratio, edge-state, term-usage curves
Key summary.json fields:
| Field | Meaning |
|---|---|
best_after_finetune |
headline accuracy (or lowest RMSE for regression) |
global_cost_ratio |
active / dense analytic compute of the hardened model |
quant_bits, joint_bit_cost_ratio |
cost_ratio × bits / 32 when quantized |
structured_mode |
which Stage-2c pattern was applied |
learned_bits, learned_bit_cost_ratio |
QAS Level 3 per-layer bit-widths |
qat_from_epoch |
QAS Level 2 (joint gate + quantization learning) |
summary.json exists iff the run finished. Every executor in this
repository uses that as its resume rule, so any script can be killed and
re-invoked without redoing completed work.
metrics.csv rows carry a stage column
(train / after_hard_prune / after_topk / after_structured /
after_quantize / finetune), the per-epoch losses and λs, train/val/test
metrics, and the full per-layer sparsity report (gate means, active ratios,
edge states full/base-only/branch-only/pruned, active counts, estimated
active/dense cost, cost ratio).
Everything in the paper is a modification of these two commands.
Dense anchor — all sparsity pressure off, gates initialized open. Uses the same epoch budget as the sparse runs so comparisons are budget-matched:
python train.py --variant efficientkan --dataset mnist --hidden-dims 256 128 \
--epochs 80 --warmup-epochs 20 --ramp-epochs 20 --finetune-epochs 20 \
--max-lambda-base 0 --max-lambda-branch 0 --max-lambda-term 0 \
--entropy-weight 0 \
--init-branch-gate-value 0.999 --init-term-gate-value 0.999 \
--seed 42 --run-name e1_dense_s42Base sparse — the verified recipe used throughout:
python train.py --variant efficientkan --dataset mnist --hidden-dims 256 128 \
--epochs 80 --warmup-epochs 20 --ramp-epochs 20 --finetune-epochs 20 \
--max-lambda-base 1e-2 --max-lambda-branch 5e-3 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 \
--seed 42 --run-name e2a_base_s42Flag by flag: the three --max-lambda-* are the peak pressures on the base /
branch / per-term gates; they are 0 for --warmup-epochs, ramp linearly over
--ramp-epochs, then hold. --sparsity-reg-type cost makes the penalty the
differentiable active cost (each gate weighted by what it gates) rather
than a plain gate mean. --entropy-weight pushes gates away from 0.5 so the
0/1 decision at Stage 2 is decisive.
Each subsection is independent and only needs the Section-6 base-sparse
checkpoint (where noted). Set RUN once:
RUN=runs_unified/mnist_efficientkan/e2a_base_s42Occupancy-weighted global importance — rank every term in the network by
rho × |c| × g and close the lowest-importance fraction:
python ablate.py --run-dir $RUN --ablation importance_prune \
--keep-ratios 1.0 0.9 0.75 0.5 0.25 0.15 0.1 0.05 --finetune-epochs 5Per-edge top-k (the controlled local baseline). Run all three score modes:
for MODE in gate coeff occupancy_coeff; do
python ablate.py --run-dir $RUN --ablation topk_harden \
--topk-values 7 6 5 4 3 2 --score-mode $MODE --finetune-epochs 5
mv $RUN/topk_harden_ablation.csv $RUN/topk_harden_${MODE}.csv
doneThe rename matters. The output CSV name is fixed and rows append, so without renaming you cannot tell the score modes apart.
python run_posthoc.py --stages e3does all of Section 7.1 with the renames handled — prefer it for real runs.
coeff is plain magnitude pruning, itself a paper baseline. Fairness rule:
use identical --finetune-epochs in every compared sweep.
These are train.py flags that add Stages 2b/2c to the pipeline.
In-training top-k (unstructured pattern — analytic cost drops, dense-hardware time does not):
python train.py --variant efficientkan --dataset mnist --hidden-dims 256 128 \
--epochs 80 --warmup-epochs 20 --ramp-epochs 20 --finetune-epochs 20 \
--max-lambda-base 1e-2 --max-lambda-branch 5e-3 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 \
--topk-after-prune 4 --seed 42 --run-name e4_topk4_s42Shared-k — all edges share the same k basis terms, the pattern physical
compaction can actually harvest:
python train.py --variant kagn_conv --dataset cifar10 --arch conv_kan \
--hidden-dims 32 64 128 \
--epochs 120 --warmup-epochs 30 --ramp-epochs 30 --finetune-epochs 30 \
--max-lambda-base 1e-2 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 \
--structured-mode shared_k --shared-k 3 --share-axis input \
--structured-score-mode occupancy_coeff --amp \
--seed 42 --run-name c10_sharedk3_s42A large accuracy dip right after [Stage 2c] is normal; judge the healed
number at the end of Stage 3. --share-axis layer shares one term set across
the whole layer (harsher, most hardware-friendly).
Neuron/channel and block modes:
python train.py ... --structured-mode neuron --neuron-keep-ratio 0.5 \
--structured-score-mode occupancy_coeff --run-name e5a_neuron05_s42
python train.py ... --structured-mode block --block-size 8 --block-keep-ratio 0.5Scoring rule. After shared-
k, never usegatescoring for neuron cuts: the gate means become numerically identical across units, so the cut degenerates to a random selection. Usecoefforoccupancy_coeff. This is the designed failure experiment reported in the ablations.
Everything above multiplies by 0/1 masks. Compaction slices the tensors:
python compact_pruned_model.py \
--run-dir runs_unified/cifar10_kagn_conv/c10_sharedk3_s42 \
--then-neuron-keep-ratio 0.7 --neuron-stages 2 --stage-heal-epochs 10 \
--neuron-score coeff --finetune-epochs 20 --save compact_n0.7.ptIn order, this performs a staged post-hoc channel cut (each stage cuts to
keep^(s/N) then heals with cosine LR), dead-unit removal with cross-layer
propagation, and shared-k term gathering ([out,in,4] -> [out,in,3] plus a
kept_term_idx buffer; conv spatial kernels too). It prints the per-layer
shape changes, the new parameter count, the parameter reduction, and the
healed accuracy of the compact model.
Add --quant-bits 4 to get a deployable sparse + compact + quantized
artifact. Compact-vs-masked forward outputs differ by ~5e-5 (residual
sigmoid(-20) leakage in the masked model); the compact model is the cleaner
one.
python benchmark_cost.py \
--run-dir runs_unified/cifar10_kagn_conv/c10_sharedk3_s42 \
--compact-ckpt compact_n0.7.pt \
--batch-sizes 1 32 256 1024 --repeats 50 --device cuda \
--output-csv bench_c10.csvHow to read the output: the gated (masked) latency ratio sits at ≈1.0 at
every batch size — expected, since 0/1 masks do not speed up dense hardware.
The compact ratio crosses below 1.0 around batch 128–256 and keeps
improving with batch size. The compact latency ratio stays above the
parameter fraction; that gap is the un-prunable basis-construction residual
discussed in the paper. Add --device cpu --batch-sizes 1 32 --repeats 20
for the edge-deployment story.
# batch-size crossover figure
python latency_batch_sweep.py --run-dir <run> --compact-ckpt compact_n0.7.pt \
--batch-sizes 1 8 32 128 256 512 1024 2048 --device cuda --out crossover.png
# Pareto figure + paste-ready LaTeX rows
# CSV columns: label,cost_ratio,accuracy,params_reduction,latency_ratio
python plot_efficiency.py --structured-csv structured_points_cifar10.csv \
--dense-acc 0.8263 --acc-percent --title "SparseKAN on CIFAR-10" \
--out c10_efficiency.pngAdd --quant-bits 8 (or 4, 2) to any training command. The run logs an
after_quantize row and the summary reports
joint_bit_cost_ratio = cost_ratio × bits / 32.
Reporting rule: quote the accuracy delta against the same sparsity regime at FP32, which isolates the quantization cost from the pruning cost.
Level 1 — quantization-aware scoring (post-hoc; needs only a base-sparse
checkpoint). Ranks terms by rho × |Q_4(c)| × g, i.e. the coefficient as the
4-bit quantizer will see it:
python ablate.py --run-dir $RUN --ablation topk_harden \
--topk-values 7 6 5 4 3 2 --score-mode quant_coeff \
--score-quant-bits 4 --finetune-epochs 5
mv $RUN/topk_harden_ablation.csv $RUN/topk_harden_quant_coeff.csvCompare against the Section-7.1 occupancy_coeff CSV at matched k.
Level 2 — joint gate + quantization learning (one extra flag). Its control
is the identical command without --qat-from-epoch; compare healed accuracy
at the identical joint_bit_cost_ratio:
python train.py ... --topk-after-prune 4 --quant-bits 4 --qat-from-epoch 1 \
--run-name e11b_joint_s42Level 3 — learned per-layer bit-widths:
python train.py ... --learn-bits --max-lambda-bits 1e-1 \
--run-name e11c_learnbits_s42The log prints soft per-layer bit-widths during Stage 1 and the hardened
values at Stage 2; summary.json gains learned_bits and
learned_bit_cost_ratio. lambda_bits is uncalibrated — sweep
{1e-2, 1e-1, 1.0}. All QAS effects require ≤4 bits to be visible.
# hard-concrete (L0) gates: exact zeros during training, no threshold dependence
python train.py --variant efficientkan --dataset mnist --hidden-dims 128 \
--gate-type hard_concrete --sparsity-reg-type cost --run-name hc_run
# Gumbel straight-through: exactly binary forward gates
python train.py ... --gate-type gumbel_st --gate-temperature 0.667
# target a cost ratio directly instead of tuning lambdas (Lagrangian dual ascent)
for B in 0.5 0.3 0.15; do
python train.py --variant pykan --dataset mnist --hidden-dims 128 \
--cost-budget $B --dual-lr 1.0 --entropy-weight 1e-5 \
--run-name mnist_pykan_budget${B}_s42
doneWith --cost-budget, metrics.csv gains lambda_dual and
soft_cost_ratio columns and the λ ramp is bypassed.
Same recipe, different --variant:
for V in efficientkan fastkan kagn chebykan relukan wavkan pykan; do
python train.py --variant $V --dataset mnist --hidden-dims 256 128 \
--epochs 80 --warmup-epochs 20 --ramp-epochs 20 --finetune-epochs 20 \
--max-lambda-base 1e-2 --max-lambda-branch 5e-3 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 \
--seed 42 --run-name e9a_${V}_s42
doneTerm counts G differ per family (efficientkan/fastkan/relukan/
wavkan/pykan: G = 8; kagn/chebykan: G = 4), so scale
--topk-values and --shared-k accordingly. Convolutional variants use
--variant kagn_conv --arch conv_kan (--kernel-mode quantkan is the default
and the correct one; box mode is only kept as a diagnostic row).
A --preset reproduces a published KAN architecture as a sparse network —
it fixes dataset / variant / architecture / widths / layer hyperparameters,
while schedule and sparsity flags stay yours:
python train.py --preset kan_convnet_cifar10 \
--epochs 120 --warmup-epochs 30 --ramp-epochs 30 --finetune-epochs 30 \
--max-lambda-base 1e-2 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 --amp \
--seed 42 --run-name preset_convnet_c10_s42
# tabular example: the full pipeline runs in seconds
python train.py --preset kan_wine --epochs 150 --warmup-epochs 40 \
--ramp-epochs 40 --finetune-epochs 40 --batch-size 32 \
--seed 42 --run-name wine_s42Each preset prints a fidelity note at the start of the run stating exactly how
it deviates (if at all) from the published architecture. Presets are listed in
sparsekan/presets.py.
The FPGA track lives in hw/ and starts from any trained run.
Gates are folded into the exported tensors, so the accelerator contains no
gating logic at all — structural decisions become compile-time constants:
# 1. fold gates into effective tensors + emit the golden logits
python hw/export/export_for_hw.py \
--run-dir runs_unified/mnist_efficientkan/e1_efficientkan_dense_s42 \
--tag m_fp32
# 2. verify the NumPy golden reference reproduces PyTorch (gates V1-V4)
python hw/export/verify_export.py --fp32 hw/artifacts/m_fp32 \
--out phase1_verification
# 3. pack the HLS design inputs (weights.bin + params.h)
python hw/export/pack_weights.py --art hw/artifacts/m_fp32 \
--out hw/hls/fp32_dense
# 4. synthesize (needs Vitis HLS; everything above is CPU-only)
cd hw/hls/fp32_dense && vitis-run --mode hls --tcl run_hls.tclSteps 1–3 reproduce every accuracy number in the hardware tables without any
FPGA tooling. hw/README.md documents the full pipeline, the verification
gates, and the mapping from each of the 18 published designs to the training
run it came from.
The dense low-precision hardware baselines use LSQ instead of the default min-max quantizer:
python train.py ... --quant-bits 8 --quant-scheme lsq \
--run-name e1_efficientkan_dense_lsq8_s42--quant-scheme lsq learns a per-output-channel step (Esser et al., ICLR'20)
rather than deriving it from the running maximum. Both schemes produce the same
integer export spec — symmetric, per-output-channel, no zero-point — so the
hardware path is unchanged.
python eval_quantkan_baseline.py --list # available baselines
python eval_quantkan_baseline.py --baseline kan_convnet_cifar10 --dataset cifar10These evaluate the published dense models unchanged (and their post-training quantization), which is what the PTQ baseline rows in the paper report.
Every flag can come from a YAML file; CLI flags override it.
python train.py --config configs/mnist_pykan_cost.yaml # as-is
python train.py --config configs/mnist_pykan_cost.yaml --seed 123 # CLI overrides YAML
python train.py --config configs/cifar100_eff_budget.yaml # budget + HC + AMPexperiments_manifest.py generates every training run in the paper as a
structured record, filterable by tier / block / dataset / seed:
python experiments_manifest.py --count # per-block totals
python experiments_manifest.py --tier A --count # 324 Tier-A runs
python experiments_manifest.py --tier A --block E2a --list # the exact commands
python experiments_manifest.py --tier A --seeds 42 --json # machine-readableCurrent totals: 362 training runs (324 Tier A + 35 Tier B + 3 Tier C),
each over seeds {42, 43, 44} where applicable. Blocks:
| Block | Tier | Runs | What it establishes |
|---|---|---|---|
E1 |
A | 39 | Dense anchors (7 MLP families on MNIST; CIFAR-10/100 MLP + conv; CIFAR-100 narrow vs wide) |
E2a |
A | 18 | Base sparse checkpoints — the shared substrate every post-hoc block reuses |
E2b |
B | 15 | Sensitivity: λ grid, entropy weight, mean/global/cost regularizer, schedule |
E4 |
A | 72 | Term structuring: top-k (G/2, G/4) and shared-k (3G/4, G/2) × input/layer axes |
E5a |
A | 36 | In-training neuron/channel hardening at keep 0.75 / 0.5 / 0.25 |
E5d |
B | 8 | Block mode (sizes 4/8, keep 0.75/0.5) |
E6a |
A | 6 | Composition: top-k and neuron in one run |
E7a |
A | 90 | 8-bit and 4-bit QAT across five sparsity regimes |
E7d |
C | 1 | 2-bit stress test |
E9a |
A | 15 | Basis-family transfer (the five families not in E2a) |
E9c |
B | 2 | Conv kernel-mode ablation (box diagnostic) |
E11b |
B | 9 | QAS Level 2: joint vs two-stage at 4 bits |
E11c |
B | 3 | QAS Level 3: learned bit-widths, lambda_bits sweep |
E12a/b/c |
A | 48 | Tabular + time-series rung (Wine, Dry Bean, JSC, Traffic) × dense / sparse / q8 / q4 |
Canonical schedules (epochs, warmup, ramp, finetune) and widths are fixed in
SCHED and HIDDEN at the top of experiments_manifest.py:
| Dataset / arch | Schedule | Hidden dims |
|---|---|---|
| MNIST, classifier | 80 / 20 / 20 / 20 | [256, 128] |
| MNIST, conv_kan | 60 / 15 / 15 / 15 | [16, 32] |
| CIFAR-10, classifier | 100 / 25 / 25 / 25 | [1024, 512] |
| CIFAR-10, conv_kan | 120 / 30 / 30 / 30 | [32, 64, 128] |
| CIFAR-100, conv_kan | 120 / 30 / 30 / 30 | [32, 64, 128] (wide: [64, 128, 256]) |
run_experiments.sh runs any slice of the manifest. It is sharded,
parallel, and resumable — a run is complete iff its summary.json exists,
one failure never aborts the session, and per-run logs land in
runs_unified/_logs/<run>.log:
DRY=1 TIER=A ./run_experiments.sh # print the plan, run nothing
WORKERS=6 TIER=A BLOCKS="E1 E2a" ./run_experiments.sh # anchors, 6 jobs in parallel
SHARD=0 NSHARDS=2 TIER=A ./run_experiments.sh # machine 0 of 2
SEEDS="42" TIER=A ./run_experiments.sh # 1-seed pilot passrun_posthoc.py drives every post-hoc stage. It is idempotent (skips work
whose output already exists) and safe to re-invoke after every session:
python run_posthoc.py --stages e3 # importance + topk sweeps on E2a/E9a
python run_posthoc.py --stages e11a # QAS Level-1 sweeps
python run_posthoc.py --stages e5c e6b e6c e7b e8
# e5c neuron-score ablation (gate vs coeff)
# e6b structured-first ordering control
# e6c dual-axis compaction (shared-k run -> post-hoc neuron -> compact)
# e7b quantized compact artifacts (q8 / q4)
# e8 CUDA + CPU benchmarks on every compact checkpoint
python run_posthoc.py --stages e3 --dry # preview any stageThe exp_scripts/ wrappers run the campaign in the order it was originally
executed. Each is resumable; rerun any of them after a crash.
# 0. Preview everything without running it
DRY=1 ./exp_scripts/exp_all_tierA.sh
# 1. Anchors: E1 dense + E2a base sparse + the tabular rung.
# Prints decision gate G1 (the CIFAR-100 conv width choice).
WORKERS=6 ./exp_scripts/exp_1_anchors.sh
# 2. Structuring: e3 + e11a post-hoc sweeps, then E4 / E5a / E9a.
# Prints decision gate G2 (importance-vs-topk framing).
WORKERS=6 ./exp_scripts/exp_2_structuring.sh
# 3. Composition: E6a / E7a, then e5c / e6b / e6c / e7b / e8 post-hoc.
# Prints decision gate G3 (how much 4-bit QAT costs -> whether to run QAS).
WORKERS=6 ./exp_scripts/exp_3_composition.sh
# 4. Tier B, optional QAS blocks, and aggregation.
WITH_QAS=1 ./exp_scripts/exp_4_tierb_and_report.sh
# 5. Large presets and the tabular/time-series rung
./exp_scripts/exp_scale_presets.sh
# 6. As-is baseline PTQ rows
./exp_scripts/exp_baselines_ptq.sh
# 7. Final aggregation -> paper tables and the Pareto figure
python summarize.py --name unified --runs-root runs_unified \
--output-dir unified_summaryEvery script accepts WORKERS=N (parallel jobs), SHARD=i NSHARDS=n
(split across machines), SEEDS="42" (single-seed pilot), and DRY=1.
Steps 1–3 print decision gates (G1, G2, G3) — points where the
campaign branched on an observed result. They are reproduced verbatim so the
reported decision procedure is auditable, not so the reader must make a
choice; following the defaults reproduces the paper.
This chain produces the structured-Pareto + latency story on its own:
# 1. dense anchor (+ 8-bit QAT)
python train.py --variant kagn_conv --dataset cifar10 --arch conv_kan \
--hidden-dims 32 64 128 --epochs 120 --warmup-epochs 30 --ramp-epochs 30 \
--finetune-epochs 30 --max-lambda-base 1e-2 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 --quant-bits 8 --amp \
--seed 42 --run-name c10_dense_q8_s42
# 2. the shared-k3 structured run (same flags, no --quant-bits)
python train.py --variant kagn_conv --dataset cifar10 --arch conv_kan \
--hidden-dims 32 64 128 --epochs 120 --warmup-epochs 30 --ramp-epochs 30 \
--finetune-epochs 30 --max-lambda-base 1e-2 --max-lambda-term 1e-3 \
--sparsity-reg-type cost --entropy-weight 1e-5 \
--structured-mode shared_k --shared-k 3 --share-axis input \
--structured-score-mode occupancy_coeff --amp \
--seed 42 --run-name c10_sharedk3_s42
# 3. two compaction operating points off the SAME run
for KR in 0.7 0.5; do
python compact_pruned_model.py \
--run-dir runs_unified/cifar10_kagn_conv/c10_sharedk3_s42 \
--then-neuron-keep-ratio $KR --neuron-stages 2 --stage-heal-epochs 10 \
--neuron-score coeff --finetune-epochs 20 --save compact_k3_n${KR}.pt
done
# 4. measured latency for each point
for KR in 0.7 0.5; do
python benchmark_cost.py \
--run-dir runs_unified/cifar10_kagn_conv/c10_sharedk3_s42 \
--compact-ckpt compact_k3_n${KR}.pt \
--batch-sizes 1 32 256 1024 --repeats 50 --device cuda \
--output-csv bench_k3.csv
done
# 5. figure + LaTeX rows
python plot_efficiency.py --structured-csv structured_points_cifar10.csv \
--dense-acc 0.8263 --acc-percent --out c10_efficiency.png# one run
python -c "import json,sys; d=json.load(open(sys.argv[1])); [print(k,':',v) for k,v in d.items()]" \
runs_unified/mnist_efficientkan/e2a_base_s42/summary.json
head -3 runs_unified/mnist_efficientkan/e2a_base_s42/metrics.csv
# everything
python summarize.py --name unified --runs-root runs_unified \
--output-dir unified_summaryThe summarizer groups runs by dataset × variant × method (method inferred
from config.json as dense, sparse_<reg>, or sparse_<reg>_topk<k>) and
writes:
| File | Contents |
|---|---|
unified_summary/summary_all_runs.csv |
one row per run |
unified_summary/summary_grouped.csv |
mean ± std per group across seeds |
unified_summary/summary_grouped.tex |
the paper-ready LaTeX table |
unified_summary/pareto_accuracy_cost.png |
accuracy-vs-cost scatter |
These are not optional if you want numbers comparable to the reported ones:
- Seeds. Frontier results use
{42, 43, 44}; the follow-up paired ablation campaign uses{42, 43, 44, 45, 46}. Report mean ± sample standard deviation. - Matched budgets. Every arm of a comparison uses the same epoch and fine-tune budget.
- Paired design. Post-hoc comparisons share the same Stage-1 checkpoint;
the tooling does this naturally because all post-hoc stages read
best_after_prune.ptfrom the same run directory. - The triplet. Report pre-prune / post-prune / healed, and headline the healed number.
- Masked vs compact. Label analytic cost ratio (masked) separately from physical parameters / MB / latency (compact). Never present one as the other.
- Latency. Always quote device and batch size.
- Dominated points (e.g. shared-
k=2) are reported as diagnostics, not hidden. - No single-seed promotion. Exploratory single-seed numbers never enter a table as a claim.
| Paper artifact | Produced by |
|---|---|
| Dense and Stage-1 anchor table | E1 + E2a → summarize.py |
| Basis-family transfer table | E9a (+ E2a for efficientkan/kagn) |
| Selection-control comparison (learned vs truncation vs random) | run_posthoc.py --stages e3 |
Term-structuring table (top-k, shared-k, axes) |
E4 |
| Neuron / block structured tables | E5a, E5d |
| Ordering control (structured-first vs term-first) | E6a + run_posthoc.py --stages e6b |
| Quantization tables (8-bit, 4-bit, 2-bit stress) | E7a, E7d |
| Dense 4-bit PTQ baseline rows | exp_scripts/exp_baselines_ptq.sh |
| QAS Levels 1 / 2 / 3 | run_posthoc.py --stages e11a; E11b; E11c |
| Physical compaction table (params, MB, top-1 agreement) | run_posthoc.py --stages e6c e7b |
| Latency tables and the batch-size crossover figure | run_posthoc.py --stages e8; latency_batch_sweep.py |
| Accuracy-vs-cost Pareto figure | plot_efficiency.py, summarize.py |
| Tabular / time-series rung | E12a, E12b, E12c |
| Sensitivity appendix (λ, entropy, regularizer, schedule) | E2b |
| Cost-model reweighting robustness | benchmark_cost.py + summarize.py |
The full campaign is 362 training runs plus the post-hoc stages. Practical guidance:
- A single-seed Tier-A pilot (
SEEDS="42") is roughly a third of the cost and reproduces every qualitative conclusion. - MNIST and tabular runs are minutes each; CIFAR-10/100 conv runs at
120 epochs dominate the budget. Use
--ampon CUDA (the manifest already adds it for CIFAR conv runs). WORKERS=4..6saturates a single modern data-center GPU for the MNIST and tabular blocks; use fewer for the CIFAR conv blocks.- Post-hoc stages (
e3,e5c,e6b,e6c,e7b,e8) are minutes-to-hours each and reuse existing checkpoints — they never retrain from scratch. - Everything is resumable, so the campaign can be spread across sessions and
machines (
SHARD/NSHARDS) without bookkeeping.
Exact wall-clock figures for the hardware used in our runs are reported in the technical appendix.
| Symptom | Cause and fix |
|---|---|
CostRatio stays ≈0.99 through Stage 1 |
Normal and by design (Section 1). Compression happens at Stages 2/2b/2c. |
Large accuracy dip right after [Stage 2c] |
Normal. Judge the healed number. If healed is far below dense, raise k (k=2 → k=3 is a large gain at nearly equal compression). |
| Post-hoc neuron cut collapses to ~random accuracy | gate scoring used after shared-k. Use --neuron-score coeff. |
| Compact vs masked forward differs by ~5e-5 | Expected sigmoid(-20) leakage residue in the masked model; the compact model is the correct one. |
| Ablation CSVs mix score modes | The CSV name is fixed and rows append. Rename between sweeps, or use run_posthoc.py. |
| Conv model stuck near 50% on MNIST | --kernel-mode box was used. The default quantkan mode is the correct one. |
| Dataset download fails behind a firewall | Pre-download into --data-dir, or use python tests/make_synthetic_mnist.py for mechanics-only checks. |
--amp warns and runs full precision |
No CUDA device; this is the intended fallback. |
A run has config.json but no summary.json |
It crashed. Delete that run directory and re-issue the same executor command; the run is redone and completed runs are skipped. |
Slow or NaN runs on relukan / wavkan / pykan |
Pilot at one seed first and lower --lr; these families are more sensitive than the spline and polynomial ones. |
Run this before archiving the directory for distribution.
# 1. Byte-code caches. REQUIRED: .pyc files embed the absolute source path of
# the machine that compiled them, which de-anonymizes the archive.
find . -name "__pycache__" -type d -prune -exec rm -rf {} +
find . -name "*.pyc" -delete
# 2. Downloaded datasets and experiment outputs (all regenerable).
rm -rf data runs_unified runs_smoke unified_summary
# 3. Generated hardware artifacts and HLS build products (all regenerable;
# the arch.json / manifest.json provenance records are kept).
find hw/artifacts -type f ! -name "*.json" -delete
rm -rf hw/hls/*/proj_* hw/hls/*/logs hw/hls/*/*.bin
rm -f hw/reports/*.json hw/reports/*.md hw/reports/logs/*.logThen verify nothing identifying survived:
grep -rIl -e "$USER" -e "$HOME" . 2>/dev/null # expect no output.gitignore covers all of the above for git, but a plain zip/tar of the
working directory does not — hence the explicit cleanup.
sparsekan/quantkan_vendor/ and sparsekan/variants/ contain published
reference implementations vendored unmodified so that the sparse layers can
be verified against them (tests/test_layer_equivalence.py) and so that the
baseline numbers come from the original code. Original copyright headers are
preserved in place; the corresponding citations appear in the paper.