Designing Custom Noise Schedules

Standard noise schedules, such as linear and cosine, provide a common starting point and are widely used. However, they may not be optimal for all datasets or generative tasks. Recognizing the limitations of these fixed schedules, it is important to design custom noise schedules tailored to specific requirements. The primary objective is to control the rate at which noise is added during the forward process, which directly influences the reverse denoising process learned by the model.

Recall that the forward process variance schedule, typically denoted by $\beta_t$ for timesteps $t=1, ..., T$, dictates the entire diffusion process. From $\beta_t$, we derive $\alpha_t = 1 - \beta_t$ and the cumulative product $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$. The choice of $\beta_t$ values directly impacts how quickly information from the initial data $x_0$ is obscured. A schedule that adds noise too quickly might destroy fine details early on, making it harder for the model to recover them. Conversely, a schedule that adds noise too slowly might require a very large number of timesteps $T$ or lead to inefficient learning in later stages where the signal is already weak.

Designing a custom schedule often involves defining a function or a sequence for $\beta_t$ that deviates from the standard linear or cosine forms.

Motivation for Custom Schedules

Why move past established schedules?

1. Dataset Characteristics: Datasets with specific structures (e.g., images with predominantly fine textures vs. large smooth areas) might benefit from schedules that add noise differently across timesteps.
2. Task Specificity: Certain tasks, like inpainting or image editing, might require preserving information more carefully in early diffusion steps compared to unconditional generation.
3. Sample Quality Focus: Custom schedules can be designed to improve specific aspects of sample quality, such as reducing blurriness or enhancing detail fidelity, potentially by concentrating the model's learning capacity on specific noise level ranges.
4. Efficiency: Tailoring the schedule might allow for achieving comparable or better results with fewer diffusion timesteps ($T$).

Approaches to Designing Noise Schedules

Instead of relying on predefined formulas like linear or cosine, we can define $\beta_t$ using other functional forms or rules:

1. Polynomial Schedules: Generalize the linear schedule $\beta_t \propto t$ or the cosine schedule. We can define $\beta_t$ using higher-order polynomials, for instance:

   $$\beta_t = c_1 \left(\frac{t}{T}\right)^p + c_0$$

   where $p$ is the polynomial degree, and $c_1, c_0$ are constants chosen to ensure $\beta_t$ stays within a desired range (e.g., $10^{-4}$ to $0.02$) and maintains monotonicity. Varying $p$ allows for different curvatures in the noise addition rate.

2. Piecewise Schedules: Define different functions for $\beta_t$ over different intervals of $t$. For example, a schedule could be linear for the first $T/2$ steps and then transition to a cosine-like decay for the remaining steps. This allows fine-grained control over noise addition at different phases of the diffusion process.

3. Signal-to-Noise Ratio (SNR) Based Design: A more principled approach involves designing the schedule based on the desired Signal-to-Noise Ratio (SNR) at each timestep $t$. The SNR is defined as:

   $$\text{SNR}(t) = \frac{\mathbb{E}[ ( \sqrt{\bar{\alpha}_t} x_0 )^2 ]}{\mathbb{E}[ ( \sqrt{1 - \bar{\alpha}_t} \epsilon )^2 ]} = \frac{\bar{\alpha}_t \mathbb{E}[x_0^2]}{(1 - \bar{\alpha}_t) \mathbb{E}[\epsilon^2]}$$

   Assuming $\mathbb{E}[x_0^2] \approx 1$ and $\mathbb{E}[\epsilon^2] = 1$ after normalization, $\text{SNR}(t) \approx \frac{\bar{\alpha}_t}{1 - \bar{\alpha}_t}$. We can work backward: define a target SNR($t$) function (e.g., exponentially decaying) and derive the required $\bar{\alpha}_t$, and subsequently $\beta_t$. This connects the schedule design directly to the information content remaining at each step. For example, ensuring the SNR drops smoothly might lead to more stable training.

4. Logarithmic Schedules: Schedules can also be defined in log-space, often focusing on the log-SNR. This can provide better control over the dynamics, especially when dealing with very small or very large $\bar{\alpha}_t$ values.

Visualizing Schedule Differences

Let's compare the cumulative noise level, represented by $\sqrt{1-\bar{\alpha}_t}$ (which indicates how much noise dominates the signal), for different schedule types over $T=1000$ timesteps. A higher value means more noise.

A comparison showing how quickly noise accumulates under linear, cosine, and a custom (quadratic-like) schedule. The custom schedule adds noise slowly initially, then accelerates, compared to the cosine schedule which adds noise faster early on.

Implementation

When implementing a custom schedule, you typically need to pre-compute the $\beta_t$, $\alpha_t$, and $\bar{\alpha}_t$ values for $t=1, ..., T$. These are then stored and used during both training (for sampling $x_t$ given $x_0$) and inference (for the denoising steps).

Steps involve:

1. Define $\beta_t$: Choose a function or sequence for $\beta_t$ over the range $t \in [1, T]$. Ensure $\beta_t$ values are sensible (e.g., small, positive, and generally non-decreasing, though strict monotonicity isn't always required). The start and end values ($\beta_1$, $\beta_T$) significantly impact behavior near pure data and pure noise.
2. Compute Derived Values: Calculate $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$. Also pre-compute other values often needed, like $\sqrt{\bar{\alpha}_t}$, $\sqrt{1 - \bar{\alpha}_t}$, and terms used in the denoising step (which depend on the specific parameterization, e.g., involving $\tilde{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t$).
3. Integration: Use these pre-computed values within your diffusion model's training loop and sampling procedures.

Evaluation and Trade-offs

Evaluating a custom noise schedule requires empirical testing. Train your diffusion model using the new schedule and compare:

• Training Dynamics: Does the loss converge smoothly? Are there instabilities?
• Sample Quality: Assess generated samples using metrics like FID (Fréchet Inception Distance), IS (Inception Score), or domain-specific metrics. Qualitative evaluation is also important.
• Sampling Speed: Does the custom schedule allow for good results with fewer inference steps ($N < T$) when using samplers like DDIM?

Designing custom schedules is often an iterative process of proposing a schedule, training, evaluating, and refining. There's a trade-off between the complexity of the schedule design and the potential gains. While standard schedules work well generally, a carefully designed custom schedule can provide significant advantages for specific applications or push the boundaries of sample quality. This understanding paves the way for exploring schedules that are not just designed, but learned, which we will cover in the next section.