Parameter estimation¶
The effects are exact recurrences.
Thus you can backpropagate through them.
The classic example follows dasp-pytorch’s quickstart.
Render a target with known compressor settings.
Then recover those settings from only the audio.
The example runs in approximately one second on the CPU (pip install optax):
from flax import nnx
import jax
import jax.numpy as jnp
import optax
from faustax import Compressor
comp = Compressor(sample_rate=44100)
rngs = nnx.Rngs(0)
# A noise burst with a level drop: level changes are what make compression
# audible (and learnable). The raw module speaks (channels, samples).
x = rngs.uniform((1, 22050), minval=-1.0, maxval=1.0)
x = x * jnp.where(jnp.arange(22050) < 11025, 1.0, 0.1)
# Target: audio rendered with the settings we want to recover
y = comp.module(x, params={"threshold_db": -36.0, "ratio": 8.0})
# Estimate those two parameters in normalized [0, 1] space; parameters not
# named in the dict keep their Faust defaults automatically
def loss_fn(theta):
y_hat = comp.module(
x, normalized_params={"threshold_db": theta[0], "ratio": theta[1]}
)
return jnp.mean((y_hat - y) ** 2)
theta = jnp.array([0.5, 0.5])
opt = optax.adam(1e-2)
opt_state = opt.init(theta)
@jax.jit
def step(theta, opt_state):
loss, grads = jax.value_and_grad(loss_fn)(theta)
updates, opt_state = opt.update(grads, opt_state)
return optax.apply_updates(theta, updates), opt_state, loss
for _ in range(1000):
theta, opt_state, loss = step(theta, opt_state)
meta = comp.module.get_parameter_metadata()
recovered = comp.module.unnormalize_params(
{"threshold_db": theta[0], "ratio": theta[1]}
)
print({meta[zone]["shortname"]: round(float(v), 2) for zone, v in recovered.items()})
# {'ratio': 8.0, 'threshold_db': -36.0} — exact recovery, final loss 0.0
Two lessons apply beyond this example:
Optimize in the normalized space. In this space, each parameter has the range
[0, 1]. Thus one learning rate is sufficient for all parameters. Values that move out of range continue to receive useful gradients. The clip back to[0, 1]is a straight-through estimator (magic clamp). This estimator passes the gradient when the update moves the value back toward the valid range. Use the physicalparams=path only for values that you trust. This path applies no clip at the module level (Processor.processvalidates concrete values against the slider ranges).Identifiability is a property of the signal, not of the optimizer. A joint estimate of all five compressor parameters on this same burst also matches the audio (loss ~1e-6). But the estimate finds a different, equivalent curve (a lower threshold and a more gentle ratio). For parameter forensics, fix the parameters that you know, or use program material that excites each parameter. For style transfer and automatic mixing, a match of the sound is the goal, and this is sufficient.
Time-varying parameters (learnable automation)¶
The parameters above are constant in each call.
A time-varying parameter needs a different construction: the slider becomes an input channel, and the automation becomes ordinary audio that carries a gradient.
Faust widget modulation makes this change.
LowpassAutomated (src/faustax/dsp/lowpass_automated.dsp) is the smallest complete example:
cutoff = hslider("cutoff", 440.0, 20.0, 20000.0, 0.01);
replace = !,_;
process = ["cutoff":replace -> fi.lowpass(1, cutoff)];
["cutoff":replace -> FX] prepends one input channel to FX and routes it to the place of the cutoff slider.
The replace = !,_ composition discards the slider value and keeps the incoming channel.
Write * or + in that position instead to scale or to offset the slider value.
Each modulated widget adds one input channel, in the order of the widgets inside the brackets, before the inputs of the effect.
This DSP therefore takes 2 channels: channel 0 is the cutoff in Hz, and channel 1 is the audio.
The slider disappears from the interface, so the generated module reports zero parameters.
examples/parameter_estimation/automation.py is the reference for the JAX side, and it records the three choices that decide the result.
The automation lives at control rate (1/100th of the audio rate) and is upsampled with differentiable linear interpolation.
The normalized value maps to Hz geometrically, as a [scale:log] slider does.
magic_clamp holds the value in range and still returns a gradient that points back into the range.
The result recovers a hidden 140 Hz–5.8 kHz filter sweep to sub-cent accuracy in ~8 s on the CPU.
Run the example with --compile to compile the .dsp file at run time with faustax.compile_file(), which is the faster loop while you author a new DSP.
Learning a knob’s response curve (known control, learnable mapping)¶
This procedure applies when you have many (knob setting, audio) pairs from a reference device, with one render for each knob position.
The goal is to recover the mapping from the knob to the DSP, not the knob values, because you know the knob values.
Make the endpoints of the mapping into their own sliders with it.remap:
drive = hslider("drive", 0.5, 0, 1, 0.01); // KNOWN per item
gain_db = drive : it.remap(0, 1,
hslider("gain_min_db", 0, -24, 24, 0.01), // LEARNED
hslider("gain_max_db", 24, -24, 48, 0.01)); // LEARNED
examples/learnable_structure/learn_control_response.py recovers a hidden gain law (src/faustax/dsp/learnable_control_response.dsp) exactly.
The example sets the drive value for each batch item, and this value does not enter the parameter vector.
The example calculates the gradient only with respect to the endpoints.
One jax.vmap over (audio_item, drive_item) shares the learned endpoints across the batch.
Each item constrains the curve at its own drive value.
Together, the items recover the full mapping.
For a curved response, add more control points (more it.remap segments, or a learnable drive : pow(_, taper)).
To fit several knobs that interact (drive/tone/level) at the same time, batch over a knob grid with one target render for each cell.
Fitting over a diverse parameter space (random sampling + minibatching)¶
A batch over a knob grid is a possible extension, but the grid is sparse and its size increases combinatorially.
A grid with k points for each knob and d knobs needs k**d renders.
The fit only extrapolates to the intermediate settings.
An added input-level axis multiplies the render count again.
The scalable alternative is to sample the settings continuously at random, and fit on random minibatches.
examples/parameter_estimation/random_sampling_fit.py recovers the two learnable laws of src/faustax/dsp/learnable_two_knob.dsp (a drive pre-gain into ma.tanh, and a tone lowpass cutoff, each a known-control it.remap):
Pre-render many targets at random continuous settings. The example draws each knob
~ U[0, 1]and the input level~ Uover a dB range. Thus the fit covers the full N-dimensional volume, not only the lattice corners. You select the sample count, and this count is independent of the axis count.The input level is its own axis. The
drivegain applies before thetanh. Thus the quantity of clipping in the saturator depends on the absolute input level. A single level holds the nonlinearity at one operating point, and the drive law stays under-constrained. A variable level measures the saturator across its curve. The level is data, included in the input of each item; the level is never a learned parameter.Draw a new random minibatch at each step. This is the same procedure as SGD over a dataset. Each step is cheap and uses a different part of the space. Across the steps, the gradient averages over the full volume.
The example sets the known knob values for each item and does not optimize them.
The example optimizes only the four it.remap endpoints.
One jax.vmap renders the minibatch with the shared learned parameters, as in the previous section.
A fixed PRNG seed keeps the run deterministic.
All four endpoints recover to less than 0.1% of their slider ranges in a few seconds on the CPU.
Fitting to a real reference: the loss must be spectral¶
The three examples above fit with a paired waveform loss mean|y - t|.
This loss operates correctly only because the same Faust program with the same input renders the target.
Then y and t align at each sample.
Some references cannot receive your exact model input with a sample-aligned capture.
Examples: a plugin, hardware, a different DSP topology, or a device with its own latency, phase, or internal noise.
For such a reference, a waveform loss compares two signals that do not correspond, and the result is not usable.
Compare spectra instead.
examples/studies/audio_loss_design.py shows the failure and the correction.
The example fits the two shelf gains of ParametricEQ to a reference render.
The example supplies the model with a different noise signal than the noise signal of the reference.
This condition simulates a situation in which you only have unaligned audio through each device.
The paired waveform loss pushes the gains to the ends of the slider ranges.
A multi-resolution STFT magnitude loss recovers the gains to approximately one dB.
This occurs because the magnitude spectra depend on the EQ curve and not on the specific noise:
def multiscale_stft_loss(y, t): # sum over several FFT sizes
total = 0.0
for n in (512, 1024, 2048):
a, b = stft_mag(y, n), stft_mag(t, n)
sc = jnp.sqrt(jnp.sum((a - b) ** 2)) / (jnp.sqrt(jnp.sum(b**2)) + 1e-6) # level
logmag = jnp.mean(jnp.abs(jnp.log(a + 1e-5) - jnp.log(b + 1e-5))) # balance
total += sc + logmag
return total / 3
Copy this loss for each fit to a real reference. The list below gives two design points in the loss and three caveats around the loss. All these points come from a fit of an overdrive pedal to its first-party plugin:
The log-magnitude equalizes the bands. With
|log(A) - log(B)|, a dB error has the same weight in a quiet top band as in a loud low band. The loudest region dominates a linear magnitude MSE. Then, if the model cannot match the reference exactly, the model uses its capacity on the loud low end. The quiet top band becomes dull. Spectral convergence constrains the level; the log-magnitude constrains the balance.Use several FFT sizes at the same time. Short windows find the time position of transients. Long windows resolve partials that are close in frequency. A sum over the sizes prevents an overfit to one axis.
The magnitude discards the phase. Thus the magnitude cannot constrain the properties that are only in the phase. Examples: the sign of the asymmetry of a nonlinearity, and the stereo or time alignment. (A mirror flip of an asymmetric clipper changes its harmonic magnitudes only a small amount.) When these properties are important, add a cheap signed summary that the magnitude loss cannot see. Examples: the waveform skewness (3rd moment) for the asymmetry, and the spectral centroid for the brightness that the shelves can exchange. Use the summary as a strong regularizer, not as the main loss.
Match each type of material separately, not in aggregate. A probe can mix transient and sustained material (drum hits + a held note). Then match the level of each segment separately. In an average across the segments, the sustained part dominates the crest of the transient, and the fit under-drives the saturator.
Know the limit of the model. A loss term can stop at a floor value, even when its weight makes the term dominate the loss. This limit comes from the capacity of the model, not from the optimizer. A physical model possibly cannot match the harmonic character of a reference. No loss weight can give the model a capability that the topology does not have. A diagnosis of this limit is itself a result. Do not continue the optimization against this limit.
Recovering an instrument’s physical parameters (modal synthesis)¶
The physical models of Faust also compile through the NNX backend.
Thus the synthesis side of a fit can be a real instrument.
src/faustax/dsp/modal_piano.dsp is a differentiable modal piano string: a bank of pm.modeFilter resonators that a hammer impulse strikes.
Physical quantities parameterize the model: the fundamental f0, the string inharmonicity (the partials stretch by √(1 + B·k²)), the decay t60_0 and its per-partial falloff, and the strike brightness.
Modal resonators differentiate cleanly, because their biquad coefficients are smooth in frequency and decay; the delay length of a waveguide is not smooth.
Thus examples/parameter_estimation/fit_modal_piano.py recovers those quantities exactly from the sound by gradient descent (the inharmonicity, decay, and brightness all to <1% at a known pitch):
from faustax.modules import ModalPiano
model = ModalPiano(sample_rate=8000)
y = model(num_samples=4000, normalized_params={"inharmonicity": ..., "t60_0": ...})
Two conditions keep the fit well-posed.
The window must show the decay, so that t60 is visible.
The signal must have sufficient high-partial energy, so that the inharmonic stretch is visible.
The pitch is known and fixed, because the played note gives the pitch.
Pass --plot out.png for an A/B spectrum and spectrogram (this option needs the viz extra, pip install faustax[viz]).
This model is the Faust-authored equivalent of ops.diag_state_space (a modal bank is a set of complex-pole resonators).
The compiler capstone (docs/source/future-work.md) will route such tf2 recurrences through ops.allpole.
Then these fits will get the memory decrease of the custom VJP with no extra work.
Saving trained settings¶
You can write the learned values back into the nnx.Param leaves of the module.
Then you can save the values as a safetensors preset.
The file records the DSP identity (the class name plus the zone-to-label map).
Thus a load of the file into an incorrect effect causes an exception; the load does not assign the parameters incorrectly without a warning.
The code below continues from the estimation example:
from faustax.realtime import RealtimeEffect
# Write the learned normalized values into the module state...
for name, value in zip(["threshold_db", "ratio"], theta):
getattr(comp.module, comp.module.label_to_zone(name))[...] = value
comp.module.save_params("compressor_preset.safetensors")
# ...and later, e.g. in the live rig:
live = Compressor(sample_rate=48000)
live.module.load_params("compressor_preset.safetensors")
fx = RealtimeEffect(live, block_size=512) # streams with the trained curve
(The copy step is necessary because the loop above optimized an external theta.
NNX-native training with nnx.grad updates the nnx.Param leaves of the module in place, and then save_params reads them directly.)