claroche1/sparse-nsnet2-checkpoints

sparse-nsnet2 — checkpoints

Best-PESQ checkpoints from the eco8-neaixt compression sweep: NSNet2 speech enhancement with the FC and GRU layers swappable between dense, Butterfly, block-diagonal, and genuine two-factor Monarch factorizations.

Trained on VoiceBank-DEMAND-16k, n_fft 512, batch 256. PESQ on the full 824-utterance test split.

⚠️ Two corrections you should read before quoting these numbers

1. Naming: the old monarch_* runs were NOT Monarch. They are a single block-diagonal factor (one block-diagonal matrix per projection, zero cross-block mixing). A genuine Monarch is a two-factor construction (block-diagonal × permutation × block-diagonal) with full cross-channel mixing. The old runs have therefore been renamed to blockdiag_*, and the monarch_* names now hold genuinely Monarch models:

old name (was mislabeled)	now	monarch_* today
monarch_8	→ blockdiag_8	genuine 2-factor Monarch
monarch_full	→ blockdiag_full	genuine 2-factor Monarch
monarch_fc	→ blockdiag_fc	genuine 2-factor Monarch
wide_monarch	→ wide_blockdiag	genuine 2-factor Monarch

If you previously pinned monarch_8, you now get a different (genuinely Monarch) model — the block-diagonal one you had is at blockdiag_8.

2. int8: the previously published int8 ONNX never quantized the structured weights. The structure-preserving export lowers each block-diagonal / Monarch matmul to an Einsum, and onnxruntime ships no QDQ handler for Einsum — so quantize_static skipped those nodes entirely. Only activations and the residual dense MatMuls were int8; the dominant weights stayed FP32. Every int8 file here has been re-quantized with the structured weights genuinely int8 (fixed in nsnet2/qdq_einsum_quantizer.py upstream). The "int8 is loss-free" property does survive the fix — but it had never actually been tested before it.

Results (int8 = genuinely quantized weights)

Genuine two-factor Monarch

run	params	FP32 PESQ	int8 PESQ	Δ (FP32→int8)
wide_monarch	3.64 M	2.881	2.884	−0.003
monarch_8	0.55 M	2.861	2.856	+0.005
monarch_fc	2.38 M	2.843	2.831	+0.012
monarch_full	1.10 M	2.838	2.846	−0.009

Block-count sweep, both families (nblocks 5 → 40)

One knob varies. Every run below is the *_8 architecture — hidden 400, fc 600, both FCs and both GRU projections structured — with only nblocks changed, so parameters move and nothing else does. Dense baseline for reference: 2.845 FP32 / 2.834 int8 at 2.78 M.

nblocks	blockdiag params	FP32	int8	Δint8	monarch params	FP32	int8	Δint8
5	0.563 M	2.826	2.793	+0.033	0.880 M	2.852	2.858	−0.007
8	0.355 M	2.832	2.825	+0.007	0.553 M	2.861	2.856	+0.005
10	0.285 M	2.772	2.744	+0.028	0.443 M	2.849	2.842	+0.007
20	0.146 M	2.719	2.627	+0.092	0.225 M	2.849	2.854	−0.005
40	0.077 M	2.608	2.455	+0.153	0.117 M	2.837	2.837	0.000

Block-diagonal collapses as blocks narrow; Monarch does not. Over nblocks 5→40 blockdiag loses 0.218 PESQ in FP32, and its int8 penalty grows from 0.033 to 0.153 (0.338 total in int8 terms). Monarch moves 0.015 in FP32 and stays int8-loss-free throughout — exactly 0.000 at nblocks 40.

The separating variable is connectivity, not capacity. A block-diagonal factor never mixes across blocks, so raising nblocks splits the network into narrower non-communicating bands; Monarch's permutation restores full cross-channel reach in one step. Two checks:

• At matched parameters: blockdiag_5 (0.563 M) 2.826 vs monarch_8 (0.553 M) 2.861 — +0.035 for Monarch at equal size.
• Monarch wins while smaller: monarch_40 (0.117 M) beats blockdiag_20 (0.146 M) by 0.130 FP32 and 0.227 int8.

monarch_40 reaches dense parity with 24× fewer parameters (2.837 vs 2.845, inside metric noise) and is loss-free in int8, where the dense baseline itself gives up 0.012. It is the model to take unless you have a reason not to.

Param-matched dense controls

Structured-vs-structured comparisons cannot tell "Monarch's mixing is what matters" apart from "any model of this width would do". The control is a plain narrower dense NSNet2 at the same parameter count (dense_h*, hidden and fc scaled together at the original 1.5 ratio, each within ~1% of its pair):

