Latent Scaling Round Trip

mediumFREEGuidance and Latent DiffusionDiffusion ModelsDiffusion Models 1

Problem Statement

Implement the latent normalisation that sits between Stable Diffusion's VAE and its U-Net, including the empirical derivation of the famous 0.18215 constant.

Background

A latent diffusion model runs the diffusion process on VAE latents, not pixels. But the forward process assumes unit-variance data: the schedule is built so that $x_t$ has variance $\bar{\alpha}_t \cdot \mathrm{Var}(x_0) + (1-\bar{\alpha}_t)$ , which only lands on 1 at $t = T$ if $\mathrm{Var}(x_0) = 1$ . A raw VAE latent has a standard deviation of roughly 5.5, so it must be normalised first:

$z_{\text{scaled}} = (z - \text{shift}) \cdot s, \qquad z = \frac{z_{\text{scaled}}}{s} + \text{shift}$

The scaling factor is estimated once from the training set as $s = 1/\mathrm{std}(z)$ -- which for SD 1.x came out at 0.18215, i.e. a latent std of about 5.49. SD 3 and Flux add a non-zero shift_factor to centre the latents too.

The pipeline is: encode, scale, diffuse; then at the end unscale and decode. Forgetting either half is a classic bug -- skip the scale and the U-Net sees data at 5x the variance it was trained on; skip the unscale and the VAE decodes a washed-out grey image.

Your Task

Implement three functions:

def compute_scaling_factor(z):
def scale_latents(z, scaling_factor, shift_factor=0.0):
def unscale_latents(z_scaled, scaling_factor, shift_factor=0.0):

compute_scaling_factor returns a Python float, the other two return arrays shaped like their input.

Input Format

z, z_scaled: NumPy arrays of any shape.
scaling_factor (float): non-zero.
shift_factor (float): defaults to 0.0.

Output Format

A float, and two arrays.

Sample

z = np.arange(-4.0, 5.0)
f = compute_scaling_factor(z)
print(round(f, 4))
print(round(float(np.std(scale_latents(z, f))), 4))

The std of [-4 ... 4] is 2.582, so the factor is 1/2.582 = 0.3873, and scaling by it makes the std exactly 1.

Example

Input

z = np.arange(-4.0, 5.0)
f = compute_scaling_factor(z)
print(round(f, 4))
print(round(float(np.std(scale_latents(z, f))), 4))

Output

0.3873
1.0

Reasoning

np.arange(-4.0, 5.0) has mean 0 and population variance 60/9 = 6.6667, so its std is 2.582 and the scaling factor is 1/2.582 = 0.3873. Multiplying the (already centred) latents by that factor makes their std exactly 1.0 -- which is the whole reason Stable Diffusion carries the constant 0.18215.

Constraints

compute_scaling_factor uses the population standard deviation over all elements (np.std with its default ddof=0) and returns 1 / std as a plain Python float.
The shift is subtracted before multiplying, and added back after dividing.
unscale_latents(scale_latents(z, s, h), s, h) must round-trip to z.
Do not round inside any function.

Learn this in

Latent Diffusion & Stable DiffusionDiffusion Study Plan

More from Diffusion Models 1

Classifier-Free Guidance: Uncond/Cond Mixingeasy DDPM Posterior Mean and Variance from x0medium Adaptive CFG Scheduling over the Trajectoryhard Classifier Gradient Guidance Stepmedium Negative Prompt as the Unconditional Branchmedium

Concept primer

1. Background Knowledge

Latent Diffusion Models (LDMs) operate in a compressed latent space rather than pixel space to reduce computational cost. The core assumption of the diffusion process is that the input data follows a standard normal distribution $\mathcal{N}(0, 1)$ . However, the outputs of a Variational Autoencoder (VAE) encoder are typically not unit-variance. For Stable Diffusion 1.x, the raw latents have a standard deviation of approximately $5.49$ . If these raw latents are fed directly into the U-Net, the noise schedule (which assumes unit variance) will be misaligned, causing training instability or failure.

