Repository navigation
Fix Mask2Former _init_weights to match Meta's original initialization - #50596
Open
KMaruthi2002 wants to merge 1 commit into
Open
KMaruthi2002 wants to merge 1 commit into
KMaruthi2002 wants to merge 1 commit into
Conversation
Fix three initialization bugs where _init_weights diverged from Meta's original Mask2Former code: 1. nn.Embedding layers received normal_(std=0.02) from the base PreTrainedModel._init_weights, but Meta's original uses PyTorch's default nn.Embedding init (std=1.0). Now explicitly initialized with normal_(std=1.0). 2. MLP prediction head biases were zeroed by the base _init_weights. Meta's original _reset_parameters only applies xavier_uniform_ on tensors where p.dim() > 1, leaving biases at PyTorch's default Kaiming uniform. Now handled by a dedicated Mask2FormerPredictionBlock branch. 3. Mask2FormerPixelDecoder.level_embed was initialized with normal_(std=0), effectively zeroing it. Meta's original does not explicitly initialize level_embed. Changed to normal_(std=1.0). Moved super()._init_weights() from a blanket call at the top to an else branch, so it only runs for modules not handled by Mask2Former-specific logic. Uses transformers.initialization helpers throughout. Added Mask2FormerInitWeightsTest with four regression tests verifying embedding std, MLP biases, level_embed, and deformable attention init. Fixes huggingface#35877
Contributor
|
Thank you for your contribution 🤗! CI Security Gate — waiting for maintainer approvalCI is approved automatically only for authors who already have a contribution merged in this repository. A maintainer will review this PR and approve CI to run on it. |
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: mask2former |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #35877. The _init_weights method in Mask2FormerPreTrainedModel had three bugs where the initialization diverged from Meta’s original Mask2Former code:
Bug 1: nn.Embedding std was 0.02 instead of PyTorch default (1.0)
super().init_weights() applied normal(std=init_std) where init_std=0.02 to all nn.Embedding layers (query_feat, query_embed). Meta’s original creates these embeddings without explicit initialization, meaning they use PyTorch’s default nn.Embedding init (normal(0, 1)). The embeddings were 50x too small.
Bug 2: MLP biases incorrectly zeroed
super()._init_weights() zeros all nn.Linear biases. Meta’s original reset_parameters only applies xavier_uniform to tensors where p.dim() > 1, leaving biases at PyTorch’s default Kaiming uniform init. This affected Mask2FormerMLPPredictionHead (the mask embedding MLP).
Bug 3: level_embed zeroed by normal_(std=0)
Mask2FormerPixelDecoder.level_embed was initialized with init.normal_(module.level_embed, std=0), which zeroes the parameter. Meta’s original does not explicitly initialize level_embed, so it should keep its constructor value.
Changes:
Removed the blanket super().init_weights(module) call at the top; it is now called only in the else branch for modules not handled by Mask2Former-specific logic
Added explicit nn.Embedding branch with normal(std=1.0) to match PyTorch default
Added Mask2FormerPredictionBlock branch that applies xavier_uniform_ on weights and fan_in-based uniform on biases (matching Meta’s reset_parameters)
Changed Mask2FormerPixelDecoder branch from normal(std=0) to normal_(std=1.0) to match PyTorch default
Uses transformers.initialization helpers throughout, as requested by @ArthurZucker
Regression tests
Added Mask2FormerInitWeightsTest with four tests:
test_embedding_std_matches_pytorch_default: verifies embedding std >> 0.02
test_mlp_biases_not_zeroed: verifies mask_embed MLP has non-zero biases
test_level_embed_not_zeroed: verifies level_embed is not all zeros
test_deformable_attention_sampling_offsets_init: verifies the grid init pattern is preserved
Reference
Meta’s original init: mask2former_transformer_decoder.py uses reset_parameters with xavier_uniform on p.dim() > 1 and no explicit embedding/bias init.