params	block-diagonal	dense	Monarch	dense−monarch
0.88 M	—	2.806†	2.852	−0.046
0.55 M	2.826	2.840	2.861	−0.021
0.44 M	—	2.815	2.849	−0.034
0.23 M	2.719	2.784	2.849	−0.065
0.12 M	—	2.751	2.837	−0.086
0.08 M	2.608	2.749	—	—

† dense_h216 was originally published from a single seed-1234 run at 2.783, the worst of three. Repeat seeds (dense_h216_s2345, dense_h216_s3456, both in this repo) score 2.814 and 2.821, giving 2.806 ± 0.017 (n=3); the pairing's gap falls from −0.069 to −0.046. Note h216's seed sd is 2.4× that of the other dense arms, and it reaches the same training loss as the smaller h168 while scoring worse — consistent with lr=0.003 being too hot for the widest dense arm.

Monarch beats dense at all five matched sizes (mean −0.055), and the gap widens as models shrink. The small-end ordering is Monarch > dense > block-diagonal — so blockdiag_40 does not fail because 0.077 M is too small for the task; a dense model of that exact size scores 2.749 against its 2.608.

Dense quantizes cleanly at every width (|Δ| ≤ 0.017), so int8 robustness does not separate dense from Monarch; that failure is specific to narrow block-diagonal.

Confirmed with repeat seeds (three each: 1234, 2345, 3456) on the two pairings that carry the argument:

pairing	Monarch (n=3)	dense (n=3)	gap	overlap?
0.12 M	2.832 ± 0.009	2.753 ± 0.002	+0.079	none (+0.067)
0.55 M	2.858 ± 0.003	2.833 ± 0.008	+0.025	none (+0.016)

In both, the worst Monarch run beats the best dense run, and the advantage grows as models shrink. Seed noise is sd 0.002–0.009, an order of magnitude below the effect. Note that dense's scatter across sizes (0.089 band, the 0.55 M arm beating the 0.88 M one) is therefore real rather than noise — dense NSNet2 trains inconsistently at these widths in a way Monarch does not.

Latency: the trade-off runs the other way

All seven small-model arms below were re-timed back-to-back in one session on an idle box, so these RTFs are directly comparable (the RTF figures elsewhere in this file were collected across different days and are not).

params	model	int8 PESQ	int8 RTF
0.12 M	monarch_40	2.837	0.013
0.12 M	dense_h68	2.754	0.007
0.23 M	monarch_20	2.854	0.013
0.23 M	dense_h100	2.794	0.011
0.08 M	blockdiag_40	2.455	0.008
0.08 M	dense_h52	2.745	0.005

At matched parameters the dense model is faster at every small size — 1.9× at 0.12 M, 1.2× at 0.23 M. Monarch's Einsum lowering carries a fixed overhead that stops paying for itself as the models shrink; parameter count and latency are not the same axis.

So the deployment choice at 0.12 M is a real trade, not a free win: +0.083 PESQ (Monarch) against ~1.9× lower CPU latency and no torch-structured / gru-qat dependency (dense). On embedded targets the case is worse still — Monarch measured ~3× slower than dense on the RT595, and raising the block count made the STM32N6 NPU slower. Monarch wins the science; dense may well win the product.

Block-diagonal, dense, butterfly

run	params	FP32 PESQ	int8 PESQ	Δ (FP32→int8)
wide_blockdiag	2.36 M	2.864	2.847	+0.016
baseline (dense)	2.78 M	2.845	2.833	+0.012
blockdiag_8	0.36 M	2.832	2.825	+0.007
blockdiag_full	0.70 M	2.827	2.843	−0.016
blockdiag_fc	2.14 M	2.805	2.787	+0.018
butterfly_2blocks	0.36 M	2.805	2.202	+0.602
butterfly_fc	1.99 M	2.799	2.494	+0.306
butterfly_ortho	0.19 M	2.780	2.577	+0.203
butterfly_full	0.19 M	2.772	2.128	+0.644

Mixed-radix butterfly, trained natively (2026-10-17/18)

Butterfly in every layer, but with a few large orthogonally initialised blocks per stage instead of nine radix-2 levels ("kind": "mixed_radix" for the FCs, x_radices/h_radices for the GRU projections; eco8 branch nsnet2-mixed-radix-all). mr8_ortho pads every width to 512 or 1024 and uses 8·8·8 / 8·8·16; np8510_ortho and np41010_ortho use radices that factor the true widths (400 = 8·5·10 or 4·10·10, 800 = 8·10·10), so nothing is padded. One seed (1234), the same recipe as every other run here.

