Stanford CS229 Machine Learning Spring 2026 Lecture 12: Representation Learning

Diffusion Model Training and Loss Functions

Diffusion models are trained by learning to reverse a forward noising process. The goal is to parameterize a reverse distribution $p_\theta(x_{t-1} | x_t)$ that can reconstruct a clean image from noise.

The Loss Function

The formal objective for training a diffusion model is to maximize the log-likelihood of the data, which leads to an Evidence Lower Bound (ELBO). This simplifies into a sequence of terms $L_{t-1}$ for $t=1$ to $T$.

In practice, the complex coefficients derived from the ELBO are often dropped. The training objective becomes the minimization of the sum of squared differences between the true mean of the reverse process (given the original data $x_0$) and the predicted mean $\mu_\theta(x_t, t)$:

$$\text{Loss} = \mathbb{E} [\mu_{\tilde{t}}(x_t, x_0) - \mu_\theta(x_t, t)]^2$$

Noise Prediction Reparameterization

Rather than predicting the mean $\mu_\theta$ directly, practitioners typically reparameterize the network to predict the noise $\epsilon$ that was added to the original image $x_0$ to create the noisy image $x_t$.

By hardcoding the known linear relationship between $x_t$, $x_0$, and the noise $\epsilon$, the loss function simplifies to a noise prediction problem:

$$\text{Loss} = \mathbb{E} [\epsilon - \epsilon_\theta(x_t, t)]^2$$

This turns the generation task into a series of smaller, locally easier reconstruction steps. The model starts with pure white noise and gradually sharpens the image, adding detail at each step.

The Foundation Model Paradigm

Foundation models represent a shift in machine learning from task-specific training to a two-phase paradigm: pre-training and adaptation.

Pre-training

Pre-training involves training a model on massive, diverse, and often unlabeled datasets. The focus is on scale and diversity rather than strict data quality. This phase creates a "foundation"—a general-purpose model that captures broad semantic information from the data.

Adaptation

Adaptation is the process of tailoring a pre-trained foundation model to a specific downstream task. This can occur in several settings:

  • Zero-shot learning: The model solves a task based on a description without any task-specific training data.
  • Few-shot learning: The model is adapted using a very small labeled dataset (e.g., 5-10 examples).
  • Fine-tuning: The model is updated using a larger labeled dataset for a specific task.

Representation Learning and Adaptation Techniques

Representation learning focuses on training a model $f_\theta$ that maps raw data $x$ into a lower-dimensional vector (embedding or feature). This representation is then used to solve downstream tasks.

Linear Probing

Linear probing is a technique where the representation model $f_\theta$ is kept frozen, and a simple linear head (a weight vector $w$) is trained on top of the fixed embeddings to predict the target $y$. This is highly effective because the pre-trained representation often transforms non-linear relationships in the raw data into linear relationships in the embedding space.

Fine-tuning and LPFT

In full fine-tuning, both the linear head $w$ and the representation parameters $\theta$ are optimized. While the loss function may be the same as in linear probing, initializing $\theta$ with a pre-trained model typically leads to a better global minimum and superior test performance compared to training from scratch.

LPFT (Linear Probing then Fine-Tuning) is a hybrid approach where the model is first trained via linear probing (optimizing only $w$) to establish a stable connection between the head and the representation. Subsequently, both $w$ and $\theta$ are jointly optimized. This prevents random initialization of the linear head from "destroying" the pre-trained representations during the early stages of fine-tuning.

Low-Rank Adaptation (LoRA)

LoRA is an efficient adaptation method that avoids updating all billions of parameters in a large model. Instead of updating a weight matrix $W$ to $W + \Delta W$, LoRA restricts the update $\Delta W$ to be a low-rank decomposition: $\Delta W = AB$, where $A$ and $B$ are matrices with a much smaller inner dimension $r$.

Benefits of LoRA

  • Memory Efficiency: While the forward pass still requires the original frozen weights $W_0$, LoRA significantly reduces the memory needed for gradients and optimizer states (like momentum in Adam), as these are only stored for the small matrices $A$ and $B$.
  • Multi-tenant Serving: LoRA is ideal for serving many users. A single shared $W_0$ can be kept in memory, while small, user-specific adapters ($A_i, B_i$) are swapped in and out rapidly. This allows a provider to serve thousands of customized models without duplicating the massive base model for every user.

Sources