The SAME Soft-Normalisation Bottleneck

A walkthrough of the SAME autoencoder’s bottleneck — five operations that shape a latent space for downstream diffusion, and the design reasoning that may have led the authors away from the VAE.

The previous post covered how SAME compresses stereo audio — a parameter-free patching step and a Transformer Resampling Block achieve 4096× temporal compression without strided convolutions. This post picks up at the narrowest point of the autoencoder, between encoder and decoder, where the latent space gets its shape. Where Stable Audio Open — the predecessor from the same team — used a standard VAE bottleneck at d=64 and ~2048× compression, SAME doubles the compression and quadruples the latent dimension to d=256 — a combination that perhaps made the VAE’s known scaling challenges more acute, and pointed toward a different kind of bottleneck.

Why a new bottleneck?

The paper doesn’t spell out why this combination called for a different bottleneck, but a plausible line of reasoning emerges from the design choices and the references they cite. Higher compression means each latent vector must represent more audio — roughly 93 ms per timestep — so a larger d is needed to avoid losing information. But scaling a VAE to d=256 is not straightforward: the KL divergence is summed over all 256 channels, the encoder must predict 512 distribution parameters, and the risk of posterior collapse — where some channels become input-independent, carrying no information about the audio, and therefore unreachable by any auxiliary supervision — grows with dimensionality. Recent work in the image domain — which the SAME paper cites — showed that simpler, non-VAE bottlenecks can work well at high latent dimensionality, provided the latent space has good semantic structure (Zheng et al., Heek et al.).

That last condition is, I think, the key to the whole design. SAME isn’t just a better-normalised autoencoder — it’s a semantically-aligned one (the name says it). The bottleneck provides a simple, well-scaled foundation, and a set of auxiliary losses add specific musical structure: a flow-matching alignment loss that trains the latent to be generable by a small diffusion model, chroma regressors that encode pitch content, an interaural level difference (ILD) probe that encodes stereo image, and a contrastive critic that aligns the latent with both audio features and text embeddings. This is a significant departure from general-purpose codec autoencoders like EnCodec, DAC, or Stable Audio Open, which learn whatever features minimise reconstruction error and leave the latent geometry to the VAE’s KL term. SAME instead takes direct control of what the latent means — and the soft-normalisation bottleneck is likely what makes that possible, by providing a simpler foundation that doesn’t fight the auxiliary losses for control of the geometry.

The auxiliary losses are a topic for a future post. This one focuses on the bottleneck itself: five operations, each familiar on its own, whose composition produces a latent space ready to receive semantic structure.

Where the bottleneck sits

After the encoder TRB has done its work (see the previous post), what arrives at the bottleneck is a 2D tensor of shape (d × T′) — for a 4.46-second training clip, that’s (256 × 48). The bottleneck’s job is to ensure this tensor has the right statistical properties before being passed to the decoder during training, or consumed by a downstream diffusion model during generation.

The pipeline at a glance

The bottleneck consists of five operations applied sequentially. Three transform the tensor (per-channel affine, EMA division, and rescaling); one is a training-only loss (the dual-axis KL) that shapes the latent through gradients without altering it; and one injects Gaussian noise — substantial during training, near-zero at inference:

The structure is a normalise → process → rescale sandwich, and the steps are tightly interdependent. Steps 1 and 2 together normalise — a learnable per-channel affine followed by a division by a single global running standard deviation — producing a latent where every channel sits at roughly unit scale, with the KL loss and the auxiliary losses all backpropagating through the normalisation into the affine and the encoder. Step 4 injects controlled noise into the normalised latent. Step 5 closes the sandwich, multiplying back by the running std to restore the decoder’s expected scale.

The encoder output: a tensor with a scale problem

Before diving into each step, it’s worth seeing the concrete problem they exist to solve. The encoder output z is a 2D tensor of shape (d × T′). Each row is a latent channel — a learned feature dimension — and each column is a temporal position representing roughly 93 ms of audio. These 256 channels are the output of a final linear projection from the TRB’s transformer dimension (1536) down to the latent dimension (256) — a raw matrix multiply with no activation or normalisation after it. Different rows of the weight matrix can have very different norms, and the transformer features they multiply can have very different distributions, so the 256 output channels land at whatever scale the optimiser finds useful for reconstruction. Consider a simplified (d=4, T′=6) example:

The bottleneck’s first job is to eliminate this heterogeneity — steps 1 and 2 do this directly, with the per-channel affine handling the per-channel differences and the EMA division correcting any residual global scale drift. Its second job is to shape what remains: step 3 applies a KL-like regularisation loss that penalises deviation from zero mean and unit variance, and the auxiliary losses add semantic structure on top. Without the first job, the second becomes intractable — if step 3 sees channel 1 at σ=33 and channel 3 at σ=0.22, it produces gradients that differ by orders of magnitude across channels, a badly conditioned optimisation. And the downstream diffusion model, operating over the full tensor jointly, would have to cope internally with features whose natural scales are wildly different.

