NVlabs / NVlabs/FastGen

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

Open Beginner friendly
#42 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1k
Forks
83
PR merge metrics
No merged PRs in 30d

Description

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.

Image

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"

?

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with fastgen/configs/experiments/WanT2V/config_mf.py and the "zero" initialization branch in fastgen/networks/Wan/network.py. Run the Wan MeanFlow configuration and inspect r_embedder gradients after the first backward pass. Done means the r pathway initially produces zero output without leaving its trainable hidden layers with degenerate zero gradients.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
78/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.