README.md

---
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: `<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:

```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 <dense g_best>
```

## 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).