The five steps: interactive walkthrough

The example below tracks a (d=4, T′=6) tensor through all five steps. All values — the raw tensor, the affine parameters, σ_ema, and the noise — are illustrative, chosen to make the effect legible rather than derived from a real model. Selecting each stage reveals the transformation and the resulting tensor state:

Step 1: Per-channel affine transform

Here γ and β are vectors of length d — one scale and one bias per channel. Each γ value multiplies every element in its channel’s row; each β value is added to every element. The operation is a broadcast: no parameters depend on the temporal dimension, so the transform works identically for any audio length. The tensor shape is unchanged: (d, T′) → (d, T′).

The sample values of γ and β illustrate this effect. For our toy tensor, γ = [0.8, 0.03, 0.15, 4.0] and β = [-0.2, 1.5, 0.0, -0.5]. The tiny γ₁ = 0.03 shrinks channel 1’s values by ~33×, while the large γ₃ = 4.0 amplifies channel 3’s. The worked cell for ch 1, t=0: z′[1,0] = 0.03 × 42 + 1.5 = 2.76. After the affine, all four channels live in roughly the same range of ±1 to ±3.

A natural question: what forces the network to learn these particular γ and β values? Nothing in this step alone — the pressure comes from the KL regularisation loss (step 3), which penalises deviation from unit variance in the normalised latent, and from the auxiliary losses, which also backpropagate gradients into γ, β, and the encoder. If γ is too large, KL loss goes up; if γ is too small, reconstruction loss increases. The parameters settle wherever those pressures balance.

Step 2: Running standard deviation division

The result is the normalised latent z″ — the representation that the regularisation loss and noise injection operate on. It’s important to note that the EMA is updated across training iterations, not across the audio sequence, and that σ_ema is a single global scalar — not a per-channel vector. The implementation calls x.std() without a dimension argument, computing one standard deviation across all batch elements, channels, and timesteps simultaneously, then blends it into the running estimate with a 0.999/0.001 decay. That one number then divides every element of the tensor uniformly (L32). At inference, σ_ema is frozen.

Why have both the affine and the σ_ema division? The paper doesn’t say, but the implementation suggests a natural division of roles: the affine (bottleneck.py:9–10) is a per-channel learnable parameter — each channel gets its own γ and β, trained by gradients to bring that channel to a useful scale. σ_ema (L19–20) is a single global scalar, updated each step without gradients to track the overall scale of the whole tensor. So the affine handles per-channel correction; the EMA handles global drift in whatever comes out of the affine. There’s also a practical inference-time reason: at inference there is no batch, so a frozen scalar is needed for the division step — the EMA is stored as a non-gradient parameter (requires_grad=False) and provides exactly that.

Step 3: Dual-axis KL regularisation

The paper calls this a “KL-like” loss rather than a KL divergence, and the distinction is worth pausing on. In a VAE, the encoder doesn’t produce a latent directly: it produces a mean μ(x) and standard deviation σ(x) that parameterise a distribution q(z|x) = N(μ, σ²), and the latent z is sampled from it via the reparameterisation trick (z = μ + σ ⋅ ε). The VAE’s KL term then measures the divergence between that predicted distribution and the prior N(0, I), computed analytically from the predicted parameters.

In SAME, the encoder produces z directly — a single deterministic point per input, no distribution parameters, no sampling. What the “KL-like” loss measures instead are the empirical statistics of z″: the mean and variance of the values in each channel computed over the time axis, and vice versa over the channel axis. The function f(μ, σ²) = μ² + σ² − log(σ²) − 1 is the same as the VAE’s per-dimension KL divergence up to a factor of ½, but here it is applied to observed batch statistics rather than encoder-predicted distribution parameters. Same mathematical form; different meaning — hence “KL-like.”

With the latent normalised, this loss encourages its statistics to match a standard normal distribution. It has two terms, one per axis:

where (μ_t, σ_t²) are per-channel mean and variance computed over time (along each row), and (μ_c, σ_c²) are per-timestep mean and variance computed over channels (down each column). It is a training loss that shapes the encoder’s learning by backpropagating gradients through steps 2 and 1, ultimately adjusting γ, β, and the encoder weights.

The penalty function

