Advanced 12 min 2 of 4

Eight channels, not four

Widening the first convolution is one line. What its new weights start at is a real decision with three defensible answers and one right question.

What this repairs

Thinking a pretrained UNet can simply "take another input", without asking what the new conv_in weights are initialised to.

Assumes

Cross-attention, concat, ControlNet
What a 4-channel latent is · Latent Diffusion

The tensor the UNet actually receives is a channel-wise concatenation of two latents, and the model names its slots rather than slicing with magic numbers:

$\underbrace{[\,z_t^{\text{target}}\,]}_{\text{channels } 0..3} \;\Vert\; \underbrace{[\,z^{\text{hazy}}\,]}_{\text{channels } 4..7} \;\in\; \mathbb{R}^{8 \times 64 \times 64}.$

One thing about that tensor is worth more than any diagram of it: the conditioning half is never noised. The first four channels are on the schedule from the previous topics — at $t = 999$ they are indistinguishable from noise — and the last four are a clean encode of the hazy input, byte-identical at every one of the twenty-five sampler steps.

Once that is seen, most of the concept is done. The target is what diffuses; the condition is what stays still.

Insight A clean condition is not a leak

The first reaction to "four of the eight channels are clean" is that the model is being handed the answer. It is not. The clean channels are the hazy image — the input to the task. The clear image is the thing being predicted, and it is the half that gets destroyed.

Two consequences fall out of the condition being clean and constant.

It can be encoded once, outside the sampler loop, and reused for every step; which is what the model does. The condition costs one VAE encode per image, not one per step.

And the network is never asked to denoise it, so it can be relied on absolutely. A model with a noisy condition has to spend capacity deciding how much to trust its own input at each timestep; this one does not. That is a real simplification and it is the strongest practical argument for conditioning on a clean latent rather than on a jointly-noised one.

Now the problem. The pretrained SD2.1 conv_in is a convolution from four channels, because SD2.1 generates images from a latent and has nothing to condition on. It has to become a convolution from eight.

Structurally that is trivial — construct a new convolution with the same kernel size, stride, padding and output channels, copy the bias verbatim, and assign the weights. The bias needs no thought because it is per-output-channel and the output channels have not changed. The weights do, and the choice is not cosmetic: whatever the new half starts at, the very first forward pass of training runs the entire pretrained network on activations whose scale that choice determined.

Why one scale is ½ and another is 1/√2

$y = \sum_{c=0}^{3} W_c * x_c, \qquad \operatorname{Var}(y) = V$

The pretrained first layer and the activation variance it produces. Every downstream block — every normalisation, every nonlinearity — was trained against this V, so preserving it is what "not breaking the pretrained network" means here.
$y' = s\sum_{c=0}^{3} W_c * x_c \;+\; s\sum_{c=0}^{3} W_c * u_c$

The widened layer, with the pretrained kernel reused in both halves and scaled by s. The two halves see different inputs — x is the noisy target, u the second group — and the whole question is how those two are related.
$u = x \;\Longrightarrow\; y' = 2s\,y \;\Longrightarrow\; s = \tfrac{1}{2}$

