[Bug] Wan MeanFlow zero initialization causes degenerate gradients in r_embedder
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.
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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