The atomic building block is the function f(μ, σ²) = μ² + σ² − log(σ²) − 1, which is zero if and only if μ = 0 and σ = 1, and grows in every other direction. If σ > 1 (too spread out) the σ² term dominates quadratically; if σ < 1 (too compressed) the −log(σ²) term explodes logarithmically; if μ ≠ 0 the μ² term adds a quadratic penalty. The only escape is μ = 0, σ = 1 — exactly what “well-scaled for downstream diffusion” means.

Two axes, two failure modes

The loss is computed along two axes independently, and neither alone is sufficient. Each axis catches a different failure mode that the other would miss. Pick a preset: the mini bar charts on the left sketch each channel across time; row statistics on the right show what the time-axis term sees (per-channel μ_t and σ_t); column statistics at the bottom show what the channel-axis term sees (per-timestep μ_c and σ_c). Cells are tinted green where the statistics look standard-normal and red where the loss is unhappy:

The correlated spike is a plausible pattern in audio: a loud broadband transient could trigger many channels simultaneously. Each row still has μ ≈ 0 and σ ≈ 1 — the time-axis loss is happy — but at the spike timestep the channel-axis loss sees a column whose mean has climbed to μ_c ≈ 2. The hidden channel bias is the reverse: channels 0 and 1 carry persistent offsets (μ_t ≈ ±1.3), which the time-axis loss flags on their rows — but the two offsets cancel at each timestep, leaving column statistics that look like N(0, 1), so the channel-axis loss sees nothing wrong.

Step 4: Gaussian noise injection

After normalisation and regularisation, a small amount of Gaussian noise is added to the latent:

Controlled Gaussian noise is added to the latent during both training and inference. The reason lies in what the decoder faces at inference. It never sees a clean encoder output at that point: it decodes whatever the diffusion model generated, which lands close to a real latent but not exactly on one — it carries small approximation errors. A decoder trained only on exact encoder outputs can be brittle to this, turning a slightly-off input into noticeably wrong audio. Adding noise during training is the fix: the decoder spends training seeing perturbed latents, so it learns to map a whole neighbourhood of nearby points to essentially the same audio and treats the diffusion model’s small errors as harmless.

The two noise scales differ by 50×. During training σ_noise = 5 × 10⁻² (a value the paper states without justification), large enough to define that neighbourhood. At inference it drops to 10⁻³ — small enough to be effectively off. Why keep any noise at all rather than setting it to zero? The paper states the floor but doesn’t justify it, and at this scale the difference from zero is unlikely to matter; it’s best treated as a genuinely open detail rather than a load-bearing design choice.

One implementation detail is worth noting, because it is what makes a single noise scale valid. In the code the noise is added after rescaling (in the decode method), not in normalised space — but because it is scaled by σ_ema (noise = randn() × σema × scale), the effect is equivalent to adding flat noise of amplitude scale in normalised space. This works precisely because normalisation put every channel on the same scale: one noise distribution can be drawn and added to every entry of the time×channel matrix with the same relative effect. Skip normalisation and no single scale fits — that fixed noise would be negligible for the ±33 channel from earlier and destructive for the ±0.22 one. That is why normalisation has to come first.

Step 5: Rescale for the decoder

The final step reverses the normalisation from step 2 by multiplying by the same running standard deviation. This restores the latent to the magnitude range the decoder expects (roughly the affine output of step 1, since the noise perturbation is small relative to the tensor scale). The values aren’t exactly the originals — they carry the small perturbation from noise injection — but they’re at the right scale. Tracing a single value through all five steps with σ_ema = 1.1 (one scalar for the whole tensor):

Putting it all together

The soft-normalisation bottleneck replaces the VAE’s learned-distribution-plus-sampling mechanism with a more direct pipeline: normalise the encoder’s output deterministically (affine + σ_ema), apply a regularisation loss to shape the statistics (dual-axis KL), add controlled noise so the decoder tolerates nearby latents, and rescale for the decoder. As a foundation it has several appealing properties: it sidesteps posterior collapse (there’s no “posterior” to collapse); the normalisation mechanism (affine + EMA) and the regularisation target (KL loss) are separate components that can be tuned independently; and the KL loss directly pressures the latent toward the roughly Gaussian marginals that diffusion models assume.

The ablation in the SAME paper does, however, reveal a critical nuance: the soft-norm bottleneck alone, without the auxiliary losses (flow-matching alignment, semantic regression, contrastive alignment), underperforms a VAE bottleneck on generation metrics. It is the combination of the simpler bottleneck and the auxiliary losses that produces the best results. The bottleneck provides a clean, well-scaled canvas; the auxiliary losses paint the semantic structure onto it. That interplay is what makes the latent space actually work for music generation — and is where the next post in this series will pick up.

The SAME paper, model weights, and bottleneck implementation are available online.