If the two halves carry the same content, the layer computes the same thing twice and the outputs add coherently — amplitudes sum, so halving each restores the original. This is the assumption behind the ×0.5 initialisation.
$u \perp x \;\Longrightarrow\; \operatorname{Var}(y') = 2s^{2}V \;\Longrightarrow\; s = \tfrac{1}{\sqrt{2}}$

If the two halves are independent, the outputs add incoherently — variances sum, not amplitudes — and the preserving scale is 1/√2. This is the assumption stated in the code beside the single-pass constant: variance-preserving for two independent noisy latents.
$\frac{1/\sqrt{2}}{1/2} = \sqrt{2} \approx 1.41$

The two answers differ by 41 %, and neither is a rounding of the other. They are the same question — what scale preserves the pretrained activation statistics — answered under opposite assumptions about the correlation between the two halves.

Pitfall Neither constant is right for the whole schedule

The derivation above has a consequence the codebase does not claim and that is worth stating: the correlation between the two halves is a function of the timestep, so a single scalar cannot be correct everywhere.

In the eight-channel layout the halves are the noisy target and the hazy condition, and the whole thing collapses to one line. Write $r$ for the correlation between the two clean latents; only the signal half of $x_t$ has anything in common with a clean latent, so the correlation the layer actually sees is $\rho_t = \sqrt{\bar{\alpha}_t}\,r$ , and the preserving scale is $s^{\star}_t = 1/\sqrt{2 + 2\rho_t}$ . Both constants are its endpoints: $\rho = 1$ gives $\tfrac{1}{2}$ and $\rho = 0$ gives $1/\sqrt{2}$ .

Since $\bar{\alpha}_t$ falls to zero, $\rho_t$ does too, and $1/\sqrt{2}$ is exactly right at the noisy end of the chain whatever $r$ is. What happens at the clean end depends on $r$ — and $r$ is measurable. Correlating the two real latents exported for the figure gives 0.33 on that scene: the clear image and the hazy image of it encode to latents far closer to independent than to identical, which is not the intuition and is the whole reason to measure. At that value $\tfrac{1}{2}$ is $18\%$ low at $t = 0$ and $29\%$ low at $t = 999$ — never right anywhere on the chain, where $1/\sqrt{2}$ is exact at one end and $16\%$ high at the other.

The twelve-channel layout has a separate arithmetic problem. Its two target slots really are drawn with independent noise by default, so $1/\sqrt{2}$ is the preserving scale for those two — which is exactly what the code's own comment claims for it. But the shipped joint_marigold puts the conditioning slot at the same $1/\sqrt{2}$ , and three independent slots at $1/\sqrt{2}$ give $\sqrt{3/2}$ : $22\%$ hot in the very limit the constant was derived in. The zero-init variant, which leaves the conditioning slot at $0$ , is the one that is exactly preserving there.

None of this is a defect — an activation-scale error of this size at initialisation is recoverable, and training recovers it. But "the two constants are the same argument" is not true, the coherent limit that would make $\tfrac{1}{2}$ correct is not a limit this data reaches, and a reader who understands which assumption each constant encodes has understood the whole issue.

Convention The three initialisations, and what each believes

The eight-channel path offers three modes, and each is defensible-sounding with materially different training behaviour:

Mode	Target half	Condition half	What it assumes
marigold	pretrained × 0.5	pretrained × 0.5	the two halves are coherent, so amplitudes sum
zero_cond	pretrained	0	start blind to the condition and learn to see it
duplicate	pretrained	pretrained	that nothing needs scaling, which is the one clearly wrong answer

marigold is the shipped default. zero_cond is the conservative choice: at initialisation the network is exactly the pretrained unconditional generator, because the condition contributes nothing, so the first training steps cannot be worse than the model you started from — and the condition has to be discovered, which takes longer.

duplicate is the instructive one. It looks like the natural "just copy the weights" answer and it walks straight into the derivation above: activations arrive $\sqrt{2}$ to $2$ times too large, into a stack of normalisations and nonlinearities calibrated for the pretrained scale, before a single gradient step has been taken.

Insight Only the conditioning slot may start at zero

The twelve-channel layout has three slots — noisy RGB, noisy depth, hazy condition — and its zero-init variant zeroes only the third. The comment in the code is explicit that zeroing the depth slot would blind the encoder to the noisy depth latent, so the depth branch could not denoise early in training.

The asymmetry is the lesson, and it generalises. Zeroing an input the network must predict from destroys the task: with no path from $z_t^{\text{depth}}$ to the output there is nothing to denoise, and the gradient that would teach the layer to open that path has to arrive through a layer that currently transmits nothing. Zeroing an input the network may learn to use only delays: the conditional path starts closed, the model works as an unconditional generator, and the condition earns its way in as soon as it reduces the loss.

Same operation, two slots, opposite consequences — decided entirely by whether the slot is on the input side or the answer side of the task.

Check yourself

You widen conv_in with the duplicate mode and training loss is worse than the marigold run from step zero — before any learning could have happened. Why?

Show answer

Because the damage is at initialisation, not in the learning.

duplicate leaves both halves at full pretrained magnitude, so the first convolution's output is between $\sqrt{2}$ and $2$ times its trained scale depending on how correlated the two halves are at that timestep. Every subsequent block sees inputs outside the range its normalisation statistics and nonlinearities were fitted on, and the error compounds with depth rather than staying put.

The useful diagnostic is that this is visible before training: run one forward pass under each init mode and compare the activation statistics at the first few blocks. If the histograms have moved, you have found it, and you have found it in seconds rather than in a training run. It is also why marigold and the $1/\sqrt{2}$ constant exist at all — both are attempts to make that first forward pass land where the pretrained network expects.