Training in progress - step 1000

projectors.py

@@ -1,10 +1,9 @@
 """Audio projector modules for bridging encoder and decoder embeddings.
 
 This module contains all projector architectures:
-- LinearProjector: Simple avg pool + linear (Chinese Dialects paper, best for Stage 1)
-- MLPAudioProjector: 2-layer MLP with frame stacking downsampling
-- MOSAProjector: MOSA-style dense mixture of experts (arXiv:2508.18998)
-- MoEAudioProjector: Shared expert + sparse routed experts
+- MLPAudioProjector: Simple 2-layer MLP with frame stacking downsampling
+- MOSAProjector: MOSA-style dense mixture of experts
+- SharedMoEAudioProjector: Shared expert + sparse routed experts
 - QFormerAudioProjector: BLIP-2 QFormer with learnable queries (Granite-style)
 """
 
@@ -16,51 +15,6 @@ import torch.nn.functional as F  # noqa: N812
 from transformers import AutoModel, Blip2QFormerConfig
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
 
-# =============================================================================
-# Linear Projector (Chinese Dialects paper style)
-# =============================================================================
-
-
-class LinearProjector(nn.Module):
-    """Simple linear projector with average pooling downsampling.
-
-    Based on Chinese Dialects paper (arXiv:2505.21138) which found this
-    outperformed Conv1D, Transformer, and Q-Former in Stage 1 (projector-only).
-
-    Architecture: AvgPool(4x) -> Linear(encoder_dim, llm_dim)
-    """
-
-    def __init__(self, config):
-        super().__init__()
-
-        self.encoder_dim = getattr(config, "encoder_dim", 768)
-        self.llm_dim = getattr(config, "llm_dim", 2048)
-        self.pool_stride = getattr(config, "projector_pool_stride", 4)  # 4x = 12.5Hz
-
-        # Single linear projection (no hidden layers, no activation)
-        self.linear = nn.Linear(self.encoder_dim, self.llm_dim)
-
-    def get_output_length(self, input_length: int) -> int:
-        """Calculate output sequence length given input length."""
-        return input_length // self.pool_stride
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: [batch, seq_len, encoder_dim]
-
-        Returns:
-            [batch, seq_len // pool_stride, llm_dim]
-        """
-        # Average pooling for downsampling (better than frame stacking for linear)
-        # Transpose for avg_pool1d: [B, S, D] -> [B, D, S]
-        x = x.transpose(1, 2)
-        x = F.avg_pool1d(x, kernel_size=self.pool_stride, stride=self.pool_stride)
-        x = x.transpose(1, 2)  # [B, S//k, D]
-
-        return self.linear(x)
-
-
 # =============================================================================
 # MLP Projector
 # =============================================================================
@@ -85,18 +39,21 @@ class MLPAudioProjector(nn.Module):
         self.linear_2 = nn.Linear(hidden_dim, llm_dim)
 
     def get_output_length(self, input_length: int) -> int:
-        """Calculate output sequence length given input length."""
-        return input_length // self.k
+        """Calculate output sequence length given input length (matches GLM-ASR)."""
+        # GLM-ASR formula: (L - merge_factor) // merge_factor + 1
+        return (input_length - self.k) // self.k + 1
 
     def forward(self, x):
         """
         x: [Batch, Seq_Len, Dim]
-        Returns: [Batch, Seq_Len // k, llm_dim]
+        Returns: [Batch, (Seq_Len - k) // k + 1, llm_dim]
         """
         batch, seq, dim = x.shape
-        # Reshape to combine k frames: [B, S, D] -> [B, -1, D*k]
-        # -1 infers sequence length, implicitly downsampling by factor k
-        x = x.reshape(batch, -1, dim * self.k)
+        # Truncate to match GLM-ASR: use (seq - k) // k + 1 frames
+        # This drops trailing frames that don't fill a complete k-frame window
+        out_len = (seq - self.k) // self.k + 1
+        x = x[:, : out_len * self.k, :]  # Truncate to exact multiple
+        x = x.reshape(batch, out_len, dim * self.k)
 
         x = self.linear_1(x)
         x = self.act(x)
@@ -109,12 +66,12 @@ class MLPAudioProjector(nn.Module):
 
 
 class SimpleAdapter(nn.Module):
-    """Simple 2-layer ReLU adapter (from MOSA paper, arXiv:2508.18998)."""
+    """Simple 2-layer GELU adapter (from MOSA paper)."""
 
     def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
         super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
-        self.act = nn.ReLU()
+        self.act = nn.GELU()
         self.fc2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -139,27 +96,19 @@ class MOSAProjector(nn.Module):
 
     Based on "MOSA: Mixtures of Simple Adapters" (arXiv:2508.18998).
     Uses softmax gating over all experts (dense MoE) with only cross-entropy loss.
-    Uses conv-based downsampling (2x Conv1d stride-2) as described in the paper.
+    Uses frame-stacking for downsampling (like MLP projector).
     """
 
     def __init__(self, config):
         super().__init__()
         self.encoder_dim = getattr(config, "encoder_dim", None) or 1280
         self.llm_dim = getattr(config, "llm_dim", None) or 2048
+        self.k = getattr(config, "projector_pool_stride", 4)
         self.num_experts = getattr(config, "num_experts", None) or 4  # MOSA-Base uses 4
         adapter_hidden = getattr(config, "adapter_hidden_dim", None) or 4096
 
-        # --- Conv-based downsampling (paper: 2x Conv1d, kernel=3, stride=2) ---
-        # Total 4x downsampling: 50Hz -> 12.5Hz
-        self.conv_downsample = nn.Sequential(
-            nn.Conv1d(self.encoder_dim, self.encoder_dim, kernel_size=3, stride=2, padding=1),
-            nn.ReLU(),
-            nn.Conv1d(self.encoder_dim, self.encoder_dim, kernel_size=3, stride=2, padding=1),
-            nn.ReLU(),
-        )
-
-        # Input dim to adapters is now just encoder_dim (not encoder_dim * k)
-        in_dim = self.encoder_dim
+        # Frame stacking: concat k adjacent frames then project
+        in_dim = self.encoder_dim * self.k
 
         # --- 1. Simple Router (MOSA-Base: 2 layers with ReLU) ---
         # Maps encoder_dim -> 512 -> num_experts
@@ -171,7 +120,7 @@ class MOSAProjector(nn.Module):
         )
 
         # --- 2. Experts (Simple 2-layer GELU adapters) ---
-        # Each expert: encoder_dim -> hidden -> llm_dim
+        # Each expert: in_dim (stacked frames) -> hidden -> llm_dim
         self.experts = nn.ModuleList(
             [SimpleAdapter(in_dim, adapter_hidden, self.llm_dim) for _ in range(self.num_experts)]
         )
@@ -180,27 +129,32 @@ class MOSAProjector(nn.Module):
         # x: (B, S, encoder_dim)
         batch_size, seq_len, dim = x.shape
 
-        # --- 1. Conv downsampling ---
-        # Conv1d expects (B, C, S), so transpose
-        x_conv = x.transpose(1, 2)  # (B, D, S)
-        x_conv = self.conv_downsample(x_conv)  # (B, D, S//4)
-        x_downsampled = x_conv.transpose(1, 2)  # (B, S//4, D)
+        # Truncate to match GLM-ASR: use (seq - k) // k + 1 frames
+        out_len = (seq_len - self.k) // self.k + 1
+        x = x[:, : out_len * self.k, :]
+
+        # --- 1. Router Branch ---
+        # Mean pool encoder outputs for routing decisions
+        x_pooled = x.reshape(batch_size, out_len, self.k, self.encoder_dim).mean(dim=2)  # (B, out_len, D)
 
-        # --- 2. Router Branch ---
-        # Router operates on downsampled features
-        routing_weights = F.softmax(self.router(x_downsampled), dim=-1)  # (B, S//4, num_experts)
+        # Router logits and softmax gating (dense MoE)
+        routing_weights = F.softmax(self.router(x_pooled), dim=-1)  # (B, out_len, num_experts)
+
+        # --- 2. Frame stacking for experts ---
+        # Reshape to combine k frames: [B, S, D] -> [B, out_len, D*k]
+        x_stacked = x.reshape(batch_size, out_len, dim * self.k)
 
         # --- 3. Expert Mixture (Dense Execution) ---
         # Run all experts and compute weighted sum
         expert_outputs = torch.stack(
-            [expert(x_downsampled) for expert in self.experts]
-        )  # (E, B, S//4, D)
+            [expert(x_stacked) for expert in self.experts]
+        )  # (E, B, out_len, D)
         return torch.einsum("ebsd, bse -> bsd", expert_outputs, routing_weights)
 
     def get_output_length(self, input_length: int) -> int:
-        """Calculate output sequence length given input length."""
-        # Two stride-2 convs = 4x downsampling
-        return input_length // 4
+        """Calculate output sequence length given input length (matches GLM-ASR)."""
+        # GLM-ASR formula: (L - merge_factor) // merge_factor + 1
+        return (input_length - self.k) // self.k + 1
 
 
 # =============================================================================
@@ -497,7 +451,6 @@ class QFormerAudioProjector(nn.Module):
 # =============================================================================
 
 PROJECTOR_CLASSES = {
-    "linear": LinearProjector,
     "mlp": MLPAudioProjector,
     "mosa": MOSAProjector,
     "moe": MoEAudioProjector,