Karras Sigma Schedule

hardFREEForward Process and Noise SchedulesDiffusion ModelsDiffusion Models 1

Problem Statement

The EDM / Karras samplers parameterize diffusion by noise level sigma instead of a discrete beta schedule, and place the sampling sigmas on a warped grid controlled by rho. Build that schedule.

Background

Karras et al. (2022) define N sampling steps whose sigmas interpolate between sigma_max and sigma_min in a rho-warped space (higher rho concentrates steps near sigma_min):

$\sigma_i = \left(\sigma_{\max}^{1/\rho} + \frac{i}{N-1}\big(\sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho}\big)\right)^{\rho}, \quad i = 0, \dots, N-1$

so sigma_0 = sigma_max and sigma_{N-1} = sigma_min. A final sigma = 0 is appended for the last denoising jump to a clean sample.

Your Task

Implement:

def karras_sigmas(N, sigma_min, sigma_max, rho=7.0):

Return a list of N + 1 sigmas (the N warped values followed by 0.0), rounded to 4 decimals.

Input Format

N (int): number of sampling steps, N >= 2.
sigma_min, sigma_max (float): 0 < sigma_min < sigma_max.
rho (float): warp exponent.

Output Format

A list of N + 1 floats rounded to 4 decimals.

Sample

print(karras_sigmas(3, 0.1, 10.0, 7.0))

Output:

[10.0, 1.4507, 0.1, 0.0]

Example

Input

print(karras_sigmas(3, 0.1, 10.0, 7.0))

Output

[10.0, 1.4507, 0.1, 0.0]

Reasoning

Inputs: N $= 3$ sampling steps, noise levels from sigma_max $= 10$ down to sigma_min $= 0.1$ , and warp exponent rho $= 7$ . The output therefore holds $N + 1 = 4$ values.
Warp the endpoints: take the $1/\rho$ power of each, $10^{1/7} \approx 1.38950$ and $0.1^{1/7} \approx 0.71969$ . In this warped space the gap between them is only about $0.67$ .
Evenly spaced fractions: with $N = 3$ , the fraction $\frac{i}{N-1}$ takes the values $0, \tfrac{1}{2}, 1$ , so the three warped values are the two endpoints and their midpoint $\frac{1.38950 + 0.71969}{2} \approx 1.05459$ .
Unwarp: raising back to the 7th power gives $10$ , $1.05459^7 \approx 1.4507$ and $0.1$ . The endpoints return exactly to $\sigma_{\max}$ and $\sigma_{\min}$ .
Finish: append the clean-sample level $0$ and round to 4 decimals: [10.0, 1.4507, 0.1, 0.0].
Notice: the middle sigma sits far below the arithmetic midpoint $5.05$ (which is what $\rho = 1$ would give). A large $\rho$ crowds the steps toward the small, detail-refining noise levels.

Constraints

N >= 2, 0 < sigma_min < sigma_max, rho > 0.
Interpolate in sigma^{1/rho} space, then raise to rho.
Append a trailing 0.0; round to 4 decimals.

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

In diffusion models, the forward process gradually adds Gaussian noise to a data sample over a sequence of timesteps. Traditional implementations (e.g., DDPM) use a discrete schedule of variance parameters $\beta_t$ , but the EDM (Elucidated Diffusion Models) framework by Karras et al. (2022) reparameterizes the process using a continuous noise level $\sigma$ . This simplifies the math and makes the sampling procedure more stable and interpretable.

The Karras schedule does not place sampling steps uniformly in $\sigma$ -space. Instead, it applies a power warp controlled by the exponent $\rho$ . By working in the transformed space $\sigma^{1/\rho}$ , the schedule can concentrate more steps near the low-noise end (close to the clean data) when $\rho > 1$ . This is critical because the final denoising steps require finer resolution to produce high-quality samples. The formula interpolates linearly between $\sigma_{\max}^{1/\rho}$ and $\sigma_{\min}^{1/\rho}$ , then maps back to $\sigma$ -space by raising the result to the power $\rho$ .

A key practical detail is that the sampling loop requires a final step where $\sigma = 0$ , representing the transition from the noisiest remaining state to a fully clean sample. This zero is appended after the $N$ computed warped sigmas.

2. Algorithm Approach

This is a direct formula evaluation problem. There is no iterative search, dynamic programming, or optimization involved. The approach is:

Compute the warped endpoints: $\sigma_{\max}^{1/\rho}$ and $\sigma_{\min}^{1/\rho}$ .
For each index $i$ from $0$ to $N-1$ , compute the linear interpolation coefficient $t_i = \frac{i}{N-1}$ .
Evaluate the warped value: $w_i = \sigma_{\max}^{1/\rho} + t_i (\sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho})$ .
Map back to $\sigma$ -space: $\sigma_i = w_i^{\rho}$ .
Append $0.0$ to the list.
Round all values to 4 decimal places.

The core pattern is vectorized or loop-based evaluation of a closed-form expression.

3. Step-by-Step Strategy

Validate inputs: Ensure $N \geq 2$ , $0 < \sigma_{\min} < \sigma_{\max}$ , and $\rho > 0$ .
Precompute constants: Calculate sigma_max_pow = sigma_max ** (1.0 / rho) and sigma_min_pow = sigma_min ** (1.0 / rho).
Iterate over indices: Loop $i$ from $0$ to $N-1$ .
Compute the interpolation fraction: t = i / (N - 1).
Compute the warped intermediate value: w = sigma_max_pow + t * (sigma_min_pow - sigma_max_pow).
Compute the final sigma: sigma = w ** rho.
Append sigma to a result list.
Append zero: Add 0.0 to the end of the list.
Round: Apply round(x, 4) to every element in the list.
Return: Return the rounded list.

Example skeleton:

def karras_sigmas(N, sigma_min, sigma_max, rho=7.0):
 sigmas = []
 smax_p = sigma_max ** (1.0 / rho)
 smin_p = sigma_min ** (1.0 / rho)
 for i in range(N):
 t = i / (N - 1)
 w = smax_p + t * (smin_p - smax_p)
 sigmas.append(w ** rho)
 sigmas.append(0.0)
 return [round(s, 4) for s in sigmas]

4. Common Pitfalls

Off-by-one in interpolation: The denominator must be $N-1$ , not $N$ , so that $i=0$ gives exactly $\sigma_{\max}$ and $i=N-1$ gives exactly $\sigma_{\min}$ . Using $N$ would shift the entire schedule.
Forgetting the final zero: The output must have length $N+1$ . Omitting the trailing 0.0 is a frequent error.
Rounding precision: Round to 4 decimal places as specified. Rounding to 2 or 6 will fail test cases.
Integer division: In Python 3, / performs float division, but be cautious if porting to other languages. Ensure i / (N - 1) produces a float.
Negative base for fractional power: Since $\sigma_{\min} > 0$ and $\sigma_{\max} > 0$ , the warped values $w_i$ remain positive, so $w_i^{\rho}$ is well-defined. However, if inputs were invalid (e.g., $\sigma_{\min} \leq 0$ ), fractional powers could produce complex numbers or errors.

5. Time & Space Complexity

Time Complexity: $O(N)$ , since we perform a constant amount of arithmetic for each of the $N$ steps.
Space Complexity: $O(N)$ , to store the list of $N+1$ sigmas.

solution.py

Python 3.13

Test results