1 / 1

Time-Conditioned Residual Block

Diffusion ModelsConvolutional NetworksNeural Networks

Medium

Apply one residual block that conditions feature maps on a timestep embedding and return its output.

\begin{aligned} \tau &= \mathrm{SiLU}(e)\, W_t^{\top} + b_t \\ h &= \mathrm{conv}_1\big(\mathrm{SiLU}(\mathrm{GN}_1(x))\big) + \tau \\ \text{out} &= \mathrm{skip}(x) + \mathrm{conv}_2\big(\mathrm{SiLU}(\mathrm{GN}_2(h))\big) \end{aligned}

Here x has shape (B, C_in, H, W), $e$ is time_embedding of shape (B, E), and SiLU(z) = z times sigmoid(z). Each GN is a group norm with num_groups groups, eps 1e-5 and a per-channel affine weight and bias. Both 3x3 convolutions use stride 1 and padding 1. The first half uses the norm1 and conv1 entries of params and the second half uses the norm2 and conv2 entries. Here $W_t$ is time_weight and $b_t$ is time_bias, laid out like a linear layer, so $\tau$ has shape (B, C_out) and its value for sample b and channel c is added to every pixel of that channel. If params has a skip_weight key, skip(x) is a 1x1 convolution with skip_weight and skip_bias, and otherwise skip(x) is x itself.

key	shape
norm1_weight, norm1_bias	(C_in,)
conv1_weight	(C_out, C_in, 3, 3)
conv1_bias, time_bias	(C_out,)
time_weight	(C_out, E)
norm2_weight, norm2_bias, conv2_bias	(C_out,)
conv2_weight	(C_out, C_out, 3, 3)
skip_weight (optional)	(C_out, C_in, 1, 1)
skip_bias (optional)	(C_out,)

Return a PyTorch tensor of shape (B, C_out, H, W) with the dtype of x.

Examples

Input: x = [[[[-0.899,-0.386],[0.202,-0.533]],[[-0.582,-0.469],[0.078,0.377]]]], time_embedding = [[-0.342,-0.128]], params = {10 tensors: norm1_weight, norm1_bias (2,); conv1_weight (2, 2, 3, 3); conv1_bias (2,); time_weight (2, 2); time_bias (2,); norm2_weight, norm2_bias (2,); conv2_weight (2, 2, 3, 3); conv2_bias (2,); values in the test case}, num_groups = 2

Output: [[[[-1.064443,-0.640805],[-0.048407,-0.468474]],[[-0.859204,-1.346627],[0.949704,0.050282]]]]

Explanation: With 2 channels and 2 groups each group norm normalizes one channel. The projected time term adds about 0.1184 to channel 0 and 0.0066 to channel 1 after the first convolution, and since no skip_weight is given, x itself is added to the result.

Input: x = [[[[-0.665,-1.824],[0.034,-0.157]],[[1.199,0.149],[-1.331,-1.122]]],[[[0.899,-0.38],[0.817,-0.721]],[[-1.343,-1.786],[0.961,0.399]]]], time_embedding = [[1.248,0.725,0.153],[-1.007,0.634,-0.231]], params = {12 tensors: norm1_weight, norm1_bias (2,); conv1_weight (4, 2, 3, 3); conv1_bias (4,); time_weight (4, 3); time_bias (4,); norm2_weight, norm2_bias (4,); conv2_weight (4, 4, 3, 3); conv2_bias (4,); skip_weight (4, 2, 1, 1); skip_bias (4,); values in the test case}, num_groups = 2

Output: [[[[-1.693148,-1.807478],[1.0174,0.343204]],[[0.195915,-0.58309],[-0.031223,-0.128251]],[[1.347961,0.456708],[-0.482963,0.213237]],[[-0.812215,-0.481138],[1.901938,0.6125]]],[[[1.410261,0.912458],[0.379139,0.083182]],[[-1.122136,0.309703],[0.311363,0.5892]],[[0.634216,-1.026428],[1.166675,-0.761879]],[[1.674857,1.385977],[-1.339211,-1.401852]]]]

Input: x = [[[[-1.226,-0.302,1.299]],[[-2.034,2.737,0.122]]]], time_embedding = [[1.704,1.76]], params = {10 tensors: norm1_weight, norm1_bias (2,); conv1_weight (2, 2, 3, 3); conv1_bias (2,); time_weight (2, 2); time_bias (2,); norm2_weight, norm2_bias (2,); conv2_weight (2, 2, 3, 3); conv2_bias (2,); values in the test case}, num_groups = 1

Output: [[[[-1.449865,0.011765,1.43844]],[[-2.05555,1.846919,-0.221911]]]]

Hint 1

torch.nn.functional provides group_norm, silu, conv2d and linear, which cover every layer in the block.

Hint 2

Indexing the (B, C_out) time projection with [:, :, None, None] lets it broadcast over the spatial dimensions.

Hint 3

F.group_norm(x, num_groups, weight, bias, eps=1e-5) and F.conv2d(h, weight, bias, padding=1) are the two layer calls.

Requirements

Use PyTorch, and torch.nn.functional is allowed.
Apply SiLU to time_embedding before the time projection, and add the projection after the first convolution.
Leave x, time_embedding and every tensor in params unchanged.

Constraints

x has shape (B, C_in, H, W) with B between 1 and 3, C_in and C_out between 2 and 8, and H and W between 1 and 6.
time_embedding has shape (B, E) with E between 2 and 8.
params holds exactly the keys in the table, with skip_weight and skip_bias either both present or both absent. Both are present whenever C_in differs from C_out.
num_groups divides both C_in and C_out.
All tensors share one dtype, float32 or float64, and all values are finite.

Try Similar Problems

Spatial Self-Attention BlockMedium Sinusoidal Timestep EmbeddingEasy Tiny U-Net Forward PassHard Gaussian Mixture ScoreMedium Cosine Noise ScheduleMedium Rescale to Zero Terminal SNRMedium Posterior Mean and VarianceMedium Prediction ParameterizationMedium Autoencoder Forward PassMedium Gradient AccumulationMedium DQN Loss with Target NetworkMedium Weight InitializationMedium