[Bug Report] Gemma 4 MLPs omit `hook_pre`, `hook_pre_linear`, and `hook_post`

[hijohnnylin]

Summary

Gemma3ArchitectureAdapter maps each block's MLP with self._gated_mlp(), which creates a GatedMLPBridge and supplies the standard gated-MLP aliases:

hook_pre        -> gate.hook_out
hook_pre_linear -> in.hook_out
hook_post       -> out.hook_in

Gemma4ArchitectureAdapter maps the same gate_proj / up_proj / down_proj structure to a bare GeneralizedComponent. As a result, Gemma 4 exposes only the MLP boundary hooks (hook_in and hook_out); the three neuron-basis hooks are absent from hook_dict.

This is a mapping gap rather than an architectural limitation: Hugging Face's Gemma4TextMLP is an ordinary gated MLP, and the TransformerLens adapter already identifies all three projection submodules.

Versions

• transformer-lens==3.7.0
• transformers==5.15.0
• Python 3.12

I also checked the v3.7.1 source. The Gemma 4 MLP mapping is unchanged.

Reproduction

This is a config-only reproduction; it does not load model weights.

transformers==5.15.0 also exposes a separate Gemma 4 heterogeneous-config crash in TransformerLens. The two-line opt-in below is used only to get past that independent crash and inspect the adapter mapping. It is not a proposed runtime workaround.

pip install "transformer-lens==3.7.1" "transformers==5.15.0"
python gemma4_mlp_hooks.py

# gemma4_mlp_hooks.py
from transformers import AutoConfig
from transformer_lens.config import TransformerBridgeConfig
from transformer_lens.factories.architecture_adapter_factory import (
    ArchitectureAdapterFactory,
)
from transformer_lens.model_bridge.sources.transformers import (
    determine_architecture_from_hf_config,
    map_default_transformer_lens_config,
)

for model_id in ("google/gemma-3-270m", "google/gemma-4-E2B"):
    hf_config = AutoConfig.from_pretrained(model_id)

    # Isolate the adapter-mapping question from the separate heterogeneous-config crash.
    for config_obj in (hf_config, hf_config.get_text_config()):
        config_obj.allow_global_per_layer_attribute_access = True

    mapped = map_default_transformer_lens_config(hf_config)
    cfg = TransformerBridgeConfig.from_dict(dict(mapped.__dict__))
    cfg.architecture = determine_architecture_from_hf_config(hf_config)
    adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
    mlp = adapter.component_mapping["blocks"].submodules["mlp"]
    aliases = getattr(type(mlp), "hook_aliases", {})

    print(f"{model_id:22s} {type(adapter).__name__}")
    print(
        f"{'':22s} mlp -> {type(mlp).__name__}, "
        f"submodules {sorted(mlp.submodules)}"
    )
    print(f"{'':22s} hook_aliases {aliases or '{}'}")

Observed:

google/gemma-3-270m    Gemma3ArchitectureAdapter
                       mlp -> GatedMLPBridge, submodules ['gate', 'in', 'out']
                       hook_aliases {'hook_pre': 'gate.hook_out', 'hook_pre_linear': 'in.hook_out', 'hook_post': 'out.hook_in'}
google/gemma-4-E2B     Gemma4ArchitectureAdapter
                       mlp -> GeneralizedComponent, submodules ['gate', 'in', 'out']
                       hook_aliases {}

The two adapters identify the same three MLP projections, but only Gemma 3 installs the compatibility aliases.

Expected behavior

For the dense Gemma4TextMLP submodule, these hooks should be available with the same semantics as other gated MLPs:

blocks.N.mlp.hook_pre
blocks.N.mlp.hook_pre_linear
blocks.N.mlp.hook_post

Actual behavior

None of those names is present for Gemma 4. Code written against the Gemma 3 hook interface fails when it attempts to select the corresponding Gemma 4 MLP activation.

Root cause

Gemma 3 uses the gated-MLP helper:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/gemma3.py#L95-L126

GatedMLPBridge defines the expected aliases:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/generalized_components/gated_mlp.py#L45-L62

Gemma 4 instead uses a bare GeneralizedComponent, despite naming gate_proj, up_proj, and down_proj:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/gemma4.py#L135-L142

The Hugging Face module has exactly those three projections:

• https://github.com/huggingface/transformers/blob/v5.15.0/src/transformers/models/gemma4/modeling_gemma4.py#L962-L977

Impact

The missing aliases block the standard TransformerLens MLP-neuron interface on Gemma 4. The failure is loud (hook_dict has no matching key), but it prevents existing activation-caching and intervention code from using the same hook names that work on Gemma 3 and other gated-MLP families.

Suggested fix

Replace the bare MLP mapping with the same helper used by Gemma 3:

"mlp": self._gated_mlp(),

The submodule names already match the helper's defaults.

Suggested regression tests:

1. Assert that all three aliases are present on a Gemma 4 block.
2. Assert that they resolve to gate_proj output, up_proj output, and down_proj input.
3. Cover a standard-width layer and a double-wide MLP layer on E2B/E4B.
4. Cover a model with Gemma 4's optional MoE branch to ensure changing the dense mlp mapping does not disturb the separate router/expert components.