mazesmazes/tiny-audio-next-plus
05
1"""Audio projector modules for bridging encoder and decoder embeddings.2 3This module contains all projector architectures:4- MLPAudioProjector: Simple 2-layer MLP with frame stacking downsampling5- MOSAProjector: MOSA-style dense mixture of experts6- SharedMoEAudioProjector: Shared expert + sparse routed experts7- QFormerAudioProjector: BLIP-2 QFormer with learnable queries (Granite-style)8"""9 10import math11 12import torch13import torch.nn as nn14import torch.nn.functional as F # noqa: N81215from transformers import AutoModel, Blip2QFormerConfig16from transformers.models.llama.modeling_llama import LlamaRMSNorm17 18# =============================================================================19# MLP Projector20# =============================================================================21 22 23class MLPAudioProjector(nn.Module):24 """2-layer MLP projector with frame-stacking downsampling (matches GLM-ASR).25 26 Both RMSNorms use LlamaRMSNorm's default weight=1.0 init. A prior version27 initialized both to 0.029 (Qwen3-0.6B's embed_tokens RMS) to put projector28 outputs at residual-stream scale on step 1. Empirically, after training the29 model drifted both norms back to ~1.0 (norm) and ~1.2 (norm_2) — the small30 init wasted compute on a 35× scale-correction phase the optimizer would31 have skipped from default init.32 """33 34 def __init__(self, config):35 """Initialize MLP projector.36 37 Args:38 config: ASRConfig with encoder_dim, llm_dim, projector_pool_stride39 """40 super().__init__()41 42 encoder_dim = getattr(config, "encoder_dim", 768)43 llm_dim = getattr(config, "llm_dim", 2048)44 self.k = getattr(config, "projector_pool_stride", 4)45 46 # Frame stacking: concat k adjacent frames then project47 in_dim = encoder_dim * self.k48 # Hidden dim defaults to llm_dim, can be overridden via config49 hidden_dim = getattr(config, "projector_hidden_dim", None) or llm_dim50 self.linear_1 = nn.Linear(in_dim, hidden_dim, bias=False)51 self.norm = LlamaRMSNorm(hidden_dim, eps=1e-6)52 self.act = nn.GELU()53 self.linear_2 = nn.Linear(hidden_dim, llm_dim, bias=False)54 self.norm_2 = LlamaRMSNorm(llm_dim, eps=1e-6)55 56 def get_output_length(self, input_length: int) -> int:57 """Calculate output sequence length given input length (matches GLM-ASR)."""58 # GLM-ASR formula: (L - merge_factor) // merge_factor + 159 return (input_length - self.k) // self.k + 160 61 def forward(self, x: torch.Tensor) -> torch.Tensor:62 """Project audio features to LLM embedding space.63 64 Args:65 x: Audio encoder output of shape [batch, seq_len, encoder_dim]66 67 Returns:68 Projected features of shape [batch, (seq_len - k) // k + 1, llm_dim]69 """70 x = _frame_stack(x, self.k)71 x = self.linear_1(x)72 x = self.norm(x)73 x = self.act(x)74 x = self.linear_2(x)75 return self.norm_2(x)76 77 78# =============================================================================79# MoE Projector (MOSA-style)80# =============================================================================81 82 83def _frame_stack(x: torch.Tensor, k: int) -> torch.Tensor:84 """Stack k adjacent frames along the feature dim.85 86 Truncates trailing frames that don't fill a complete k-frame window,87 matching GLM-ASR's `(seq_len - k) // k + 1` formula.88 """89 batch, seq, dim = x.shape90 out_len = (seq - k) // k + 191 return x[:, : out_len * k, :].reshape(batch, out_len, dim * k)92 93 94class SimpleAdapter(nn.Module):95 """Simple 2-layer GELU adapter (from MOSA paper)."""96 97 def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):98 super().__init__()99 self.fc1 = nn.Linear(input_dim, hidden_dim)100 self.act = nn.GELU()101 self.fc2 = nn.Linear(hidden_dim, output_dim)102 103 def forward(self, x: torch.Tensor) -> torch.Tensor:104 return self.fc2(self.act(self.fc1(x)))105 106 107class MOSAProjector(nn.Module):108 """MOSA-Base projector: simple 2-layer ReLU router with 4 simple adapters.109 110 Based on "MOSA: Mixtures of Simple Adapters" (arXiv:2508.18998).111 Uses softmax gating over all experts (dense MoE) with only cross-entropy loss.112 Uses Conv1d for downsampling (2 layers, stride 2 each = 4x total).113 """114 115 ADAPTER_HIDDEN_DIM = 4096116 ROUTER_HIDDEN_DIM = 512117 CONV_KERNEL = 3118 CONV_STRIDE = 2119 CONV_PADDING = 1120 121 def __init__(self, config):122 """Initialize MOSA projector.123 124 Args:125 config: ASRConfig with encoder_dim, llm_dim, num_experts126 """127 super().__init__()128 self.encoder_dim = getattr(config, "encoder_dim", None) or 1280129 self.llm_dim = getattr(config, "llm_dim", None) or 2048130 self.num_experts = getattr(config, "num_experts", None) or 4 # MOSA-Base uses 4131 132 conv_kwargs = {133 "kernel_size": self.CONV_KERNEL,134 "stride": self.CONV_STRIDE,135 "padding": self.CONV_PADDING,136 }137 self.downsampler = nn.Sequential(138 nn.Conv1d(self.encoder_dim, self.encoder_dim, **conv_kwargs),139 nn.GELU(),140 nn.Conv1d(self.encoder_dim, self.llm_dim, **conv_kwargs),141 nn.GELU(),142 )143 144 self.router = nn.Sequential(145 nn.Linear(self.llm_dim, self.ROUTER_HIDDEN_DIM),146 nn.ReLU(),147 nn.Linear(self.ROUTER_HIDDEN_DIM, self.num_experts),148 )149 150 self.experts = nn.ModuleList(151 [152 SimpleAdapter(self.llm_dim, self.ADAPTER_HIDDEN_DIM, self.llm_dim)153 for _ in range(self.num_experts)154 ]155 )156 157 def forward(self, x: torch.Tensor) -> torch.Tensor:158 """Project audio features using mixture of experts.159 160 Args:161 x: Audio encoder output of shape [batch, seq_len, encoder_dim]162 163 Returns:164 Projected features of shape [batch, out_len, llm_dim]165 """166 x = self.downsampler(x.transpose(1, 2)).transpose(1, 2)167 168 routing_weights = F.softmax(self.router(x), dim=-1) # (B, out_len, num_experts)169 170 # Accumulate weighted expert outputs without materializing all experts at once.171 output = self.experts[0](x) * routing_weights[..., 0:1]172 for i, expert in enumerate(self.experts[1:], start=1):173 output = output + expert(x) * routing_weights[..., i : i + 1]174 return output175 176 def get_output_length(self, input_length: int) -> int:177 """Calculate output sequence length after Conv1d downsampling (4x reduction)."""178 length = input_length179 for _ in range(2):180 length = (length + 2 * self.CONV_PADDING - self.CONV_KERNEL) // self.CONV_STRIDE + 1181 return length182 183 184# =============================================================================185# MoE Projector (Pure PyTorch with Shared Expert)186# =============================================================================187 188 189class MoEAudioProjector(nn.Module):190 """MoE projector with shared expert (DeepSeek-style), pure PyTorch implementation.191 192 Uses 4 sparse experts with top-2 routing plus a shared expert that processes all tokens.193 No external dependencies (megablocks removed).194 195 Architecture matches main branch: norm → experts(in_dim → hidden → out_dim)196 """197 198 def __init__(self, config):199 """Initialize MoE projector.200 201 Args:202 config: ASRConfig with encoder_dim, llm_dim, num_experts, num_experts_per_tok203 """204 super().__init__()205 206 self.k = getattr(config, "projector_pool_stride", 4)207 self.aux_coef = getattr(config, "router_aux_loss_coef", 0.01)208 209 # Stability coefficients210 self.router_z_loss_coef = getattr(211 config, "router_z_loss_coef", 1e-4212 ) # Prevents logit explosion213 self.router_jitter_noise = getattr(214 config, "router_jitter_noise", 0.01215 ) # Prevents expert collapse216 217 in_dim = config.encoder_dim * self.k218 out_dim = config.llm_dim219 220 # Expert hidden dim (default = output dim)221 hidden_dim = getattr(config, "projector_hidden_dim", None) or out_dim222 223 # Number of experts and top-k selection224 self.num_experts = getattr(config, "num_experts", 4)225 self.top_k = getattr(config, "num_experts_per_tok", 2)226 227 # A. Normalize stacked input (like main branch SharedMoEBlock)228 self.norm = LlamaRMSNorm(in_dim, eps=1e-6)229 230 # B. Router (operates on stacked input)231 self.router = nn.Linear(in_dim, self.num_experts, bias=False)232 233 # C. Experts: simple 2-layer MLP (same as MLPAudioProjector)234 self.experts = nn.ModuleList(235 [SimpleAdapter(in_dim, hidden_dim, out_dim) for _ in range(self.num_experts)]236 )237 238 # D. Shared Expert (same architecture)239 self.shared_expert = SimpleAdapter(in_dim, hidden_dim, out_dim)240 241 # E. Initialize weights for stable training242 self._init_weights()243 244 self.last_aux_loss = torch.tensor(0.0)245 246 def _init_weights(self):247 """Initialize weights for stable training start."""248 with torch.no_grad():249 # Router: small weights -> uniform probability250 nn.init.normal_(self.router.weight, mean=0.0, std=0.02)251 252 # Experts: xavier for fc1, small for fc2 (output)253 for expert in [self.shared_expert, *self.experts]:254 nn.init.xavier_uniform_(expert.fc1.weight)255 nn.init.normal_(expert.fc2.weight, mean=0.0, std=0.01) # Small init256 257 def get_output_length(self, input_length: int) -> int:258 """Calculate output sequence length given input length (matches MLP projector)."""259 return (input_length - self.k) // self.k + 1260 261 def forward(self, x: torch.Tensor) -> torch.Tensor:262 """Project audio features using shared + sparse MoE.263 264 Args:265 x: Audio encoder output of shape [batch, seq_len, encoder_dim]266 267 Returns:268 Projected features of shape [batch, out_len, llm_dim]269 """270 x = _frame_stack(x, self.k)271 batch, out_len, _ = x.shape272 273 # Normalize stacked input (like main branch SharedMoEBlock)274 x = self.norm(x)275 flat_x = x.view(-1, x.size(-1)) # [tokens, in_dim]276 277 # 3. Shared Expert (compute first, creates output tensor)278 output = self.shared_expert(flat_x)279 280 # 4. Sparse Experts (in-place add to shared output)281 self.last_aux_loss = self._forward_sparse(flat_x, output)282 283 return output.view(batch, out_len, -1)284 285 def _forward_sparse(self, x: torch.Tensor, output: torch.Tensor) -> torch.Tensor:286 """Stability-hardened sparse expert dispatch (in-place add to output).287 288 Args:289 x: Flattened input of shape [tokens, dim]290 output: Output tensor to add sparse expert results into (in-place)291 292 Returns:293 Auxiliary loss tensor294 """295 # A. Router Logic with Jitter296 logits = self.router(x)297 298 if self.training and self.router_jitter_noise > 0:299 # Jitter: multiply by uniform noise (1-eps, 1+eps) to shake decision boundary300 # Prevents router from getting stuck on one expert early in training301 noise = torch.empty_like(logits).uniform_(302 1.0 - self.router_jitter_noise, 1.0 + self.router_jitter_noise303 )304 logits = logits * noise305 306 # Force float32 for softmax (bf16/fp16 exponentials can overflow)307 probs = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(x)308 309 # B. Top-K Selection310 top_k_weights, top_k_indices = torch.topk(probs, self.top_k, dim=-1)311 312 # Normalize weights so they sum to 1.0313 top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-6)314 315 # C. Aux Loss + Z-Loss316 aux_loss = torch.tensor(0.0, device=x.device)317 318 if self.training:319 # Load balancing loss (batch-size invariant)320 prob_per_expert = probs.mean(0) # [num_experts]321 target = 1.0 / self.num_experts322 balance_loss = (323 self.aux_coef * ((prob_per_expert - target) ** 2).mean() * self.num_experts324 )325 326 # Z-loss: penalty on large logits to prevent softmax saturation327 z_loss = self.router_z_loss_coef * torch.logsumexp(logits, dim=-1).pow(2).mean()328 329 aux_loss = balance_loss + z_loss330 331 # D. Dispatch Loop (in-place add to output)332 for i, expert in enumerate(self.experts):333 # Create boolean mask for tokens that selected Expert 'i'334 mask = top_k_indices == i335 336 if mask.any():337 # token_idx = which tokens, k_idx = 1st or 2nd choice338 token_idx, k_idx = torch.where(mask)339 340 # Gather inputs and compute341 expert_input = x[token_idx]342 expert_output = expert(expert_input)343 344 # Apply routing weight345 weight = top_k_weights[token_idx, k_idx].unsqueeze(-1)346 weighted_output = (expert_output * weight).type_as(output)347 348 # Scatter back in-place (index_add_ is atomic and deterministic)349 output.index_add_(0, token_idx, weighted_output)350 351 return aux_loss352 353 def get_aux_loss(self) -> torch.Tensor:354 """Return auxiliary load balancing loss."""355 return self.last_aux_loss356 357 358# =============================================================================359# QFormer Projector (Granite-style)360# =============================================================================361 362 363class QFormerAudioProjector(nn.Module):364 """365 BLIP-2 QFormer projector with learnable queries.366 367 Based on GraniteSpeechEncoderProjector - uses a QFormer model with learnable368 query embeddings to compress and project audio encoder outputs. The audio369 sequence is processed in windows and downsampled via cross-attention.370 """371 372 def __init__(self, config):373 """Initialize QFormer projector.374 375 Args:376 config: ASRConfig with encoder_dim, llm_dim, qformer_* settings377 """378 super().__init__()379 380 encoder_dim = config.encoder_dim381 llm_dim = config.llm_dim382 383 # Window and downsampling parameters (Granite defaults: window=15, downsample=5)384 self.window_size = getattr(config, "qformer_window_size", 15)385 self.downsample_rate = getattr(config, "downsample_rate", 5)386 self.num_queries = self.window_size // self.downsample_rate387 388 # QFormer hidden size (matches encoder for cross-attention)389 qformer_hidden = getattr(config, "qformer_hidden_size", None) or encoder_dim390 qformer_num_layers = getattr(config, "qformer_num_layers", 2)391 qformer_num_heads = getattr(config, "qformer_num_heads", 16)392 qformer_intermediate = getattr(config, "qformer_intermediate_size", None) or (393 qformer_hidden * 4394 )395 396 # Learnable query embeddings (Granite uses std=1.0)397 self.query = nn.Parameter(torch.zeros(1, self.num_queries, qformer_hidden))398 self.query.data.normal_(mean=0.0, std=1.0)399 400 # Optional projection if encoder dim != qformer hidden401 if encoder_dim != qformer_hidden:402 self.encoder_proj = nn.Linear(encoder_dim, qformer_hidden, bias=False)403 else:404 self.encoder_proj = None405 406 # Configure QFormer to match Granite's exact config407 qformer_config = Blip2QFormerConfig(408 hidden_size=qformer_hidden,409 num_hidden_layers=qformer_num_layers,410 num_attention_heads=qformer_num_heads,411 intermediate_size=qformer_intermediate,412 encoder_hidden_size=qformer_hidden,413 cross_attention_frequency=1,414 # Granite-specific settings415 hidden_act="gelu",416 attention_probs_dropout_prob=0.1,417 hidden_dropout_prob=0.1,418 layer_norm_eps=1e-12,419 initializer_range=0.02,420 )421 self.qformer = AutoModel.from_config(qformer_config)422 423 # Final projection to LLM dimension (Granite uses bias=True)424 self.linear = nn.Linear(qformer_hidden, llm_dim)425 426 def get_output_length(self, input_length):427 """Calculate output sequence length given input length.428 429 Accepts either Python ints or torch tensors; uses ceiling division so430 the formula is identical for both — math.ceil would block tensors.431 """432 nblocks = (input_length + self.window_size - 1) // self.window_size433 return nblocks * self.num_queries434 435 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:436 """437 Args:438 hidden_states: [batch_size, seq_len, encoder_dim]439 440 Returns:441 projected: [batch_size, num_output_tokens, llm_dim]442 """443 batch_size, seq_len, dim = hidden_states.size()444 445 # Ensure float dtype for QFormer446 target_dtype = self.query.dtype447 if hidden_states.dtype != target_dtype:448 hidden_states = hidden_states.to(target_dtype)449 450 # Optional encoder projection451 if self.encoder_proj is not None:452 hidden_states = self.encoder_proj(hidden_states)453 454 # Compute number of windows and pad to fit455 nblocks = math.ceil(seq_len / self.window_size)456 pad = nblocks * self.window_size - seq_len457 if pad > 0:458 hidden_states = F.pad(hidden_states, (0, 0, 0, pad), "constant", 0)459 460 # Reshape to process each window: [batch*nblocks, window_size, dim]461 effective_batch = batch_size * nblocks462 hidden_states = hidden_states.view(effective_batch, self.window_size, -1)463 464 # Expand queries to match batch size465 query_embeds = self.query.expand(effective_batch, -1, -1)466 467 # QFormer cross-attention468 query_output = self.qformer(469 query_embeds=query_embeds,470 encoder_hidden_states=hidden_states,471 return_dict=True,472 )473 474 # Reshape back: [batch, nblocks * num_queries, hidden]475 output_tokens = nblocks * self.num_queries476 query_proj = query_output.last_hidden_state.view(batch_size, output_tokens, -1)477 478 # Project to LLM dimension479 return self.linear(query_proj)480 481 482# =============================================================================483# Projector Registry484# =============================================================================485 486PROJECTOR_CLASSES = {487 "mlp": MLPAudioProjector,488 "mosa": MOSAProjector,489 "moe": MoEAudioProjector,490 "qformer": QFormerAudioProjector,491}492