Advanced 11 min 3 of 4

One encoder, two decoders

Not two models. One pretrained encoder whose skips are consumed twice — and in the default mode it runs twice, which is where the cost is.

What this repairs

Treating multi-task output as multiple models, when it is one shared encoder whose skips are consumed twice.

Assumes

Eight channels, not four

A UNet is an encoder, a bottleneck, and a decoder that consumes the encoder's skip connections on the way back up. The joint model splits it at the bottleneck: the first convolution, the timestep embedding, the down blocks and the mid block are shared; the up blocks, the output normalisation and the output convolution exist twice.

The dehazing decoder keeps the pretrained up blocks. The depth decoder is a deep copy of them — not a random initialisation. Both branches therefore start as the same pretrained image decoder and diverge only under training, which is the whole reason a depth decoder trained on a modest dataset works at all: it begins already knowing how to turn latent features into spatially coherent output, and only has to learn what to put there.

Insight "Shared" means two different things here

The two joint modes share the same weights and differ in how many times they are run, which is the distinction that decides the cost.

Dual pass — the default in the shipped config — runs the shared encoder twice per sampler step: once on the eight-channel concat of the noisy clear latent with the condition, once on the concat of the noisy depth latent with the same condition. Same weights, different inputs, therefore different skip connections, each consumed by its own decoder. This is not one forward pass producing two outputs.

Single pass runs the encoder once, on a twelve-channel concat of both noisy targets and the condition, and both decoders consume the same skips in lockstep. That genuinely is one forward pass producing two outputs — and it is the variant that is not the default.

The word "shared" in "shared encoder" is therefore about parameters in one mode and about activations in the other. A claim about cost that does not say which is a claim about nothing.

What sharing actually saves

$C_{\text{dual}} = 2E + 2D$

Two encoder passes and two decoder passes per sampler step, where E and D are the cost of one encoder and one decoder pass. The encoder is not amortised across tasks at all in this mode — only its parameters are.
$C_{\text{single}} = E + 2D$

One encoder pass, both decoders. The saving is exactly one encoder pass per step — not half the model, which is the intuition the phrase "shared encoder" invites.
$\frac{C_{\text{single}}}{C_{\text{dual}}} = \frac{E + 2D}{2E + 2D}$

The ratio, left in symbols because the split between encoder and decoder cost is an architectural fact rather than something this page should assert. Measure it before quoting it.
$E = D \;\Longrightarrow\; \frac{3}{4} \quad (25\% \text{ saved})$

For a roughly symmetric UNet the saving is a quarter, not a half. Worth having in mind before treating single-pass as an optimisation: it also forces one skip set to serve two tasks, which is a modelling change and not only a cost change. In parameters the split is measurable and it is not symmetric: summed off the shipped SD2.1 checkpoint, the decoder is 59 % of the UNet and the shared trunk 41 %, so the joint model comes to 1.38 billion parameters against 1.73 billion for two separate single-task models — 20 % fewer, not half.

In practice The structural invariant is asserted, not assumed

Because the depth decoder is a deep copy, the two decoders are structurally identical — and the single-pass lockstep depends on that: it advances both by a shared up-block index and fuses outputs at matching indices, which is only meaningful if block counts, per-block residual counts and output channel widths all agree.

The code asserts all three at construction rather than trusting the copy. That is the right instinct and worth borrowing. A deep copy is structurally identical until someone adds a per-task adapter, a per-task normalisation or a partially-frozen branch, and then the failure is a shape mismatch several layers into a sampler loop, at a point that says nothing about where the assumption was made. An assertion at the place the invariant is created costs a line and turns a confusing failure into an obvious one.

What sharing an encoder bets is that the two tasks want the same features.

For these two tasks that is a strong bet, and it is worth stating explicitly rather than assuming. Dehazing needs to know where surfaces are, how far away they are and how thick the haze is along each ray — because the haze thickness at a pixel is a function of its distance. Depth estimation needs the same three things. The tasks are not merely compatible; the physical model that relates them is the one in the fog topic, and the shared encoder is that relationship expressed as an architecture.

The failure mode is negative transfer, and it is what happens when the bet is wrong: two tasks that want different features contend for one encoder, and both end up worse than two separate models would have been. This is not hypothetical — the same repository ships optional per-task skip adapters, zero-initialised so they are the identity at the start, whose entire purpose is to let each decoder re-specialise the shared skips when they turn out not to suit it. They are off by default, which is the honest way to ship a mitigation for a problem you do not currently have.

Pitfall One model, two outputs, two trajectories

In dual-pass mode the two tasks do not share a sampler run. The shipped config gives dehazing 25 inference steps and depth 20, and those are two separate DDIM loops over two separate latents.

Two things follow that catch people. The outputs are not two views of one trajectory, so they can disagree — the sampled depth need not be the depth of the sampled clear image, and nothing in the architecture enforces that it is. And reproducibility does not survive a mode switch: the model's own documentation warns that because the joint loop draws both latents from a single generator in a fixed order, its random stream differs from the single-task loop's, so seeds are not comparable across joint modes or tasks.

Any experiment that compares dual-pass against single-pass at a fixed seed and attributes the difference to the architecture is measuring the RNG as well. Vary the seed and compare distributions, or compare at a level where the seed does not enter.

Check yourself

Single-pass runs the shared encoder once instead of twice. Does that make single-task inference cheaper on a single-pass model?

Show answer

No, and the model logs a message saying so, which is a good sign that this catches people.

A single-pass model has a twelve-channel first convolution, so it is shape-incompatible with the eight-channel single-task forward path. Asking it for one task routes through the full joint denoising loop — one shared twelve-channel encode per step, both decoders — and returns the branch you asked for, discarding the other. You pay for both tasks and use one.

The saving is real only when you consume both outputs from the one shared trajectory. That is the general shape of the trap with multi-task architectures: the efficiency argument is about the joint workload, and it silently inverts the moment the deployment only wants one of the outputs.