Study Plan

Toy Diffusion Training Step

Diffusion ModelsTrainingOptimization

Hard

Run one training step of a tiny noise-prediction diffusion model and return the loss, gradient norm and updated weights.

\begin{aligned} x_t &= \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \epsilon, \qquad \hat\epsilon = \operatorname{SiLU}\big([\,x_t,\ e(t)\,] W_1 + b_1\big) W_2 + b_2 \\ L &= \operatorname{mean}\big((\hat\epsilon - \epsilon)^2\big), \qquad \theta \leftarrow \theta - \eta\, \min\Big(1, \frac{c}{\lVert g \rVert + 10^{-6}}\Big)\, g \end{aligned}

Each sample uses $\bar\alpha_t$ = alpha_bars at its own timestep, $x_0$ is x0 and $\epsilon$ is noise. The embedding $e(t) = [\cos(t f), \sin(t f)]$ has width embed_dim, where half = embed_dim // 2 and the vector $f$ holds $f_i = \exp(-\ln(10000)\, i / \text{half})$ for $i = 0, \dots, \text{half} - 1$ . One zero is appended when embed_dim is odd, and the bracket concatenates along the feature axis.

The weights $W_1$ , $b_1$ , $W_2$ and $b_2$ are the entries W1, b1, W2 and b2 of params, and $\operatorname{SiLU}(z) = z \cdot \operatorname{sigmoid}(z)$ . The loss $L$ averages over all B times D entries. The vector $g$ collects the gradients of $L$ for all four tensors, $\lVert g \rVert$ is their global L2 norm, $\eta$ is learning_rate, $c$ is max_grad_norm, and each weight tensor $\theta$ is updated with its own part of $g$ .

Return a dict with "loss" ( $L$ before the update) and "grad_norm" ( $\lVert g \rVert$ before clipping) as Python floats, and "params", a dict of new tensors with the same keys, shapes and dtype.

Examples

Input: x0 = [[1.0,-1.0],[0.5,0.5]], timesteps = [0,3], noise = [[0.2,-0.4],[1.0,0.0]], alpha_bars = [0.9,0.7,0.4,0.1], params = {"W1":[[0.5,-0.3,0.2],[0.1,0.4,-0.2],[0.3,0.0,0.1],[-0.2,0.2,0.5]],"b1":[0.0,0.1,-0.1],"W2":[[0.4,-0.1],[0.2,0.3],[-0.5,0.2]],"b2":[0.0,0.05]}, embed_dim = 2, learning_rate = 0.1, max_grad_norm = 100.0

Output: {"loss":0.283354,"grad_norm":0.750735,"params":{"W1":[[0.51813,-0.296348,0.179547],[0.096968,0.401509,-0.195672],[0.292565,-0.004809,0.106971],[-0.198283,0.20057,0.498159]],"b1":[0.016776,0.10323,-0.118989],"W2":[[0.410936,-0.109202],[0.194805,0.304335],[-0.496157,0.195084]],"b2":[0.057866,0.029459]}}

Explanation: The two samples are noised at timesteps 0 and 3, the MLP predicts their noise with a loss of about 0.283354, and the gradient norm 0.750735 is below max_grad_norm, so every weight moves by exactly 0.1 times its gradient.

Input: x0 = [[2.0,0.0],[-1.0,1.5],[0.0,-2.0]], timesteps = [1,2,4], noise = [[1.0,1.0],[-0.5,0.3],[0.8,-1.2]], alpha_bars = [0.95,0.8,0.6,0.35,0.15], params = {"W1":[[0.6,-0.4],[0.2,0.5],[0.3,-0.1],[-0.2,0.4],[0.1,0.1],[0.5,-0.3]],"b1":[0.1,-0.1],"W2":[[0.7,-0.2],[0.3,0.6]],"b2":[0.0,0.0]}, embed_dim = 4, learning_rate = 0.5, max_grad_norm = 0.1

Output: {"loss":0.784072,"grad_norm":1.007376,"params":{"W1":[[0.598729,-0.380109],[0.186212,0.490759],[0.295572,-0.093402],[-0.199701,0.395091],[0.091689,0.096041],[0.500144,-0.300154]],"b1":[0.100302,-0.104911],"W2":[[0.703093,-0.166994],[0.279822,0.59106]],"b2":[0.005674,-0.00356]}}

Input: x0 = [[0.3,-0.7]], timesteps = [0], noise = [[-1.0,0.5]], alpha_bars = [0.5,0.2], params = {"W1":[[0.2,0.1],[-0.3,0.4],[0.5,0.0],[0.1,-0.2],[0.7,0.7]],"b1":[0.0,0.0],"W2":[[0.3,-0.3],[0.1,0.2]],"b2":[0.1,-0.1]}, embed_dim = 3, learning_rate = 0.05, max_grad_norm = 1.0

