← Walkthroughmapa/models/pe.py
Source snapshot · source SHA-256 9daefe51d2e6
1"""Rotary contact/time coordinates and learned anatomical region embeddings.23Main path: L1RoPE.cos_sin() builds angles; rotate() applies them to Q/K.45Input: Q/K features plus each token's contact/time coordinates.6Output: rotated Q/K with unchanged shapes; there are no learned RoPE parameters.7Next: return to attention.py to compute attention with these Q/K and unchanged V.89RoPE splits each head equally between clinical contact index and time on the1032 Hz clock. Within each half, adjacent feature pairs rotate at fixed frequencies.11Q and K are rotated; V is unchanged. For fixed Q/K content, a common coordinate12offset leaves their rotary dot product unchanged.1314Released bases: contact 8, time 64. Contact indices preserve gaps and are ordinal,15not physical distances. These settings do not imply a hard contact-count cutoff;16generalization to longer shafts or two-dimensional grids needs validation.1718RegionIdentityEmbed is a learned atlas-region lookup added to each contact's tokens.19"""2021from __future__ import annotations2223import torch24from torch import Tensor, nn252627def init_transformer_weights(m: nn.Module, init_std: float = 0.02) -> None:28 """V-JEPA 2 module init (``vision_transformer.py:130-141``): Linear weights29 ``trunc_normal_(std)`` + zero bias; LayerNorm weight 1 / bias 0. Applied via30 ``module.apply(...)``; Embeddings are skipped (they self-init in ``__init__``)."""31 if isinstance(m, nn.Linear):32 nn.init.trunc_normal_(m.weight, std=init_std)33 if m.bias is not None:34 nn.init.constant_(m.bias, 0.0)35 elif isinstance(m, nn.LayerNorm):36 nn.init.constant_(m.bias, 0.0)37 nn.init.constant_(m.weight, 1.0)383940def _rotate_half(x: Tensor) -> Tensor:41 # Each adjacent feature pair (a,b) becomes (-b,a), a 90-degree rotation.42 x_paired = x.unflatten(-1, (-1, 2)) # (..., hd/2, 2)43 return torch.stack([-x_paired[..., 1], x_paired[..., 0]], dim=-1).flatten(-2)444546class L1RoPE(nn.Module):47 """Two-axis axial RoPE: half of each head for contact, half for time."""4849 def __init__(50 self,51 head_dim: int,52 *,53 base_index: float = 8.0,54 base_time: float = 64.0,55 space: bool = True,56 ) -> None:57 super().__init__()58 if head_dim % 4 != 0:59 raise ValueError(60 f"L1RoPE needs head_dim % 4 == 0 (equal even split across 2 axes), "61 f"got {head_dim}"62 )63 self.head_dim = head_dim64 self.space = bool(space)65 pairs = head_dim // 4 # head_dim=64: 16 pairs for contact + 16 for time66 idx_freq = 1.0 / (base_index ** (torch.arange(pairs).float() / pairs))67 # Disabling space makes contact rotations identity; time rotations remain.68 if not self.space:69 idx_freq = torch.zeros_like(idx_freq)70 t_freq = 1.0 / (base_time ** (torch.arange(pairs).float() / pairs))71 self.register_buffer("idx_freq", idx_freq, persistent=False)72 self.register_buffer("t_freq", t_freq, persistent=False)7374 def cos_sin(self, idx: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:75 """The (..., seq, head_dim) fp32 cos/sin rotary table for coords (idx, t).7677 Every block holds the same schedule, so the table is the same for all of them and78 for every head. It is built once per forward and passed to :meth:`rotate` in each79 block, which is the same arithmetic as rebuilding it and a good deal less of it.80 The table stays fp32 and the result is cast down afterwards, because rounding it81 first would change the numbers."""82 # angle = coordinate * frequency (radians); frequencies are fixed buffers.83 ang_idx = idx[..., None].float() * self.idx_freq # (..., seq, pairs)84 ang_t = t[..., None].float() * self.t_freq85 ang = torch.cat([ang_idx, ang_t], dim=-1) # (..., seq, head_dim/2)86 cos = ang.cos().repeat_interleave(2, dim=-1) # (..., seq, head_dim)87 sin = ang.sin().repeat_interleave(2, dim=-1)88 return cos, sin8990 @staticmethod91 def rotate(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]:92 """Apply a precomputed (from :meth:`cos_sin`) rotary table to q, k.9394 q, k: (..., H, seq, head_dim); cos, sin: (..., seq, head_dim) — the head95 axis is broadcast in via ``unsqueeze(-3)``."""96 cos = cos.unsqueeze(-3) # (..., 1, seq, head_dim) → broadcast over heads97 sin = sin.unsqueeze(-3)98 q_out = q * cos + _rotate_half(q) * sin99 k_out = k * cos + _rotate_half(k) * sin100 return q_out, k_out101102 def forward(103 self, q: Tensor, k: Tensor, idx: Tensor, t: Tensor104 ) -> tuple[Tensor, Tensor]:105 """q, k: (..., H, seq, head_dim); idx, t: (..., seq) (no head axis)."""106 cos, sin = self.cos_sin(idx, t) # (..., seq, head_dim)107 return self.rotate(q, k, cos, sin)108109110class RegionIdentityEmbed(nn.Module):111 """Learned per-region identity embedding indexed by the DKT/DK hard tag."""112113 def __init__(114 self, n_regions: int, d_model: int, *, init_std: float = 1e-6, enabled: bool = True115 ) -> None:116 super().__init__()117 self.d_model = d_model118 self.embed: nn.Embedding | None = None119 # The disabled variant has no embedding parameters and returns zeros.120 if not enabled:121 return122 embed = nn.Embedding(n_regions, d_model)123 nn.init.trunc_normal_(embed.weight, std=init_std)124 self.embed = embed125126 def forward(self, region_id: Tensor) -> Tensor:127 """region_id: (..., seq) long → (..., seq, d_model)."""128 if self.embed is None:129 return region_id.new_zeros((*region_id.shape, self.d_model), dtype=torch.float32)130 return self.embed(region_id)