📘

Guidance Rescaling

HardDiffusion Models 📚 Diffusion Models 1

Problem Statement

Fix the over-exposure artefact of high guidance scales by rescaling the guided prediction back to the standard deviation of the conditional one, then blending.

Background

Extrapolating with $w \gg 1$ inflates the magnitude of the noise prediction: $\mathrm{std}(\hat{\varepsilon}_{\text{cfg}})$ grows roughly linearly in $w$ . Since the sampler treats that output as a unit-variance noise estimate, the extra magnitude is spent pushing the sample toward the extremes of the data range -- the blown-out, over-saturated look of high-CFG images.

Common Diffusion Noise Schedules and Sample Steps are Flawed (Lin et al., 2024) fixes this in two moves. First renormalise:

\hat{\varepsilon}_{\text{rescaled}} = \hat{\varepsilon}_{\text{cfg}} \cdot \frac{\mathrm{std}(\hat{\varepsilon}_c)}{\mathrm{std}(\hat{\varepsilon}_{\text{cfg}})}$$ Full renormalisation turns out to over-correct and flatten the image, so the paper interpolates back with a factor $\phi$ (the **guidance_rescale** argument, typically 0.7): $$\hat{\varepsilon}_{\text{final}} = \phi\,\hat{\varepsilon}_{\text{rescaled}} + (1-\phi)\,\hat{\varepsilon}_{\text{cfg}}$$ The standard deviations are computed **per sample**, over all non-batch dimensions, as the population std (**ddof=0**, NumPy's default). ## Your Task Implement: ```python def cfg_rescale(eps_uncond, eps_cond, w, phi=0.7): ``` Inputs are 2-D arrays of shape **(N, D)**; return an array of the same shape. ## Input Format - **eps_uncond**, **eps_cond**: arrays of shape **(N, D)**. - **w** (float): guidance scale. - **phi** (float): rescale blend factor in **[0, 1]**. ## Output Format An array of shape **(N, D)**. ## Sample ```python u = np.array([[0.0, 1.0, -1.0, 0.0]]) c = np.array([[0.0, 2.0, -2.0, 0.0]]) print(np.round(cfg_rescale(u, c, 3.0, 0.7), 4).tolist()) ``` The guided prediction is **[0, 4, -4, 0]** with std 2.8284, while the conditional has std 1.4142, so the rescaled version is halved to **[0, 2, -2, 0]**, and the 0.7 blend gives **[0, 2.6, -2.6, 0]**.

Example:

Input:

u = np.array([[0.0, 1.0, -1.0, 0.0]])
c = np.array([[0.0, 2.0, -2.0, 0.0]])
print(np.round(cfg_rescale(u, c, 3.0, 0.7), 4).tolist())

Output:

[[0.0, 2.6, -2.6, 0.0]]

Reasoning:

eps_cfg = u + 3*(c - u) = [0, 4, -4, 0], whose population std is 2.8284, while std(c) = 1.4142. The ratio 0.5 halves it to [0, 2, -2, 0]. Blending with phi = 0.7 gives 0.7*[0,2,-2,0] + 0.3*[0,4,-4,0] = [0, 2.6, -2.6, 0] -- most of the magnitude inflation removed, but not all.

Constraints:

Standard deviations are per row: reduce over axis 1 and keep the dimension so it broadcasts.
Use the population standard deviation (ddof=0), NumPy's default.
phi = 0.0 must reproduce plain classifier-free guidance exactly.
Both std values are guaranteed non-zero for the given inputs.
Do not round inside the function.

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

1. Background Knowledge

Classifier-Free Guidance (CFG) is a technique used in diffusion models to steer generation toward a specific condition (like a text prompt) without requiring a separate classifier. It works by interpolating between the unconditional noise prediction ( $\hat{\varepsilon}_u$ ) and the conditional noise prediction ( $\hat{\varepsilon}_c$ ). The standard formula is $\hat{\varepsilon}_{\text{cfg}} = \hat{\varepsilon}_u + w(\hat{\varepsilon}_c - \hat{\varepsilon}_u)$ , where $w$ is the guidance scale. While increasing $w$ improves adherence to the condition, it often causes over-exposure or saturation artifacts because the magnitude of the noise prediction grows disproportionately.

The issue arises because the sampler assumes the noise prediction has a specific variance (typically close to 1). When $w$ is large, the standard deviation of $\hat{\varepsilon}_{\text{cfg}}$ becomes much larger than that of $\hat{\varepsilon}_c$ . This inflated magnitude pushes pixel values toward the extremes of the data range. Guidance Rescaling corrects this by normalizing the guided prediction's standard deviation to match that of the conditional prediction, effectively keeping the "energy" of the noise within expected bounds.

However, fully normalizing can sometimes make images look too flat or dull. Therefore, a blending factor $\phi$ is introduced. The final output is a weighted average between the fully rescaled prediction and the original CFG prediction. This allows users to balance between high adherence (high $w$ ) and natural-looking variance (controlled by $\phi$ ). The standard deviations are calculated per sample across all feature dimensions, treating each row in the input array as an independent data point.

2. Algorithm Approach

The core algorithm involves three main mathematical operations applied element-wise or along specific axes:

Compute CFG Prediction: Calculate the standard guided noise vector using the linear interpolation formula.
Calculate Standard Deviations: Compute the population standard deviation ( $ddof=0$ ) for both the conditional prediction ( $\hat{\varepsilon}_c$ ) and the CFG prediction ( $\hat{\varepsilon}_{\text{cfg}}$ ) for each sample.
Rescale and Blend:

Scale $\hat{\varepsilon}_{\text{cfg}}$ by the ratio $\frac{\mathrm{std}(\hat{\varepsilon}_c)}{\mathrm{std}(\hat{\varepsilon}_{\text{cfg}})}$ .
Interpolate between this rescaled vector and the original $\hat{\varepsilon}_{\text{cfg}}$ using $\phi$ .

This approach relies heavily on vectorized operations to handle batches of samples efficiently. You should avoid explicit Python loops over the batch dimension $N$ and instead use library functions (like NumPy) that operate on entire arrays or specific axes.

3. Step-by-Step Strategy

Compute the CFG Vector: Calculate $\hat{\varepsilon}_{\text{cfg}} = \hat{\varepsilon}_u + w \cdot (\hat{\varepsilon}_c - \hat{\varepsilon}_u)$ . Ensure you perform this operation on the entire array to maintain shape $(N, D)$ .
Compute Standard Deviations:

Calculate $\mathrm{std}_c = \mathrm{std}(\hat{\varepsilon}_c, \text{axis}=1, \text{ddof}=0)$ . This results in a 1D array of shape $(N,)$ .
Calculate $\mathrm{std}_{\text{cfg}} = \mathrm{std}(\hat{\varepsilon}_{\text{cfg}}, \text{axis}=1, \text{ddof}=0)$ . This also results in a 1D array of shape $(N,)$ .
Note: Use ddof=0 to ensure you are calculating the population standard deviation, as specified.

Handle Division Safety: Although rare in practice, $\mathrm{std}_{\text{cfg}}$ could theoretically be zero if the CFG prediction is constant. In numerical implementations, it is good practice to add a small epsilon ( $\epsilon$ ) to the denominator to prevent division by zero, though for this specific problem, valid inputs usually guarantee non-zero variance.
Rescale the CFG Vector: Compute the scaling factor: $\text{scale} = \frac{\mathrm{std}_c}{\mathrm{std}_{\text{cfg}}}$ . Compute the rescaled vector: $\hat{\varepsilon}_{\text{rescaled}} = \hat{\varepsilon}_{\text{cfg}} \cdot \text{scale}$ . Crucial: You must broadcast the 1D scale array to match the 2D shape of $\hat{\varepsilon}_{\text{cfg}}$ . In NumPy, this often requires reshaping scale to $(N, 1)$ or relying on automatic broadcasting if the library supports it correctly for row-wise multiplication.
Blend the Results: Compute the final output: $\hat{\varepsilon}_{\text{final}} = \phi \cdot \hat{\varepsilon}_{\text{rescaled}} + (1 - \phi) \cdot \hat{\varepsilon}_{\text{cfg}}$ .
Return the Result: Return the resulting array of shape $(N, D)$ .

4. Common Pitfalls

Axis Confusion: When computing standard deviations, ensure you are reducing along the feature dimension (axis 1), not the batch dimension (axis 0). Reducing along axis 0 would give you a single std for the whole batch, which is incorrect.
Broadcasting Errors: When multiplying the 2D array $\hat{\varepsilon}_{\text{cfg}}$ by the 1D array of standard deviations, ensure the dimensions align. If std_cfg has shape $(N,)$ , multiplying directly might not work as expected in all libraries or might broadcast incorrectly. Reshaping std_cfg to $(N, 1)$ is a robust way to ensure row-wise scaling.
Degree of Freedom (ddof): By default, some statistical functions use ddof=1 (sample standard deviation). The problem explicitly requires ddof=0 (population standard deviation). Failing to set this will result in incorrect scaling factors.
Order of Operations: Do not rescale before computing the CFG vector. The rescaling applies to the result of the CFG interpolation, not the individual conditional/unconditional predictions.
Float Precision: While not usually a failure point, ensure intermediate calculations are done in float precision to avoid integer division issues if inputs are integers (though inputs are typically floats in diffusion contexts).

5. Time & Space Complexity

Time Complexity: $O(N \cdot D)$ .
Computing the CFG vector involves element-wise addition and multiplication: $O(N \cdot D)$ .
Computing standard deviations involves iterating over all elements to calculate mean and variance: $O(N \cdot D)$ .
Rescaling and blending involve element-wise operations: $O(N \cdot D)$ .
Since these steps are sequential, the total time complexity is linear with respect to the total number of elements in the input arrays.
Space Complexity: $O(N \cdot D)$ .
We need to store the intermediate $\hat{\varepsilon}_{\text{cfg}}$ array: $O(N \cdot D)$ .
We need to store the $\hat{\varepsilon}_{\text{rescaled}}$ array: $O(N \cdot D)$ .
The standard deviation arrays are $O(N)$ , which is negligible compared to $O(N \cdot D)$ .
The final output array is $O(N \cdot D)$ .
Thus, the auxiliary space required is proportional to the size of the input data.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.