Training in progress - step 1000

asr_config.py

@@ -51,6 +51,12 @@ class ASRConfig(transformers.PretrainedConfig):
         downsample_rate: int = 5,  # Granite default
         projector_hidden_dim: Optional[int] = None,
         projector_type: str = "mlp",  # "mlp", "mosa", "moe", "qformer"
+        # Per-time-step Bernoulli zero-mask on encoder output before the
+        # projector (training-only). 0.05–0.15 is the SpecAugment-equivalent
+        # range for frozen-encoder setups; drops whole encoder frames so
+        # the projector learns robustness to missing context. No magnitude
+        # rescaling. 0.0 disables.
+        audio_token_dropout: float = 0.0,
         # MoE-specific configuration
         num_experts: int = 4,  # Number of experts in MoE projectors
         num_experts_per_tok: int = 2,  # Top-k experts per token
@@ -117,6 +123,7 @@ class ASRConfig(transformers.PretrainedConfig):
         self.downsample_rate = downsample_rate
         self.projector_hidden_dim = projector_hidden_dim
         self.projector_type = projector_type
+        self.audio_token_dropout = audio_token_dropout
         # MoE-specific configuration
         self.num_experts = num_experts
         self.num_experts_per_tok = num_experts_per_tok

asr_modeling.py

@@ -449,11 +449,35 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             encoder_out = self.audio_tower(input_features=audio_features)
             hidden_states = encoder_out.last_hidden_state
 
+        hidden_states = self._maybe_drop_audio_tokens(hidden_states)
         audio_embeds = self.projector(hidden_states)
 
         token_counts = expected_token_counts.to(device=audio_embeds.device, dtype=torch.long)
         return _gather_audio_embeds(audio_embeds, token_counts)
 
+    def _maybe_drop_audio_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Per-time-step Bernoulli zero-mask on encoder output (train-only).
+
+        SpecAugment-equivalent for frozen-encoder setups: drops whole frames
+        from the encoder output sequence so the projector learns robustness
+        to missing context. Length-preserving (zeros, not deletions) so
+        audio token counts in the prompt stay consistent. No magnitude
+        rescaling — the projector should not learn to compensate.
+        """
+        p = float(getattr(self.config, "audio_token_dropout", 0.0))
+        if not self.training or p <= 0.0:
+            return hidden_states
+        keep = 1.0 - p
+        mask = torch.bernoulli(
+            torch.full(
+                hidden_states.shape[:-1],
+                keep,
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+        ).unsqueeze(-1)
+        return hidden_states * mask
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,

config.json

@@ -103,6 +103,7 @@
   },
   "audio_model_id": "zai-org/GLM-ASR-Nano-2512",
   "audio_sample_rate": 16000,
+  "audio_token_dropout": 0.1,
   "auto_map": {
     "AutoConfig": "asr_config.ASRConfig",
     "AutoModel": "asr_modeling.ASRModel",
@@ -340,7 +341,7 @@
   "text_model_id": "Qwen/Qwen3-0.6B",
   "top_k": null,
   "top_p": null,
-  "transformers_version": "5.6.1",
+  "transformers_version": "5.7.0",
   "use_cache": false,
   "use_lora": false,
   "vocab_size": 151670

generation_config.json

@@ -12,6 +12,6 @@
   "num_beams": 1,
   "pad_token_id": 151643,
   "repetition_penalty": 1.0,
-  "transformers_version": "5.6.1",
+  "transformers_version": "5.7.0",
   "use_cache": true
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:333242aa801151639eea5cc44531b6ffd9678d6438423fcba55f0194979c1ceb
+oid sha256:fa4b92ee0d0748185d09b2af6dea12c49805b253fd9b5350f3ae559299148424
 size 2433494416