Output: {"loss":0.93028,"grad_norm":1.537563,"params":{"W1":[[0.206443,0.099852],[-0.298159,0.399958],[0.486983,0.0003],[0.1,-0.2],[0.7,0.7]],"b1":[-0.013017,0.0003],"W2":[[0.289671,-0.29393],[0.101921,0.198871]],"b2":[0.061758,-0.077527]}}

Hint 1

Cloning each weight with detach().clone().requires_grad_(True) lets autograd compute gradients without touching the inputs.

Hint 2

The global norm is the square root of the sum of the squared entries of every gradient tensor together.

Hint 3

Apply the update inside torch.no_grad() so the new weights are plain tensors.

Hint 4

torch.nn.utils.clip_grad_norm_ uses this same global-norm formula and returns the norm before clipping.

Requirements

Use PyTorch; torch.autograd is allowed on copies of the weights.
Compute the embedding in the dtype of x0.
Leave every input unchanged, including the tensors in params.

Constraints

x0 and noise are float tensors of shape (B, D) with B between 1 and 6 and D between 1 and 3.
timesteps is an int64 tensor of shape (B,) with values between 0 and T - 1, and alpha_bars is a float tensor of shape (T,) with T between 2 and 50 and values in (0, 1).
params has exactly the keys W1 (D + embed_dim, hidden), b1 (hidden), W2 (hidden, D) and b2 (D), with hidden between 1 and 8.
embed_dim is an integer between 2 and 8, learning_rate is between 0.01 and 1, and max_grad_norm is positive.
All float tensors share one dtype, float32 or float64, and all values are finite.

Try Similar Problems

DiT Patchify and adaLN-Zero BlockHard Toy DDPM SamplerHard Min-SNR Loss WeightingMedium Frechet Distance Between Feature SetsMedium Tiny U-Net Forward PassHard DDIM StepHard Karras Euler SamplerHard Gaussian Mixture ScoreMedium GRPO Group-Relative AdvantageHard Hungarian MatchingHard Gradient ClippingHard L-BFGSHard

Toy Diffusion Training Step

Diffusion ModelsTrainingOptimization

Hard

Run one training step of a tiny noise-prediction diffusion model and return the loss, gradient norm and updated weights.

\begin{aligned} x_t &= \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \epsilon, \qquad \hat\epsilon = \operatorname{SiLU}\big([\,x_t,\ e(t)\,] W_1 + b_1\big) W_2 + b_2 \\ L &= \operatorname{mean}\big((\hat\epsilon - \epsilon)^2\big), \qquad \theta \leftarrow \theta - \eta\, \min\Big(1, \frac{c}{\lVert g \rVert + 10^{-6}}\Big)\, g \end{aligned}

Return a dict with "loss" ( $L$ before the update) and "grad_norm" ( $\lVert g \rVert$ before clipping) as Python floats, and "params", a dict of new tensors with the same keys, shapes and dtype.

Examples

Hint 1

Cloning each weight with detach().clone().requires_grad_(True) lets autograd compute gradients without touching the inputs.

Hint 2

The global norm is the square root of the sum of the squared entries of every gradient tensor together.

Hint 3

Apply the update inside torch.no_grad() so the new weights are plain tensors.

Hint 4

torch.nn.utils.clip_grad_norm_ uses this same global-norm formula and returns the norm before clipping.

Requirements

Use PyTorch; torch.autograd is allowed on copies of the weights.
Compute the embedding in the dtype of x0.
Leave every input unchanged, including the tensors in params.

Constraints

x0 and noise are float tensors of shape (B, D) with B between 1 and 6 and D between 1 and 3.
timesteps is an int64 tensor of shape (B,) with values between 0 and T - 1, and alpha_bars is a float tensor of shape (T,) with T between 2 and 50 and values in (0, 1).
params has exactly the keys W1 (D + embed_dim, hidden), b1 (hidden), W2 (hidden, D) and b2 (D), with hidden between 1 and 8.
embed_dim is an integer between 2 and 8, learning_rate is between 0.01 and 1, and max_grad_norm is positive.
All float tensors share one dtype, float32 or float64, and all values are finite.

Try Similar Problems

x0 =

Accepts: array

timesteps =

Accepts: array

noise =

Accepts: array

alpha_bars =

Accepts: array

params =

Accepts: object

embed_dim =

Accepts: integer

learning_rate =

Accepts: number

max_grad_norm =

Accepts: number