--- license: mit library_name: pytorch tags: - speech-enhancement - speech-denoising - nsnet2 - sparsity - semi-structured-sparsity - 2:4-sparsity - pruning - sparse-matmul datasets: - JacobLinCool/VoiceBank-DEMAND-16k metrics: - pesq --- # NSNet2 under semi-structured sparsity Six NSNet2 speech-enhancement checkpoints trained under **fixed semi-structured sparsity masks** — 2:4, 4:8, 1:4, 80% 1×4 blocks, 80% unstructured, and a dense control — for work on sparse-dense MatMul packing and code generation at **batch 1**. Each ships as a PyTorch checkpoint, an FP32 and a static-int8 ONNX graph, and a numpy export with explicit zeros and masks. No pattern here costs measurable quality, in FP32 or int8. Code, training recipe and export tooling: [LarocheC/eco8-neaixt](https://github.com/LarocheC/eco8-neaixt), branch `sparse-masks-rowfusion`. See `SPARSE_MATMUL_COLLAB.md` there for the full method. ## Results PESQ on the full 824-utterance VoiceBank-DEMAND test set. Every arm was fine-tuned from the same dense baseline on an identical schedule (lr 3e-4, 120 epochs), so the mask is the only variable. | directory | pattern | sparsity | PESQ | | ------------------- | --------------- | -------: | ----: | | `dense` | dense (control) | 0% | 2.777 | | `2_4` | 2:4 | 50.0% | 2.779 | | `4_8` | 4:8 | 50.0% | 2.779 | | `1_4` | 1:4 | 75.0% | 2.781 | | `unstructured_80` | unstructured | 80.0% | 2.776 | | `1x4_80` | 1×4 blocks | 80.0% | 2.770 | ### int8 Static int8 PTQ (QDQ, per-channel symmetric weights, MinMax calibration on 200 utterances), PESQ through onnxruntime on the same test split. Δ is int8 − FP32. | directory | sparsity | FP32 | int8 | Δ | int8 RTF | | ------------------- | -------: | ----: | ----: | -----: | -------: | | `dense` | 0% | 2.777 | 2.783 | +0.006 | 0.121 | | `2_4` | 50% | 2.779 | 2.781 | +0.002 | 0.125 | | `4_8` | 50% | 2.779 | 2.790 | +0.011 | 0.122 | | `1_4` | 75% | 2.781 | 2.784 | +0.003 | 0.123 | | `1x4_80` | 80% | 2.770 | 2.779 | +0.009 | 0.124 | | `unstructured_80` | 80% | 2.776 | 2.774 | −0.002 | 0.121 | **Sparsity does not make quantization harder** — every Δ is inside the ±0.01 noise band at every sparsity level, and five of six are positive. **The mask survives int8 bit-exactly.** Symmetric per-channel weight quantization maps 0.0 to exactly 0. The N:M arms conform in the int8 graph with sparsity slightly *above* target (0.5016 / 0.5011 / 0.7510 — a few small weights round to zero, which N:M permits), and `1x4_80` holds block support at exactly 0.2000 live against its 0.2000 budget. Check it yourself with `nsnet2/verify_int8_sparsity.py` from the repo. **But the sparsity buys no speed today.** int8 RTF is 0.121–0.125 across every arm, dense and 80%-sparse alike, and the int8 file is 2.78 MiB regardless — onnxruntime stores the zeros explicitly and multiplies by them like any other weight. 80% of the multiplies are gone mathematically and none of the latency is. Closing that gap is what these checkpoints are for. **None of these patterns costs measurable quality.** The spread across all six arms is 0.012 PESQ while the run-to-run variation within a single arm is ~0.010 sd, so they are statistically indistinguishable. Do not read an ordering into the table — 1:4 topping it at 75% sparsity is which validation happened to land last, not a result. One caveat. Every arm including the dense control sits ~0.07 below the [published 200-epoch baseline](https://huggingface.co/claroche1/sparse-nsnet2-checkpoints) of 2.845, because these were shortened fine-tunes with a freshly initialised discriminator; the comparison between arms is unaffected since all paid the same penalty, and all six curves were still rising at epoch 120. A full-length run would likely lift every arm. For reference, magnitude pruning *without* fine-tuning is far more pessimistic: 2:4 costs 0.378 PESQ and 1×4 at 80% costs 0.656. Almost all of it comes back, so pruning-only numbers are a poor guide to what a pattern actually costs. ## Layout Each directory holds both a runnable checkpoint and a kernel-oriented export: * `g_best`, `config.json` — PyTorch checkpoint, loadable with the `NSNet2` model in the repo above. * `g_best_fp32.onnx`, `g_best_int8.onnx` — the streaming graph in FP32 and in static int8 (QDQ). The int8 graph preserves the sparsity pattern exactly. * `weights.npz` — per matrix: `.weight` (float32, dense **with explicit zeros**), `.mask` (uint8, 1 = kept), `.bias`, and golden vectors `.ref_x` / `.ref_y` where `ref_y = W @ ref_x + bias`. * `manifest.json` — shapes, pattern, grouping axis, achieved sparsity, ragged tail counts, and N at inference vs training. `verify.py` at the top level needs only numpy: ```bash python verify.py 2_4 # shapes, mask/weight agreement, pattern conformance, # and the golden vectors ``` ## Conventions **Layout.** Every weight is row-major `(M, K)`, used as `y = W · x + b` with `x` of shape `(K, N)`. **N = 1 at deployment.** The model runs one 16 ms frame at a time, so each of these is a matrix-vector product. During training N is 256 · T. **Grouping runs along K.** For an N:M pattern the groups of M are contiguous *within a row* — along the input dimension, contiguous in memory for a row-major `(M, K)` array. This matches the NVIDIA 2:4 convention. The masking code supports grouping along the output dimension too, if a kernel wants that. **Ragged tail.** `fc_in` has K = 257 — 64 groups of 4 plus one leftover column, left dense — so it measures 49.8% sparse rather than exactly 50%. `manifest.json` reports `tail_elements` per matrix. **GRU gate packing.** `gru.weight_ih_l*` and `gru.weight_hh_l*` are `(3H, K)`: PyTorch stacks the r/z/n gates along the **output** dimension, so each gate is a contiguous block of rows and a group of 4 along K never straddles a gate boundary. Each gate submatrix independently satisfies the pattern, so a `1200×400` packs as one matrix or as three `400×400` with identical results. ## The matrices The four GRU matrices are 69% of the weights and run once per frame, so they dominate. `gru.weight_hh_l0` and `gru.weight_hh_l1` sit inside the recurrence and cannot be batched over time even in principle — the strictest N=1 case here. | matrix | M | K | params | | ------------------- | ---: | --: | ------: | | `fc_in` | 400 | 257 | 102,800 | | `gru.weight_ih_l0` | 1200 | 400 | 480,000 | | `gru.weight_hh_l0` | 1200 | 400 | 480,000 | | `gru.weight_ih_l1` | 1200 | 400 | 480,000 | | `gru.weight_hh_l1` | 1200 | 400 | 480,000 | | `fc1` | 600 | 400 | 240,000 | | `fc2` | 600 | 600 | 360,000 | | `fc_out` | 257 | 600 | 154,200 | ## Usage Kernel work — numpy only, no PyTorch: ```python import json import numpy as np npz = np.load("2_4/weights.npz") W = npz["gru.weight_hh_l0.weight"] # (1200, 400) float32, explicit zeros b = npz["gru.weight_hh_l0.bias"] # (1200,) x = npz["gru.weight_hh_l0.ref_x"] # (400,) float32 assert np.allclose(W @ x + b, npz["gru.weight_hh_l0.ref_y"], atol=1e-4) ``` Running the model: ```python import json import torch from common.env import AttrDict from nsnet2.model import NSNet2 h = AttrDict(json.load(open("2_4/config.json"))) model = NSNet2(h) model.load_state_dict(torch.load("2_4/g_best", map_location="cpu")["generator"]) ``` Reproducing a mask, or training a new one: ```bash python -m nsnet2.train --config configs/ov_2to4.json \ --checkpoint_path cp_ov_2to4 --init_from ``` ## Related * [claroche1/sparse-nsnet2-checkpoints](https://huggingface.co/claroche1/sparse-nsnet2-checkpoints) — the same model under Butterfly / block-diagonal / Monarch **structured factorizations** (a different kind of sparsity: factorized transforms rather than masked dense matrices), plus the dense 2.845 baseline these were fine-tuned from. ## Citation NSNet2: Braun & Tashev, *Towards efficient models for real-time deep noise suppression*, ICASSP 2021. Training recipe built on [MP-SENet](https://github.com/yxlu-0102/MP-SENet).