run	params	weight MACs	FP32 PESQ	int8 PESQ (g_best.onnx)	int8 PESQ, per-tensor NPU export
np8510_ortho	0.19 M	182,800	2.843	2.662	2.826
np41010_ortho	0.20 M	188,800	2.831	2.750	2.821
mr8_ortho	0.26 M	249,856	2.823	2.749	2.827

The g_best.onnx here is eco8's own static int8 graph, as for every other run. The last column is the same weights exported as a lean graph for the STM32N6 (padding folded into zero weights, no replication glue) and quantised per-tensor with the same calibration set: that export loses at most 0.018 PESQ to int8, against 0.07–0.18 for eco8's graph. On the N6 with the epoch controller the lean graphs run in 0.395 ms (np8510_ortho), 0.381 ms (np41010_ortho) and 0.450 ms (mr8_ortho), against 0.382 ms for monarch_20.

Key findings

• Quality saturates — this model class is architecture-bound, not capacity-bound. Across three structure families and ~10× parameters, every configuration lands in a 2.83–2.88 band, non-monotonically: the 0.55 M monarch_8 (2.861) beats both the 1.10 M monarch_full and the 2.38 M monarch_fc. Going 7× from monarch_8 to wide_monarch buys +0.020 PESQ. NSNet2 predicts a magnitude mask and reuses the noisy phase, which caps PESQ regardless of how expressive the mask predictor is. The dense model was already over-parameterized — which is exactly why aggressive structuring is nearly free. For deployment, take the smallest (blockdiag_8 / monarch_20). The block-count sweep extends this: 4× parameters across nblocks 5→20 moves PESQ by 0.012, within run-to-run noise.
• Genuine Monarch beats block-diagonal, but marginally (+0.011…+0.038 FP32 at matched nblocks) and it costs parameters — its second factor makes it larger. Consistent with the saturation above.
• Monarch quantizes loss-free at every block count tested (|Δ| ≤ 0.012 over nblocks 4–40). Block-diagonal only up to nblocks 8 (|Δ| ≤ 0.018) — at 20 and 40 the penalty is 0.092 and 0.153. The older unqualified "block-diagonal quantizes loss-free" claim was tested only on wide blocks, with the weights genuinely quantized.
• Butterfly with randn init degrades catastrophically under int8 (Δ up to 0.644). Use init=ortho: butterfly_ortho loses 0.203 to int8 vs 0.644 for butterfly_full — same architecture, same data, only the init differs.

Held-out-speaker retraining (vs_*, 2026-10)

The runs above keep the checkpoint with the best PESQ on the test split, which reads high. The vs_* runs retrain every configuration (same config.json, 200 epochs, eco8-neaixt c697f64, nsnet2.train --val_speakers p226,p287) with the best checkpoint chosen on two speakers held out of the training split (10,802 train / 770 validation utterances), so the 824 test utterances are used only to score. Folder names carry the seed: vs_<run>_s<seed>; np8510_ortho, monarch_20, dense_h148 and dense_h168 have seeds 1234–1236 (mean ± sd below). Each folder also holds best.json (selected step and validation PESQ) and run.json (command, commit, package versions).

PESQ-WB on the 824 VBD test utterances. int8 (g_best.onnx) is eco8's static int8 graph (MinMax, 200 calibration utterances); int8 (NPU export) is the graph built for the STM32N6 (weight granularity chosen on training utterances; mixed-radix butterflies through the lean export, per-tensor).

run	seeds	PESQ FP32	PESQ int8 (g_best.onnx)	PESQ int8 (NPU export)
vs_monarch_full	1234	2.828	2.825	2.819
vs_monarch_10	1234	2.832	2.820	2.817
vs_monarch_8	1234	2.837	2.817	2.808
vs_wide_blockdiag	1234	2.828	2.784	2.804
vs_monarch_20	1234,1235,1236	2.817 ± 0.023	2.808 ± 0.038	2.803 ± 0.025
vs_mr8_ortho	1234	2.800	2.779	2.803
vs_monarch_40	1234	2.810	2.772	2.800
vs_monarch_5	1234	2.807	2.821	2.799
vs_monarch_fc	1234	2.794	2.793	2.796
vs_np8510_ortho	1234,1235,1236	2.798 ± 0.031	2.718 ± 0.023	2.780 ± 0.022
vs_np41010_ortho	1234	2.795	2.703	2.772
vs_dense_h216	1234	2.793	2.790	2.765
vs_dense_h168	1234,1235,1236	2.779 ± 0.011	2.779 ± 0.014	2.754 ± 0.018
vs_blockdiag_fc	1234	2.777	2.768	2.753
vs_dense_h68	1234	2.748	2.733	2.745
vs_dense_h148	1234,1235,1236	2.771 ± 0.020	2.760 ± 0.011	2.744 ± 0.043
vs_blockdiag_full	1234	2.783	2.730	2.731
vs_dense_h100	1234	2.769	2.751	2.723
vs_blockdiag_5	1234	2.769	2.723	2.717
vs_blockdiag_10	1234	2.747	2.744	2.714
vs_blockdiag_8	1234	2.750	2.743	2.708
vs_butterfly_ortho	1234	2.785	2.574	2.691
vs_blockdiag_20	1234	2.718	2.635	2.652
vs_dense_h52	1234	2.732	2.707	2.652
vs_butterfly_fc	1234	2.781	2.558	2.646
vs_butterfly_2blocks	1234	2.809	2.141	2.375
vs_blockdiag_40	1234	2.589	2.269	2.344
vs_butterfly_full	1234	2.756	2.171	2.178

