← Walkthroughmapa/models/stem.py
Source snapshot · source SHA-256 b38f3e0ed0e3
1"""Per-band token embeddings from normalized magnitude STFT inputs.23Main path: PerBandStem.forward().45Input: three normalized clip tensors (B,C,F_b,T32).6Output: three token tensors (B,C,T_b,d), plus shared-clock positions.7Next: pack_r4.pack_band_tokens() combines the bands into a flat sequence.89Input order is slow, mid, fast. All inputs use a 32 Hz frame clock; strided10slicing reduces the token rates to 4, 16, and 32 Hz. Each band's linear layer11maps its frequency bins to the model width, then adds a learned band embedding.12"""1314from __future__ import annotations1516from collections.abc import Sequence1718import torch19from torch import Tensor, nn2021from mapa.models.pe import init_transformer_weights2223# (frequency bins, time stride), in slow/mid/fast order.24PER_BAND_SPECS: tuple[tuple[int, int], ...] = ((7, 8), (6, 2), (7, 1))25INPUT_CLIP_Z = (15.0, 15.0, 20.0) # Slow, Mid, Fast: published Guard 3 caps.262728def clip_band(x: Tensor, band: int) -> Tensor:29 """Bound normalized model inputs without modifying the caller's tensor."""30 cap = INPUT_CLIP_Z[band]31 return x.clamp(-cap, cap)3233# Initial standard deviation of the learned band embeddings.34BAND_EMB_INIT_STD: float = 1e-6353637class PerBandStem(nn.Module):38 """Subsample, clip, transpose, project, and add band identity.3940 For a two-second clip: (B, C, F_b, 64) -> (B, C, T_b, d),41 where F_b = 7/6/7 and T_b = 8/32/64 for slow/mid/fast.42 Projection weights are shared across contacts and times within each band.43 Returned positions use the original 32 Hz clock, not the subsampled index.44 """4546 def __init__(47 self,48 d_model: int = 256,49 *,50 bands: Sequence[tuple[int, int]] = PER_BAND_SPECS,51 band_emb_std: float = BAND_EMB_INIT_STD,52 ) -> None:53 super().__init__()54 self.specs = tuple((int(nb), int(st)) for nb, st in bands)55 self.projs = nn.ModuleList(nn.Linear(nb, d_model) for nb, _ in self.specs)56 # One learned width-d vector per band, broadcast over contacts and time.57 self.band_type_emb = nn.Parameter(torch.empty(len(self.specs), d_model))58 self.projs.apply(init_transformer_weights) # V-JEPA trunc_normal(0.02)+zero-bias59 nn.init.trunc_normal_(self.band_type_emb, std=band_emb_std)6061 @staticmethod62 def decimate(x: Tensor, stride: int) -> Tensor:63 """Strided time slice ``x[..., ::stride]`` — the band's own-rate frames. Requires64 the 32 Hz length to be an exact multiple of ``stride`` (64 % {1,2,8} == 0)."""65 t32 = x.shape[-1]66 if t32 % stride != 0:67 raise ValueError(f"32 Hz length {t32} not a multiple of stride {stride}")68 return x[..., ::stride]6970 def forward(71 self, band_inputs: Sequence[Tensor]72 ) -> tuple[tuple[Tensor, ...], tuple[Tensor, ...]]:73 """Bands ``[(...,F_b,T32)]`` → (per-band tokens ``(..., T_b, d)``, per-band lattice74 positions ``(T_b,)`` long). Order is SLOW, MID, HGA (the spec order)."""75 if len(band_inputs) != len(self.specs):76 raise ValueError(f"expected {len(self.specs)} bands, got {len(band_inputs)}")77 tokens: list[Tensor] = []78 positions: list[Tensor] = []79 for b, (x, (n_bins, stride), proj) in enumerate(80 zip(band_inputs, self.specs, self.projs)81 ):82 if x.shape[-2] != n_bins:83 raise ValueError(84 f"band {b} has {x.shape[-2]} freq bins, expected {n_bins}"85 )86 # 1. Subsample time and clip normalized inputs.87 xd = self.decimate(x, stride) # (..., F_b, T_b)88 xd = clip_band(xd, b) # same shape; cap extreme normalized values89 # 2. Put frequency last, then project F_b -> d and add band identity.90 xd = xd.transpose(-1, -2) # (..., T_b, F_b)91 tok = proj(xd) # (..., T_b, F_b) -> (..., T_b, d)92 tok = tok + self.band_type_emb[b] # add the same band vector to each token93 t_b = xd.shape[-2]94 # 3. Keep positions on the shared clock: slow 0,8,...; mid 0,2,...; fast 0,1,...95 pos = torch.arange(t_b, device=x.device, dtype=torch.long) * stride # (T_b,)96 tokens.append(tok)97 positions.append(pos)98 return tuple(tokens), tuple(positions)