Initial upload: Gemma 4 audio encoder (304.8M USM-style Conformer)

README.md

@@ -0,0 +1,165 @@
+---
+language:
+- en
+- multilingual
+license: apache-2.0
+library_name: transformers
+tags:
+- feature-extraction
+- audio
+- speech
+- conformer
+- gemma4
+- usm
+- google
+pipeline_tag: feature-extraction
+base_model: google/gemma-4-E2B-it
+---
+
+# Gemma 4 Audio Encoder (USM-style Conformer)
+
+Standalone extraction of the audio encoder from Google's [Gemma 4](https://huggingface.co/google/gemma-4-E2B-it) multimodal model family. This is a 304.8M parameter USM-style Conformer that converts audio waveforms (via 128-bin mel spectrogram) into embeddings.
+
+**License:** Apache 2.0 (inherited from Gemma 4 — no restrictions)
+
+## Architecture
+
+| Property | Value |
+|---|---|
+| Total parameters | 304.8M |
+| Architecture | USM-style Conformer (Macaron-net) |
+| Hidden dimension | 1024 |
+| Output dimension | 1536 (via `output_proj` Linear + bias) |
+| Conformer layers | 12 |
+| Attention heads | 8 (128 dim per head) |
+| FFW intermediate | 4096 (4× expansion) |
+| Depthwise conv kernel | 5 |
+| Subsampling conv channels | [128, 32] |
+| Input | 128-bin mel spectrogram @ 16kHz |
+| Conformer activation | SiLU |
+| Subsampling activation | ReLU |
+| Conformer normalization | RMSNorm (eps=1e-6) |
+| Subsampling normalization | LayerNorm |
+| Residual weight | 0.5 (Macaron half-step) |
+| Attention type | Chunked causal (chunk_size=12, left_context=13, right_context=0) |
+| Clipped linears | Yes (quantization-ready input_min/max, output_min/max per layer) |
+| Temporal downsampling | 4× (two stride-2 Conv2d layers) |
+
+### Conformer Block Structure
+
+Each of the 12 conformer blocks follows the Macaron-net pattern:
+
+```
+Input
+  → FFW1: pre_layer_norm → Linear(1024→4096) → SiLU → Linear(4096→1024) → post_layer_norm
+  → + 0.5 × residual
+  → Self-Attention: norm_pre_attn → Q/K/V proj (1024→1024) → relative position → post proj → norm_post_attn
+  → + residual
+  → LightConv1d: pre_layer_norm → Linear(1024→2048, gated) → DepthwiseConv1d(k=5) → conv_norm → Linear(1024→1024)
+  → + residual
+  → FFW2: pre_layer_norm → Linear(1024→4096) → SiLU → Linear(4096→1024) → post_layer_norm
+  → + 0.5 × residual
+  → norm_out
+```
+
+### Input/Output Shapes
+
+- **Input:** `(batch, time_frames, 128)` — 128-bin mel features, time-first
+- **Output:** `(batch, time_frames/4, 1536)` — 4× temporal downsampling, projected to 1536
+- For 4 seconds of 16kHz audio: input ~(1, 399, 128) → output ~(1, 100, 1536)
+
+## Usage
+
+```python
+import torch
+import numpy as np
+from safetensors.torch import load_file
+from transformers import Gemma4AudioModel, Gemma4AudioConfig, AutoProcessor
+
+# Load config and weights
+audio_cfg = Gemma4AudioConfig.from_pretrained("rnagabh/gemma4-audio-encoder")
+audio_tower = Gemma4AudioModel(audio_cfg)
+
+state_dict = load_file("path/to/model.safetensors")  # or use from_pretrained
+audio_tower.load_state_dict(state_dict)
+audio_tower = audio_tower.to(dtype=torch.bfloat16, device="cuda")
+audio_tower.eval()
+
+# Use feature extractor from the parent model
+processor = AutoProcessor.from_pretrained("google/gemma-4-E2B-it")
+feature_extractor = processor.feature_extractor
+
+import numpy as np
+waveform = np.random.randn(64000).astype(np.float32)  # 4s @ 16kHz
+inputs = feature_extractor([waveform], sampling_rate=16000, return_tensors="pt")
+
+with torch.no_grad():
+    output = audio_tower(inputs["input_features"].to(dtype=torch.bfloat16, device="cuda"))
+    embeddings = output.last_hidden_state  # (1, 100, 1536)
+```
+
+## Critical: The AutoModel Loading Gotcha
+
+⚠️ **`AutoModel.from_pretrained("google/gemma-4-E2B-it")` silently fails to load audio tower weights.**
+
+All audio tower parameters initialize as random (std ≈ 0.02). The model runs without errors, produces outputs of the correct shape, but the outputs are meaningless.
+
+**Root cause:** The checkpoint stores keys with a `model.` prefix (e.g., `model.audio_tower.layers.0...`). `AutoModel` builds the module tree expecting keys without the prefix. The mismatch causes every key to be both UNEXPECTED and MISSING. Transformers loads with `strict=False` by default, so this silently initializes everything fresh.
+
+**Fix:** Use `AutoModelForMultimodalLM` instead:
+
+```python
+# ❌ WRONG — audio tower weights are randomly initialized
+model = AutoModel.from_pretrained("google/gemma-4-E2B-it")
+audio_tower = model.audio_tower  # RANDOM WEIGHTS
+
+# ✅ CORRECT — audio tower weights load properly
+model = AutoModelForMultimodalLM.from_pretrained("google/gemma-4-E2B-it")
+audio_tower = model.model.audio_tower  # TRAINED WEIGHTS
+```
+
+**How to verify:**
+
+```python
+w = audio_tower.output_proj.weight.float()
+print(f"std={w.std().item():.6f}")
+# ✅ Trained: std ≈ 0.031250
+# ❌ Random:  std ≈ 0.019884
+```
+
+## E2B and E4B Share Identical Audio Weights
+
+The audio encoder weights are **byte-for-byte identical** between Gemma 4 E2B and E4B. This was verified empirically — all 751 parameter tensors match exactly.
+
+Gemma 4's E2B is a MatFormer sub-model nested inside E4B. The MatFormer architecture only affects the text decoder's feed-forward dimensions. The audio tower sits outside the MatFormer nesting and is a shared module.
+
+**Implication:** There is no reason to prefer E4B over E2B for audio encoder extraction. E2B is a smaller download (~10GB vs ~16GB).
+
+## Files in This Repo
+
+| File | Description | Size |
+|---|---|---|
+| `config.json` | Audio tower config (Gemma4AudioConfig) | <1 KB |
+| `model.safetensors` | Audio tower weights (304.8M params, BF16) | 609.7 MB |
+| `preprocessor_config.json` | Mel spectrogram feature extractor config | <1 KB |
+| `embed_audio.safetensors` | Audio→text embedding projection (1536→1536) | 4.7 MB |
+
+## Limitations
+
+- **End-to-end trained for LLM decoding:** The encoder was trained to produce features for Gemma 4's text decoder, not as a general-purpose audio encoder. For standalone feature extraction, the 1024-dim pre-projection output (before `output_proj`) may be more useful than the 1536-dim post-projection output.
+- **Causal chunked attention:** The encoder uses right_context=0, meaning it cannot look ahead. This limits its use in offline/non-streaming settings compared to bidirectional encoders.
+- **Multi-layer fusion doesn't help:** Unlike wav2vec2/W2v-BERT where combining multiple hidden layers improves downstream performance, this encoder's Macaron half-step residuals and causal attention mean only the final layer output is useful.
+- **Subsampling frontend uses ReLU + LayerNorm** (not SiLU + GroupNorm as in some USM descriptions).
+
+## Extraction Details
+
+- Extracted from `google/gemma-4-E2B-it` using `AutoModelForMultimodalLM`
+- Weights saved in BF16 as safetensors
+- Forward pass verified: extracted model produces outputs with **0.0 max absolute difference** from the original
+- All architecture specs independently verified against the live model
+
+## References
+
+- [Gemma 4 on HuggingFace](https://huggingface.co/google/gemma-4-E2B-it)
+- [Gemma 4 Blog Post](https://huggingface.co/blog/gemma4)
+- [Google USM Paper](https://arxiv.org/abs/2303.01037) — "Google USM: Scaling Automatic Speech Recognition Beyond 100 Languages"

config.json

@@ -0,0 +1,64 @@
+{
+  "transformers_version": "5.5.0",
+  "architectures": [
+    "Gemma4AudioModel"
+  ],
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": "bfloat16",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "hidden_size": 1024,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 8,
+  "hidden_act": "silu",
+  "subsampling_conv_channels": [
+    128,
+    32
+  ],
+  "conv_kernel_size": 5,
+  "residual_weight": 0.5,
+  "attention_chunk_size": 12,
+  "attention_context_left": 13,
+  "attention_context_right": 0,
+  "attention_logit_cap": 50.0,
+  "attention_invalid_logits_value": -1000000000.0,
+  "use_clipped_linears": true,
+  "rms_norm_eps": 1e-06,
+  "gradient_clipping": 10000000000.0,
+  "output_proj_dims": 1536,
+  "initializer_range": 0.02,
+  "_name_or_path": "",
+  "model_type": "gemma4_audio",
+  "output_attentions": false,
+  "torch_dtype": "bfloat16",
+  "_verified_total_params": 304824608,
+  "_verified_hidden_dim": 1024,
+  "_verified_output_dim": 1536,
+  "_verified_num_layers": 12,
+  "_verified_num_heads": 8,
+  "_verified_head_dim": 128,
+  "_verified_ffn_intermediate": 4096,
+  "_verified_conv_kernel": 5,
+  "_verified_subsampling_channels": [
+    128,
+    32
+  ],
+  "_verified_subsampling_norm": "LayerNorm",
+  "_verified_subsampling_activation": "ReLU",
+  "_verified_conformer_activation": "SiLU",
+  "_verified_conformer_norm": "RMSNorm",
+  "_verified_conformer_norm_eps": 1e-06,
+  "_verified_temporal_downsample": 4,
+  "_source_model": "google/gemma-4-E2B-it",
+  "_extraction_note": "Audio tower weights are identical between E2B and E4B variants"
+}

embed_audio.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:966aa3b0b2a3513c4d4cfcdb0396fb802d291dfe68edc07f7efaa5487ee8f099
+size 4718696

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecbc8887cb18418036b56b67c6b6ae43f3a75a8d363fb0d63a89307684709988
+size 609732608

preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+  "feature_size": 128,
+  "sampling_rate": 16000,
+  "padding_value": 0.0,
+  "padding_side": "right",
+  "return_attention_mask": true,
+  "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+  "fft_length": 512,
+  "frame_length": 320,
+  "hop_length": 160,
+  "min_frequency": 0.0,
+  "max_frequency": 8000.0,
+  "preemphasis": 0.0,
+  "preemphasis_htk_flavor": true,
+  "fft_overdrive": false,
+  "dither": 0.0,
+  "input_scale_factor": 1.0,
+  "mel_floor": 0.001,
+  "per_bin_mean": null,
+  "per_bin_stddev": null
+}