[Bug] Wan MeanFlow zero initialization causes degenerate gradients in r_embedder

[YunkangTao]

Description

In fastgen/configs/experiments/WanT2V/config_mf.py, the Wan MeanFlow model uses:

config.model.net.r_embedder_init = "zero"

This triggers the following code in fastgen/networks/Wan/network.py:

elif embedder_init == "zero":
    for param in embedder.parameters():
        param.data.zero_()

This sets every parameter of the multi-layer r_embedder to zero. As a result, several weight tensors receive zero gradients because of the zero symmetry (only some bias parameters may update). The r-conditioning pathway therefore remains effectively inactive and is difficult to learn.

Reproduction

Run the Wan MeanFlow configuration and inspect the gradients of model.net.transformer.r_embedder after the first backward pass.

The weights of layers such as time_embedder.linear_1 and time_embedder.linear_2.weight remain zero or receive zero gradients.

Expected behavior

The r pathway should initially produce zero output without initializing all of its trainable layers to zero.

Suggested fix

The existing "random" branch appears to implement the intended behavior: it randomly initializes the hidden layers and zero-initializes only the output projections.

Could the "zero" branch be changed to use this initialization strategy, or could the configuration be changed to:

config.model.net.r_embedder_init = "random"

?

[juliusberner]

Hi @YunkangTao,

Thanks for the great catch! With every r-embedder parameter zeroed, linear_2.weight only ever sees
g ⊗ SiLU(linear_1(x)) = 0 and linear_1 only linear_2.weightᵀ g = 0, so both stay at zero
for the whole run and the r path can learn nothing but an r-independent bias. The following fix keeps the
copied hidden layer and zeroes only the output layers:

elif embedder_init == "zero":
    # Zero only the output layers: no-op at init, but the copied hidden layer
    # keeps the gradients flowing (all-zero weights would never leave zero).
    for linear in (embedder.time_embedder.linear_2, embedder.time_proj):
        torch.nn.init.zeros_(linear.weight)
        if linear.bias is not None:
            torch.nn.init.zeros_(linear.bias)

remb and time_proj(act(remb)) are still exactly zero at init and linear_2.weight receives a
gradient on the first backward pass. I will include this fix in the next code update.

Thanks again!