Posted on

Table of Contents

Back in March I wanted to clean noisy speech live, from a microphone, in Rust. NVIDIA has Real-time RE-USE: a streaming Mamba model that takes one 20 ms frame of audio at a time. Small, causal, made for exactly this.

I was sure it had to run on the GPU. Neural network, real time, GPU. Because that's just how it goes, right? So I exported it to ONNX and ran it through ort.

This post is how that went, and why the thing I ended up shipping, reuse-rs, is plain Rust on the CPU.

Before the story, here's what it does: the same recording, clean original, noisy, and cleaned by reuse-rs (from VoiceBank-DEMAND). The spectrogram fills in as the audio plays.

Original (clean)Noisyreuse-rs, 8 layers

The ONNX road

Mamba's core is a selective scan: a small state per channel that gets decayed and updated every step. ONNX has no op for that. The closest thing is com.microsoft.LinearAttention, a contrib op in ONNX Runtime, so I mapped the scans onto it.

It worked. Sort of. I ran into problems inside ONNX Runtime on the way and reported what I found (onnxruntime#27796, #27996).

Then the numbers, per 20 ms frame on my M3:

  • WebGPU: 16.9 ms with 4 layers, 20.6 ms with 8. That's at the edge: anything over 20 ms means you fall behind the microphone.
  • CPU (ONNX Runtime): about 70 ms.
  • CoreML: the graph got split, fell back to CPU, 168–198 ms.

I also tried NVIDIA's bigger offline model (30 bidirectional layers, their best quality). 0.19–0.33× real time on WebGPU, 3.3–3.9 GB of memory. At one point I ran several backends in the same process to compare them and watched it eat 26 GB of RAM. Lesson learned.

What if I just... don't

I never considered writing the model myself. It felt like the "serious" path was the GPU and the runtime. But the more I looked at it, the less sense it made: every step is one frame. 321 frequency bins. You can't batch the future, because a live mic hasn't produced it yet, and the model needs the state from the previous frame anyway.

So for each tiny frame, the GPU has to launch kernels, run a bit of math, and sync back. My suspicion was that the launching and syncing was most of the bill.

The WebGPU numbers already said so, I just hadn't read them that way: going from 4 to 8 layers only added 3.7 ms. Double the work, almost the same time. Most of those ~17 ms were fixed per-frame overhead, not the model.

So I wrote it in plain Rust. The convolutions are im2col plus one matrix multiply (Accelerate on macOS, the matrixmultiply crate elsewhere). The scan is the interesting part. RE-USE uses a state size of 16 everywhere, so the state is a fixed [f32; 16] and the update is three tiny loops the compiler can vectorize:

fn scan16(s: &mut [f32; 16], a: &[f32; 16], delta: f32, dx: f32,
          b: &[f32; 16], c: &[f32; 16]) -> f32 {
    let mut decay = [0f32; 16];
    for n in 0..16 {
        decay[n] = fast_exp(delta * a[n]);
    }
    let mut out = [0f32; 16];
    for n in 0..16 {
        s[n] = decay[n] * s[n] + dx * b[n];
        out[n] = s[n] * c[n];
    }
    let mut acc = [0f32; 4];
    for n in 0..16 {
        acc[n % 4] += out[n];
    }
    (acc[0] + acc[1]) + (acc[2] + acc[3])
}

fast_exp is a polynomial exp (about 1 ulp) written so it vectorizes too, because f32::exp in that loop stops the compiler. The scans for different channels are independent, so they run in parallel with rayon.

Is it actually right?

A denoiser that "sounds fine" can still be wrong in ways you don't hear. So I compared every output against a float64 PyTorch reference.

Difference from the float64 reference
reuse-rs1.5e-6
ONNX export through ort2.4e-3

1.5e-6 is about as close as PyTorch's own float32 run gets.

The numbers

Same M3 (8 cores, 16 GB), idle, time per 20 ms frame:

Layersreuse-rs (CPU)ONNX Runtime (WebGPU)
48.4 ms16.9 ms
611.5 ms17.4 ms
814.7 ms20.6 ms

Memory went from about 420 MB to about 150 MB. And the CPU numbers grow with the work (+6.3 ms from 4 to 8 layers), which is what you want: the time goes into the model, not around it.

Without Accelerate (the matrixmultiply path) 8 layers takes 21.4 ms on the same machine, so on other platforms 4–6 layers is the safer pick. And the machine has to be reasonably free: when I ran it next to a game, frame times went up.

Fewer layers?

The model was trained so you can stop after any of 3 to 12 layers, each with its own decoders. I couldn't hear a difference between 4 and 8, so I measured it on the same VoiceBank-DEMAND clip, against the clean original (SI-SDR, higher is better):

SI-SDR
Noisy4.5 dB
4 layers16.4 dB
6 layers16.6 dB
8 layers16.7 dB

0.3 dB between 4 and 8, for 1.75× the speed. That's one clip, but it matches my ears.

4 layers6 layers

Using it

use reuse_rs::Denoiser;

let mut denoiser = Denoiser::load("realtime_reuse/model.safetensors", 4)?;
// 16 kHz mono samples, in whatever chunks they arrive.
let clean = denoiser.process(&noisy_chunk)?;
// At the end of the stream:
let rest = denoiser.flush()?;

Output is delayed by exactly one frame (20 ms). I'm also adding it to GStreamer as a reusedenoise element, so it can sit in any pipeline.

Where it falls short

  • When the background noise is other people talking, it struggles. Which makes sense: from the model's point of view, that noise is speech too. More examples are in the README. If you don't need real time, NVIDIA's offline model does better, and SpeechLens runs it on Apple Silicon.
  • Only 16 kHz and no look ahead are tested.
  • The code is MIT/Apache-2.0, but the weights are NVIDIA's, under a non-commercial license.

The thing I keep thinking about: the GPU wasn't wrong, it was just the wrong shape for the job. For one stream of tiny frames, the CPU was sitting right there the whole time.


Cite:

Altun, E. (2026, October 20). I forced a speech model onto the GPU for months. The CPU was faster.. Retrieved from https://altunenes.github.io/posts/reuse/

CC BY 4.0