To bridge this gap, a normalization step is applied. This involves two parameters: a scaling factor ( $s$ ) and a shift factor (often called shift_factor or mean subtraction). The scaling factor is derived empirically from the training dataset. Specifically, $s = 1 / \sigma_z$ , where $\sigma_z$ is the standard deviation of the latent space. For SD 1.x, this results in $s \approx 0.18215$ . The shift factor centers the data by subtracting the mean, ensuring the distribution is centered at zero before scaling.

The transformation is linear: $z_{\text{scaled}} = (z - \text{shift}) \cdot s$ . The inverse operation, used during decoding, is $z = z_{\text{scaled}} / s + \text{shift}$ . This ensures that the data entering the diffusion process has unit variance and zero mean, satisfying the theoretical requirements of the forward diffusion process.

2. Algorithm Approach

The problem requires implementing three distinct functions based on basic statistical operations and linear algebra.

compute_scaling_factor: This function calculates the inverse of the standard deviation of the input array. It treats the input as a sample from a distribution and computes its empirical standard deviation.
scale_latents: This function applies the forward normalization. It subtracts the shift_factor from the input and multiplies by the scaling_factor.
unscale_latents: This function applies the inverse normalization. It divides by the scaling_factor and adds the shift_factor.

The approach relies on NumPy's vectorized operations. Since the operations are element-wise, they can be applied to arrays of any shape without explicit loops. The key is to ensure that the standard deviation calculation in compute_scaling_factor uses the correct degrees of freedom (typically population std for this specific constant derivation, but check if ddof=0 or ddof=1 is implied by the sample output; the sample output suggests ddof=0 or large N approximation, but standard np.std defaults to ddof=0).

3. Step-by-Step Strategy

**Implement **compute_scaling_factor(z)****:

Calculate the standard deviation of the input array z using np.std(z).
Return the reciprocal: 1.0 / std_z.
Note: Ensure the result is a Python float.

**Implement **scale_latents(z, scaling_factor, shift_factor=0.0)****:

Subtract shift_factor from z.
Multiply the result by scaling_factor.
Return the resulting array.

**Implement **unscale_latents(z_scaled, scaling_factor, shift_factor=0.0)****:

Divide z_scaled by scaling_factor.
Add shift_factor to the result.
Return the resulting array.

Verification:

Test with the provided sample. If z = np.arange(-4.0, 5.0), the std is $\approx 2.582$ . The factor should be $\approx 0.3873$ .
Verify that np.std(scale_latents(z, f)) is approximately $1.0$ .

4. Common Pitfalls

Degrees of Freedom: np.std defaults to ddof=0 (population standard deviation). Some statistical contexts use ddof=1 (sample standard deviation). The problem description implies the empirical constant 0.18215 derived from a large dataset, where the difference is negligible, but for small arrays like the sample, ddof=0 is the standard NumPy behavior and matches the sample output.
Order of Operations: In scale_latents, ensure you subtract the shift before scaling. In unscale_latents, ensure you divide before adding the shift. Reversing these leads to incorrect normalization.
Data Types: Ensure the output of compute_scaling_factor is a Python float, not a NumPy scalar, if strict type checking is applied. Use float() to cast if necessary.
In-place Modification: Avoid modifying the input array z in-place. NumPy operations like z * s return a new array, but be cautious if using augmented assignment like *z = s.

5. Time & Space Complexity

Time Complexity:
compute_scaling_factor: $O(N)$ , where $N$ is the number of elements in z, to compute the mean and variance.
scale_latents and unscale_latents: $O(N)$ , as each element is processed once for arithmetic operations.
Space Complexity:
compute_scaling_factor: $O(1)$ auxiliary space, as it returns a single scalar.
scale_latents and unscale_latents: $O(N)$ space to store the output array of the same shape as the input.

The operations are highly optimized in NumPy, leveraging vectorization for efficient execution on large latent tensors.

solution.py

Python 3.13

Test results