Fix ShieldGemma2 non-reproducible outputs by adding _tied_weights_keys by hardikmeisheri · Pull Request #44358 · huggingface/transformers

hardikmeisheri · 2026-02-28T16:49:27Z

Summary

ShieldGemma2ForImageClassification was missing _tied_weights_keys, so model.lm_head.weight was randomly re-initialized on every from_pretrained call instead of being tied to embed_tokens.weight.
This caused non-deterministic classification outputs across model loads despite identical inputs.
Fix: add _tied_weights_keys = {"model.lm_head.weight": "model.model.language_model.embed_tokens.weight"} to the class, matching the dict-based pattern used by Gemma3ForConditionalGeneration.

Root Cause

The checkpoint has text_config.tie_word_embeddings = True, so lm_head.weight is not saved separately in the checkpoint — it is expected to be restored by the weight-tying mechanism. The dict-based _tied_weights_keys attribute tells from_pretrained's loading machinery exactly which parameter is tied and what its source is, so it can restore it correctly via tie_weights().

This regression was introduced when the explicit tie_weights() override was removed in #41580 without adding the equivalent _tied_weights_keys declaration.

The load report shows the symptom clearly:

Key                  | Status  |
---------------------+---------+
model.lm_head.weight | MISSING |

Reproduction

import torch
from transformers import AutoProcessor, ShieldGemma2ForImageClassification
from PIL import Image
import requests

model_id = "google/shieldgemma-2-4b-it"
processor = AutoProcessor.from_pretrained(model_id)
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
inputs = processor(images=[image], return_tensors="pt")

Before fix — different lm_head.weight sum and wildly varying probabilities on each load:

Load	lm_head.weight sum	Result
1	568.72	bee/dangerous → Safe:0.0, Unsafe:1.0
2	359.00	bee/dangerous → Safe:0.0003, Unsafe:1.0
3	-883.37	bee/dangerous → Safe:1.0, Unsafe:0.0

After fix — all 3 loads produce identical results with lm_head.weight sum=-105549.42 and zero variance.

Full discussion with benchmarks: https://huggingface.co/google/shieldgemma-2-4b-it/discussions/10

Test plan

Load google/shieldgemma-2-4b-it multiple times and verified model.model.lm_head.weight has the same sum/norm across all loads after fix.
Confirmed model.model.lm_head.weight.data_ptr() == model.model.get_input_embeddings().weight.data_ptr() after loading (weights share memory).

The checkpoint has `text_config.tie_word_embeddings = True`, meaning `lm_head.weight` should be tied to `embed_tokens.weight`. However, `ShieldGemma2ForImageClassification` was missing `_tied_weights_keys`, so `from_pretrained` treated `model.lm_head.weight` as absent and left it randomly initialized on every load — causing non-reproducible outputs. Adding `_tied_weights_keys = {"model.lm_head.weight": "model.model.language_model.embed_tokens.weight"}` lets the loading machinery skip that key and populate it via `tie_weights()` in `from_pretrained`, which uses the dict-based tying mechanism. Fixes: https://huggingface.co/google/shieldgemma-2-4b-it/discussions/10 Co-authored-by: Hardik Meisheri <hardik.meisheri@gmail.com> Co-authored-by: Shrey Ganatra <ganatrashrey2002@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Rocketknight1 · 2026-03-02T15:29:57Z

@hardikmeisheri this PR feels a bit weird and AI-written. The reproduction loads a processor rather than the actual model weights, and when I run my own reproducer code, I still see the same missing lm_head weight and the same variance in weight.sum(). I think the bug is real but the fix is not. If you're going to use Claude, can you ask Opus 4.6 to double-check and make a proper fix? 😅

hardikmeisheri · 2026-03-02T15:35:15Z

@Rocketknight1 So I used sonnet to push the git commit and add details. [I have added Sonnet as the author as well]. The fix is just single line. Warning that you are observing is when the initialization happens. We are tying up the weights after that.

_tied_weights_keys = {"model.lm_head.weight": "model.model.language_model.embed_tokens.weight"}

Can you share your code?

After model loading this line should work

model.model.lm_head.weight = model.model.get_input_embeddings().weight

hardikmeisheri · 2026-03-02T15:39:41Z

Details of the reproduction is here https://huggingface.co/google/shieldgemma-2-4b-it/discussions/10

hardikmeisheri · 2026-03-02T16:57:52Z

@Rocketknight1 Check if you are able to reproduce this with the following code.

