claroche1/nsnet2-sparse-rowfusion

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, 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 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: <name>.weight (float32, dense with explicit zeros), <name>.mask (uint8, 1 = kept), <name>.bias, and golden vectors <name>.ref_x / <name>.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:

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:

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:

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:

python -m nsnet2.train --config configs/ov_2to4.json \
    --checkpoint_path cp_ov_2to4 --init_from <dense g_best>

Related

• 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.