The radix-2 butterflies lose far more to int8 than any other family. Re-encoding their levels as three merged 8·8·8 stages (exact in float, fewer requantisation points) helps but does not close the gap; training the merged form natively (vs_mr8_ortho, vs_np8510_ortho) does. butterfly_2blocks is not merged (two blocks per projection).

run	PESQ FP32	PESQ int8, per-level export	PESQ int8, merged 8·8·8
vs_butterfly_ortho	2.785	2.691	2.717
vs_butterfly_fc	2.781	2.646	2.683
vs_butterfly_full	2.756	2.178	2.309

Layout

One subdirectory per run: the generator (g_best), the streaming FP32 ONNX, the static int8 ONNX, and the exact config.json it was trained with.

baseline/          blockdiag_5/     monarch_5/       butterfly_fc/
blockdiag_fc/      blockdiag_8/     monarch_8/       butterfly_full/
blockdiag_full/    blockdiag_10/    monarch_10/      butterfly_ortho/
wide_blockdiag/    blockdiag_20/    monarch_20/      butterfly_2blocks/
wide_monarch/      blockdiag_40/    monarch_40/      mr8_ortho/
monarch_fc/                                          np8510_ortho/
                                                     np41010_ortho/
vs_<run>_s<seed>/   held-out-speaker retraining (36 runs, see above)

  each: {g_best, g_best_fp32.onnx, g_best.onnx, config.json}

Loading

git clone https://github.com/LarocheC/eco8-neaixt && cd eco8-neaixt && uv sync

import json, torch
from huggingface_hub import hf_hub_download
from common.env import AttrDict
from nsnet2.model import NSNet2

REPO = "claroche1/sparse-nsnet2-checkpoints"
RUN  = "monarch_8"   # or any run from the tables

cfg  = json.load(open(hf_hub_download(REPO, f"{RUN}/config.json")))
ckpt = torch.load(hf_hub_download(REPO, f"{RUN}/g_best"),
                  map_location="cuda", weights_only=False)

model = NSNet2(AttrDict(cfg)).cuda().eval()
model.load_state_dict(ckpt["generator"])

Note for the monarch_* runs: their GRU was trained through gru-qat's fused Monarch Triton kernel ("gru": {"kind": "triton_monarch"}), so their state_dict uses gru-qat module names. Loading them requires gru-qat >= 0.5.0 and torch-structured >= 1.3.0 (both pulled in by uv sync). The blockdiag_* / butterfly_* / baseline runs use the native path and have no such requirement.

ONNX (FP32 or int8)

Streaming-shape: one frame (B, n_freq) plus GRU state (num_layers, B, hidden) per session call, threaded across frames. End-to-end pipeline in nsnet2/inference_onnx.py.

import onnxruntime as ort
from huggingface_hub import hf_hub_download
REPO, RUN = "claroche1/sparse-nsnet2-checkpoints", "monarch_8"
sess = ort.InferenceSession(hf_hub_download(REPO, f"{RUN}/g_best.onnx"),
                            providers=["CPUExecutionProvider"])   # int8

Citations

@inproceedings{braun2021nsnet2,
    title={Towards efficient models for real-time deep noise suppression},
    author={Braun, Sebastian and Tashev, Ivan},
    booktitle={ICASSP}, year={2021}
}
@inproceedings{dao2019butterfly,
    title={Learning fast algorithms for linear transforms using butterfly factorizations},
    author={Dao, Tri and Gu, Albert and Eichhorn, Matthew and Rudra, Atri and R{\'e}, Christopher},
    booktitle={ICML}, year={2019}
}
@inproceedings{dao2022monarch,
    title={Monarch: Expressive structured matrices for efficient and accurate training},
    author={Dao, Tri and Chen, Beidi and Sohoni, Nimit S and Desai, Arjun and Poli, Michael and Grogan, Jessica and Liu, Alexander and Rao, Aniruddh and Rudra, Atri and R{\'e}, Christopher},
    booktitle={ICML}, year={2022}
}