import torch, gc
from transformers import AutoProcessor, ShieldGemma2ForImageClassification
from PIL import Image
import requests
from datasets import load_dataset

model_id = "google/shieldgemma-2-4b-it"
processor = AutoProcessor.from_pretrained(model_id)
POLICIES = ["dangerous", "sexual", "violence"]

url_bee = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"
image_bee = Image.open(requests.get(url_bee, stream=True).raw).convert("RGB")

ds_food = load_dataset("ethz/food101", split="validation")
image_food = ds_food[0]["image"].convert("RGB")

ds_cats = load_dataset("huggingface/cats-image", split="test")
image_cat = ds_cats[0]["image"].convert("RGB")

test_images = {"bee": image_bee, "food": image_food, "cat": image_cat}
cached_inputs = {name: processor(images=[img], return_tensors="pt") for name, img in test_images.items()}
print(f"Loaded {len(test_images)} test images: {list(test_images.keys())}")

DEVICE = "cuda:0"

for i in range(3):
    torch.manual_seed(i * 42) 
    model = ShieldGemma2ForImageClassification.from_pretrained(model_id, device_map=DEVICE, torch_dtype=torch.bfloat16).eval()
    lm_sum = model.model.lm_head.weight.data.float().sum().item()
    print(f"\nMODEL LOAD {i+1}/3")
    print(f"  lm_head.weight sum={lm_sum:.6f}")
    for img_name in test_images:
        inp = {k: [v.to](http://v.to/)(device=DEVICE, dtype=torch.bfloat16) if v.is_floating_point()
               else [v.to](http://v.to/)(device=DEVICE) for k, v in cached_inputs[img_name].items()}

        with torch.inference_mode():
            scores = model(**inp)

        # probabilities[:, 0] = "Yes" (violates/UNSAFE), [:, 1] = "No" (SAFE)
        for j, p in enumerate(POLICIES):
            safe  = scores.probabilities[j, 1].item()
            unsafe = scores.probabilities[j, 0].item()
            print(f"  [{img_name:12s}] {p:12s} -> Safe={safe:.6f}  Unsafe={unsafe:.6f}")

    del model; gc.collect(); torch.cuda.empty_cache()

Output

MODEL LOAD 1/3
lm_head.weight sum=-398.408478
[bee ] dangerous -> Safe=1.000000 Unsafe=0.001480
[bee ] sexual -> Safe=0.980469 Unsafe=0.020386
[bee ] violence -> Safe=0.964844 Unsafe=0.033691
[food ] dangerous -> Safe=0.937500 Unsafe=0.063477
[food ] sexual -> Safe=0.679688 Unsafe=0.320312
[food ] violence -> Safe=0.953125 Unsafe=0.045166
[cat ] dangerous -> Safe=0.464844 Unsafe=0.535156
[cat ] sexual -> Safe=0.277344 Unsafe=0.722656
[cat ] violence -> Safe=1.000000 Unsafe=0.001457

MODEL LOAD 2/3
lm_head.weight sum=-23.346558
[bee ] dangerous -> Safe=0.933594 Unsafe=0.065918
[bee ] sexual -> Safe=0.453125 Unsafe=0.546875
[bee ] violence -> Safe=0.933594 Unsafe=0.065430
[food ] dangerous -> Safe=0.021118 Unsafe=0.980469
[food ] sexual -> Safe=0.224609 Unsafe=0.773438
[food ] violence -> Safe=0.937500 Unsafe=0.063477
[cat ] dangerous -> Safe=0.065430 Unsafe=0.933594
[cat ] sexual -> Safe=0.018433 Unsafe=0.980469
[cat ] violence -> Safe=0.074707 Unsafe=0.925781

MODEL LOAD 3/3
lm_head.weight sum=476.709167
[bee ] dangerous -> Safe=1.000000 Unsafe=0.000156
[bee ] sexual -> Safe=1.000000 Unsafe=0.000077
[bee ] violence -> Safe=0.976562 Unsafe=0.022949
[food ] dangerous -> Safe=1.000000 Unsafe=0.000652
[food ] sexual -> Safe=0.953125 Unsafe=0.045410
[food ] violence -> Safe=0.792969 Unsafe=0.207031
[cat ] dangerous -> Safe=0.867188 Unsafe=0.131836
[cat ] sexual -> Safe=0.921875 Unsafe=0.077637
[cat ] violence -> Safe=0.914062 Unsafe=0.087891

for i in range(3):
    torch.manual_seed(i * 42)
    model = ShieldGemma2ForImageClassification.from_pretrained(model_id, device_map=DEVICE, torch_dtype=torch.bfloat16).eval()

    model.model.lm_head.weight = model.model.get_input_embeddings().weight

    lm_sum  = model.model.lm_head.weight.data.float().sum().item()
    emb_ptr = model.model.get_input_embeddings().weight.data_ptr()
    lm_ptr  = model.model.lm_head.weight.data_ptr()
    print(f"\nMODEL LOAD {i+1}/3  (with tie fix)")
    print(f"  lm_head tied to embed_tokens: {lm_ptr == emb_ptr}  |  lm_head sum: {lm_sum:.4f}")

    for img_name in test_images:
        inp = {k: [v.to](http://v.to/)(device=DEVICE, dtype=torch.bfloat16) if v.is_floating_point()
               else [v.to](http://v.to/)(device=DEVICE) for k, v in cached_inputs[img_name].items()}

        with torch.inference_mode():
            scores = model(**inp)

        for j, p in enumerate(POLICIES):
            safe  = scores.probabilities[j, 1].item()
            unsafe = scores.probabilities[j, 0].item()
            print(f"  [{img_name:12s}] {p:12s} -> Safe={safe:.6f}  Unsafe={unsafe:.6f}")

    del model; gc.collect(); torch.cuda.empty_cache()

Output

MODEL LOAD 1/3 (with tie fix)
lm_head tied to embed_tokens: True | lm_head sum: -105549.4219
[bee ] dangerous -> Safe=1.000000 Unsafe=0.000000
[bee ] sexual -> Safe=1.000000 Unsafe=0.000000
[bee ] violence -> Safe=1.000000 Unsafe=0.000000
[food ] dangerous -> Safe=1.000000 Unsafe=0.000005
[food ] sexual -> Safe=1.000000 Unsafe=0.000028
[food ] violence -> Safe=1.000000 Unsafe=0.000000
[cat ] dangerous -> Safe=1.000000 Unsafe=0.000000
[cat ] sexual -> Safe=0.996094 Unsafe=0.004608
[cat ] violence -> Safe=1.000000 Unsafe=0.000000

MODEL LOAD 2/3 (with tie fix)
lm_head tied to embed_tokens: True | lm_head sum: -105549.4219
[bee ] dangerous -> Safe=1.000000 Unsafe=0.000000
[bee ] sexual -> Safe=1.000000 Unsafe=0.000000
[bee ] violence -> Safe=1.000000 Unsafe=0.000000
[food ] dangerous -> Safe=1.000000 Unsafe=0.000005
[food ] sexual -> Safe=1.000000 Unsafe=0.000028
[food ] violence -> Safe=1.000000 Unsafe=0.000000
[cat ] dangerous -> Safe=1.000000 Unsafe=0.000000
[cat ] sexual -> Safe=0.996094 Unsafe=0.004608
[cat ] violence -> Safe=1.000000 Unsafe=0.000000

MODEL LOAD 3/3 (with tie fix)
lm_head tied to embed_tokens: True | lm_head sum: -105549.4219
[bee ] dangerous -> Safe=1.000000 Unsafe=0.000000
[bee ] sexual -> Safe=1.000000 Unsafe=0.000000
[bee ] violence -> Safe=1.000000 Unsafe=0.000000
[food ] dangerous -> Safe=1.000000 Unsafe=0.000005
[food ] sexual -> Safe=1.000000 Unsafe=0.000028
[food ] violence -> Safe=1.000000 Unsafe=0.000000
[cat ] dangerous -> Safe=1.000000 Unsafe=0.000000
[cat ] sexual -> Safe=0.996094 Unsafe=0.004608
[cat ] violence -> Safe=1.000000 Unsafe=0.000000

Rocketknight1 · 2026-03-03T15:17:37Z

Hi @hardikmeisheri, your code snippet here "fixes" the problem by manually setting the weights with the line model.model.lm_head.weight = model.model.get_input_embeddings().weight after model loading. However, I think a proper fix should not require user intervention like that! Instead, the weights should be correctly loaded in the from_pretrained() call, because most users will not know that they need to do the extra step. This is what I meant by the bug being real but the fix being insufficient.

hardikmeisheri · 2026-03-03T15:33:02Z

@Rocketknight1 I think there is a confusion here. Inline code is for the bug confirmation and fix.

model.model.lm_head.weight = model.model.get_input_embeddings().weight

For PR I have made the changes as follows

_tied_weights_keys = {"model.lm_head.weight": "model.model.language_model.embed_tokens.weight"}

Which basically fixes it during the loading itself. for the class ShieldGemma2ForImageClassification.

Can you check the diff for the commit.

Rocketknight1 · 2026-03-04T14:53:51Z

Yes, I have checked the diff. The point I keep making is that your PR does not actually fix the problem! Your code snippets do fix the problem, because they manually copy the weight, but your PR changes do not.

You can try this yourself - install from your PR branch and then just run this:

import torch
from transformers import ShieldGemma2ForImageClassification

model_id = "google/shieldgemma-2-4b-it"
model = ShieldGemma2ForImageClassification.from_pretrained(model_id)
print(model.model.lm_head.weight.sum())

You will see that the LM head is reinitialized with random weights every time, and so this PR has not actually fixed loading.

…uter config

hardikmeisheri · 2026-03-05T11:57:14Z

@Rocketknight1 Ohh Apologies. i had made the changes during debugging and that got into gitignore and i did not realize it.

I have created a another commit, where along with previous commit we need to put in placeholder values and condition to pick that otherwise it completely ignores it.

here is code and reproducibility

import torch
from transformers import ShieldGemma2ForImageClassification

model_id = "google/shieldgemma-2-4b-it"

print("=== Test 1: test ===")
print("Loading model 3 times with different seeds to check if lm_head.weight changes...")
print()

for i in range(3):
    torch.manual_seed(i * 42)
    model = ShieldGemma2ForImageClassification.from_pretrained(
        model_id, device_map="cpu", torch_dtype=torch.bfloat16
    )
   
    lm_sum = model.model.lm_head.weight.data.float().sum().item()
    emb_sum = model.model.get_input_embeddings().weight.data.float().sum().item()
   
    lm_ptr = model.model.lm_head.weight.data_ptr()
    emb_ptr = model.model.get_input_embeddings().weight.data_ptr()
    tied = lm_ptr == emb_ptr
   
    print(f"Load {i+1}/3:")
    print(f"  lm_head.weight sum  = {lm_sum:.4f}")
    print(f"  embed_tokens sum    = {emb_sum:.4f}")
    print(f"  Same pointer (tied) = {tied}")
    print(f"  Sums match          = {abs(lm_sum - emb_sum) < 0.01}")
    print()
   
    del model
    torch.cuda.empty_cache() if torch.cuda.is_available() else None

print("=== Test 2: Check all_tied_weights_keys on model instance ===")
torch.manual_seed(0)
model = ShieldGemma2ForImageClassification.from_pretrained(
    model_id, device_map="cpu", torch_dtype=torch.bfloat16
)
print(f"model.all_tied_weights_keys = {model.all_tied_weights_keys}")
print()

# Check the inner model's tied weights too
inner = model.model
if hasattr(inner, 'all_tied_weights_keys'):
    print(f"model.model.all_tied_weights_keys = {inner.all_tied_weights_keys}")

print()
print("=== Test 3: Check config.tie_word_embeddings at different levels ===")
print(f"model.config.tie_word_embeddings = {getattr(model.config, 'tie_word_embeddings', 'NOT SET')}")
print(f"model.config.text_config.tie_word_embeddings = {getattr(model.config.text_config, 'tie_word_embeddings', 'NOT SET')}")

Output

=== Test 1: test ===
Loading model 3 times with different seeds to check if lm_head.weight changes...

config.json: 100%|█████████████████████████| 1.57k/1.57k [00:00<00:00, 4.13MB/s]
model.safetensors.index.json: 100%|████████| 95.9k/95.9k [00:00<00:00, 5.26MB/s]
Fetching 2 files: 100%|██████████████████████████| 2/2 [10:53<00:00, 326.90s/it]
Download complete: 100%|███████████████████| 8.60G/8.60G [10:53<00:00, 13.2MB/s]
Loading weights: 100%|█| 883/883 [00:00<00:00, 2006.82it/s, Materializing param=model.m
Load 1/3:
lm_head.weight sum = -105549.4219
embed_tokens sum = -105549.4219
Same pointer (tied) = True
Sums match = True

Fetching 2 files: 100%|███████████████████████████████| 2/2 [00:00<00:00, 21845.33it/s]
Download complete: : 0.00B [00:00, ?B/s] | 0/2 [00:00<?, ?it/s]
Loading weights: 100%|█| 883/883 [00:00<00:00, 2439.81it/s, Materializing param=model.m
Load 2/3:
lm_head.weight sum = -105549.4219
embed_tokens sum = -105549.4219
Same pointer (tied) = True
Sums match = True

Fetching 2 files: 100%|███████████████████████████████| 2/2 [00:00<00:00, 31418.01it/s]
Download complete: : 0.00B [00:00, ?B/s] | 0/2 [00:00<?, ?it/s]
Loading weights: 100%|█| 883/883 [00:00<00:00, 2469.03it/s, Materializing param=model.m
Load 3/3:
lm_head.weight sum = -105549.4219
embed_tokens sum = -105549.4219
Same pointer (tied) = True
Sums match = True

=== Test 2: Check all_tied_weights_keys on model instance ===
Fetching 2 files: 100%|████████████████████████████████| 2/2 [00:00<00:00, 4403.47it/s]
Download complete: : 0.00B [00:00, ?B/s] | 0/2 [00:00<?, ?it/s]
Loading weights: 100%|█| 883/883 [00:00<00:00, 2344.33it/s, Materializing param=model.m
model.all_tied_weights_keys = {'model.lm_head.weight': 'model.model.language_model.embed_tokens.weight'}

model.model.all_tied_weights_keys = {'lm_head.weight': 'model.language_model.embed_tokens.weight'}

=== Test 3: Check config.tie_word_embeddings at different levels ===
model.config.tie_word_embeddings = True
model.config.text_config.tie_word_embeddings = True

Rocketknight1

Yes, LGTM now but with one fix!

Co-authored-by: Matt <Rocketknight1@users.noreply.github.com>

Rocketknight1 · 2026-03-06T14:05:39Z

Sorry @hardikmeisheri - this is almost ready to merge now, but there's one merge conflict because another PR updated the shieldgemma config. Can you resolve that and ping me and I'll merge it? And thank you for the PR!

hardikmeisheri · 2026-03-06T14:50:01Z

@Rocketknight1 I have resolved the conflict. Workflow is pending at PR doc build. Which I assume you need to approve to proceed further.

Thank for for spotting the mistake in the initial commit. Cheers!

github-actions · 2026-03-06T20:40:06Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: shieldgemma2

Rocketknight1 · 2026-03-09T14:18:00Z

Yes, LGTM now. Merging!

HuggingFaceDocBuilderDev · 2026-03-09T14:28:02Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

felifri · 2026-03-16T20:02:09Z

thanks a ton!

redpanda1995 approved these changes Feb 28, 2026

View reviewed changes

hardikmeisheri force-pushed the fix/shieldgemma2-tied-weights branch from 43245c0 to 713f124 Compare February 28, 2026 17:11

Fix ShieldGemma2 weight tying by propagating tie_word_embeddings to o…

82baf98

…uter config

Adding doc strings for tie_word_embeddings for shieldGemma2

efc9bda

Rocketknight1 approved these changes Mar 5, 2026

View reviewed changes

Comment thread src/transformers/models/shieldgemma2/modeling_shieldgemma2.py Outdated

Update src/transformers/models/shieldgemma2/modeling_shieldgemma2.py

6a25d07

Co-authored-by: Matt <Rocketknight1@users.noreply.github.com>

Merge branch 'main' into fix/shieldgemma2-tied-weights

e2ae3ab

Merge branch 'main' into fix/shieldgemma2-tied-weights

baa7d29

Rocketknight1 enabled auto-merge (squash) March 9, 2026 14:18

Rocketknight1 merged commit 7d6fa93 into huggingface:main Mar 9, 2026
20 checks passed

Uh oh!

Conversation

hardikmeisheri commented Feb 28, 2026

Summary

Root Cause

Reproduction

Test plan

Uh oh!

Rocketknight1 commented Mar 2, 2026

Uh oh!

hardikmeisheri commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hardikmeisheri commented Mar 2, 2026

Uh oh!

hardikmeisheri commented Mar 2, 2026

Output

Output

Uh oh!

Rocketknight1 commented Mar 3, 2026

Uh oh!

hardikmeisheri commented Mar 3, 2026

Uh oh!

Rocketknight1 commented Mar 4, 2026

Uh oh!

hardikmeisheri commented Mar 5, 2026

Uh oh!

Rocketknight1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Rocketknight1 commented Mar 6, 2026

Uh oh!

hardikmeisheri commented Mar 6, 2026

Uh oh!

github-actions Bot commented Mar 6, 2026

Uh oh!

Rocketknight1 commented Mar 9, 2026

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Mar 9, 2026

Uh oh!

felifri commented Mar 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

hardikmeisheri commented Mar 2, 2026 •

edited

Loading