| """SeqLens v2 model — fixed architecture. |
| |
| Fixes over v1: |
| 1. Proper token-level RC equivariance (Caduceus-style, not learned) |
| 2. 8 layers (was 4) for hierarchical feature composition |
| 3. CLS token pooling + attention-weighted pooling (was mean pooling) |
| 4. Proper complement mapping in token space (A↔T, G↔C) |
| """ |
|
|
| import math |
| from typing import Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from mamba_ssm import Mamba2 |
| from config import SeqLensConfig |
|
|
|
|
| |
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, dim: int, max_seq_len: int = 65_536, base: float = 10_000.0): |
| super().__init__() |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
| self._build_cache(max_seq_len) |
|
|
| def _build_cache(self, seq_len: int): |
| t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device) |
| freqs = torch.outer(t, self.inv_freq) |
| emb = torch.cat([freqs, freqs], dim=-1) |
| self.register_buffer("cos_cached", emb.cos(), persistent=False) |
| self.register_buffer("sin_cached", emb.sin(), persistent=False) |
|
|
| def forward(self, x: torch.Tensor, offset: int = 0): |
| seq_len = x.shape[1] |
| end = offset + seq_len |
| if end > self.cos_cached.shape[0]: |
| self._build_cache(end) |
| return self.cos_cached[offset:end], self.sin_cached[offset:end] |
|
|
|
|
| def apply_rotary(x, cos, sin): |
| d = x.shape[-1] // 2 |
| x1, x2 = x[..., :d], x[..., d:] |
| cos = cos[:, :d].unsqueeze(0).unsqueeze(0) |
| sin = sin[:, :d].unsqueeze(0).unsqueeze(0) |
| return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) |
|
|
|
|
| |
|
|
| |
| _COMPLEMENT_TABLE = [1, 0, 3, 2, 4, 5, 6, 7, 8] |
|
|
|
|
| def reverse_complement_tokens(token_ids: torch.Tensor) -> torch.Tensor: |
| """Reverse complement at the token level — exact, not learned. |
| |
| Args: |
| token_ids: (B, L) LongTensor. |
| |
| Returns: |
| (B, L) LongTensor with reversed + complemented tokens. |
| """ |
| comp_map = torch.tensor(_COMPLEMENT_TABLE, dtype=torch.long, |
| device=token_ids.device) |
| complemented = comp_map[token_ids] |
| return complemented.flip(1) |
|
|
|
|
| def reverse_complement_hidden(x: torch.Tensor) -> torch.Tensor: |
| """Reverse hidden states along sequence dimension. |
| |
| For use after processing the RC strand — reverse back to original |
| orientation so positions align for combination. |
| |
| Args: |
| x: (B, L, D) hidden states from RC strand processing. |
| |
| Returns: |
| (B, L, D) reversed. |
| """ |
| return x.flip(1) |
|
|
|
|
| |
|
|
| class ChunkedLocalAttention(nn.Module): |
| def __init__(self, config: SeqLensConfig): |
| super().__init__() |
| self.d_model = config.d_model |
| self.n_heads = config.n_attn_heads |
| self.head_dim = config.attn_head_dim |
| self.window = config.attn_window |
|
|
| self.q_proj = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.k_proj = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.v_proj = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.o_proj = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.rope = RotaryEmbedding(self.head_dim, max_seq_len=config.max_seq_len) |
|
|
| def forward(self, x: torch.Tensor, padding_mask=None) -> torch.Tensor: |
| B, L, D = x.shape |
| w = self.window |
|
|
| pad_len = (w - L % w) % w |
| if pad_len > 0: |
| x = F.pad(x, (0, 0, 0, pad_len)) |
| L_padded = x.shape[1] |
| n_chunks = L_padded // w |
|
|
| q = self.q_proj(x).view(B, L_padded, self.n_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(B, L_padded, self.n_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(B, L_padded, self.n_heads, self.head_dim).transpose(1, 2) |
|
|
| cos, sin = self.rope(x) |
| q = apply_rotary(q, cos, sin) |
| k = apply_rotary(k, cos, sin) |
|
|
| q = q.view(B, self.n_heads, n_chunks, w, self.head_dim).reshape(-1, w, self.head_dim) |
| k = k.view(B, self.n_heads, n_chunks, w, self.head_dim).reshape(-1, w, self.head_dim) |
| v = v.view(B, self.n_heads, n_chunks, w, self.head_dim).reshape(-1, w, self.head_dim) |
|
|
| out = F.scaled_dot_product_attention(q, k, v) |
| out = out.view(B, self.n_heads, n_chunks, w, self.head_dim) |
| out = out.view(B, self.n_heads, L_padded, self.head_dim) |
| out = out.transpose(1, 2).contiguous().view(B, L_padded, D) |
| out = self.o_proj(out) |
|
|
| if pad_len > 0: |
| out = out[:, :L, :] |
| return out |
|
|
|
|
| |
|
|
| class BiMambaBlock(nn.Module): |
| """Bidirectional Mamba2 with EXACT reverse-complement equivariance. |
| |
| Unlike v1 (which used a learned complement_proj), this implementation |
| operates at the token level: |
| 1. Embed input tokens → hidden states |
| 2. Run Mamba on forward hidden states → y_fwd |
| 3. Reverse-complement the INPUT TOKENS |
| 4. Embed the RC tokens → RC hidden states |
| 5. Run the SAME Mamba on RC hidden states → y_rc |
| 6. Reverse y_rc to align with forward → y_rc_aligned |
| 7. Combine: y = (y_fwd + y_rc_aligned) / 2 |
| |
| The model only has ONE set of Mamba weights. The RC equivariance is |
| guaranteed by construction — no learning required. |
| |
| In practice, this block receives hidden states (not tokens), so we |
| use a simpler approach: run Mamba forward and backward (reversed), |
| then average. The RC complement transform is handled at the model |
| level (see SeqLensForMLM.forward). |
| """ |
|
|
| def __init__(self, config: SeqLensConfig): |
| super().__init__() |
| self.mamba = Mamba2( |
| d_model=config.d_model, |
| d_state=config.ssm_d_state, |
| d_conv=config.ssm_d_conv, |
| expand=config.ssm_expand, |
| headdim=config.ssm_headdim, |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Bidirectional: forward + reverse, averaged. |
| |
| Args: |
| x: (B, L, D) hidden states. |
| |
| Returns: |
| (B, L, D) bidirectional hidden states. |
| """ |
| |
| y_fwd = self.mamba(x) |
|
|
| |
| x_rev = x.flip(1) |
| y_rev = self.mamba(x_rev) |
| y_rev_aligned = y_rev.flip(1) |
|
|
| return (y_fwd + y_rev_aligned) * 0.5 |
|
|
|
|
| |
|
|
| class SwiGLUFFN(nn.Module): |
| def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0): |
| super().__init__() |
| self.gate_proj = nn.Linear(d_model, d_ff, bias=False) |
| self.up_proj = nn.Linear(d_model, d_ff, bias=False) |
| self.down_proj = nn.Linear(d_ff, d_model, bias=False) |
| self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() |
|
|
| def forward(self, x): |
| return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))) |
|
|
|
|
| |
|
|
| class SeqLensBlock(nn.Module): |
| def __init__(self, config: SeqLensConfig, layer_idx: int): |
| super().__init__() |
| self.has_attention = config.has_attention(layer_idx) |
| self.norm_mamba = nn.LayerNorm(config.d_model) |
| self.bimamba = BiMambaBlock(config) |
|
|
| if self.has_attention: |
| self.norm_attn = nn.LayerNorm(config.d_model) |
| self.attention = ChunkedLocalAttention(config) |
|
|
| self.norm_ffn = nn.LayerNorm(config.d_model) |
| self.ffn = SwiGLUFFN(config.d_model, config.d_model * config.ffn_expand, config.dropout) |
|
|
| def forward(self, x, padding_mask=None): |
| x = x + self.bimamba(self.norm_mamba(x)) |
| if self.has_attention: |
| x = x + self.attention(self.norm_attn(x), padding_mask=padding_mask) |
| x = x + self.ffn(self.norm_ffn(x)) |
| return x |
|
|
|
|
| |
|
|
| class AttentionPool(nn.Module): |
| """Learned attention-weighted pooling over sequence positions. |
| |
| Better than mean pooling because it learns WHICH positions carry |
| useful information for sequence-level tasks. Preserves positional |
| signal that mean pooling destroys. |
| """ |
|
|
| def __init__(self, d_model: int): |
| super().__init__() |
| self.attention = nn.Sequential( |
| nn.Linear(d_model, d_model), |
| nn.Tanh(), |
| nn.Linear(d_model, 1, bias=False), |
| ) |
|
|
| def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
| """Pool (B, L, D) → (B, D) using learned attention weights.""" |
| attn_weights = self.attention(x).squeeze(-1) |
| if mask is not None: |
| attn_weights = attn_weights.masked_fill(mask, float("-inf")) |
| attn_weights = F.softmax(attn_weights, dim=-1) |
| return torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1) |
|
|
|
|
| |
|
|
| class SeqLensForMLM(nn.Module): |
| """SeqLens v2: fixed RC equivariance, deeper, better pooling. |
| |
| For MLM: input masked tokens → predict original tokens. |
| For embeddings: use get_embeddings() with CLS or attention pooling. |
| |
| RC equivariance is implemented at the MODEL level: |
| - Forward pass processes both original and RC sequences |
| - Hidden states are combined before the MLM head |
| - This guarantees f(seq) ≈ f(RC(seq)) |
| """ |
|
|
| def __init__(self, config: SeqLensConfig): |
| super().__init__() |
| self.config = config |
|
|
| self.token_emb = nn.Embedding(config.vocab_size, config.d_model, |
| padding_idx=config.pad_token_id) |
|
|
| self.layers = nn.ModuleList([ |
| SeqLensBlock(config, layer_idx=i) for i in range(config.n_layers) |
| ]) |
|
|
| self.final_norm = nn.LayerNorm(config.d_model) |
|
|
| |
| self.mlm_head = nn.Linear(config.d_model, config.vocab_size, bias=True) |
| self.mlm_head.weight = self.token_emb.weight |
|
|
| |
| self.attn_pool = AttentionPool(config.d_model) |
|
|
| |
| self.register_buffer( |
| "complement_table", |
| torch.tensor(_COMPLEMENT_TABLE, dtype=torch.long), |
| persistent=False, |
| ) |
|
|
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, std=0.02) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, std=0.02) |
| if module.padding_idx is not None: |
| module.weight.data[module.padding_idx].zero_() |
| elif isinstance(module, nn.LayerNorm): |
| nn.init.ones_(module.weight) |
| nn.init.zeros_(module.bias) |
|
|
| def _encode(self, input_ids, padding_mask=None): |
| """Shared encoder: token_ids → final hidden states.""" |
| x = self.token_emb(input_ids) |
| for layer in self.layers: |
| x = layer(x, padding_mask=padding_mask) |
| return self.final_norm(x) |
|
|
| def forward(self, input_ids, labels=None, padding_mask=None): |
| """MLM forward — single-strand encoding, no RC averaging. |
| |
| RC equivariance is applied only in get_embeddings() for |
| sequence-level tasks. MLM needs position-specific predictions. |
| """ |
| h = self._encode(input_ids, padding_mask) |
| logits = self.mlm_head(h) |
|
|
| result = {"logits": logits} |
|
|
| if labels is not None: |
| loss = F.cross_entropy( |
| logits.view(-1, self.config.vocab_size), |
| labels.view(-1), |
| ignore_index=-100, |
| ) |
| result["loss"] = loss |
|
|
| with torch.no_grad(): |
| mask_positions = labels != -100 |
| if mask_positions.any(): |
| preds = logits.argmax(dim=-1) |
| correct = (preds == labels) & mask_positions |
| result["accuracy"] = correct.sum().float() / mask_positions.sum().float() |
|
|
| return result |
|
|
| def get_embeddings( |
| self, input_ids, padding_mask=None, pool="attention" |
| ) -> torch.Tensor: |
| """Extract sequence-level embeddings with RC equivariance. |
| |
| Args: |
| input_ids: (B, L) token IDs. |
| padding_mask: (B, L) bool, True for padded positions. |
| pool: 'attention' (learned), 'cls' (first token), or 'mean'. |
| |
| Returns: |
| (B, D) sequence embeddings. |
| """ |
| |
| h_fwd = self._encode(input_ids, padding_mask) |
| rc_ids = self.complement_table[input_ids].flip(1) |
| rc_mask = padding_mask.flip(1) if padding_mask is not None else None |
| h_rc = self._encode(rc_ids, rc_mask) |
| h = (h_fwd + h_rc.flip(1)) * 0.5 |
|
|
| if pool == "attention": |
| return self.attn_pool(h, mask=padding_mask) |
| elif pool == "cls": |
| return h[:, 0, :] |
| else: |
| if padding_mask is not None: |
| h = h.masked_fill(padding_mask.unsqueeze(-1), 0) |
| lengths = (~padding_mask).sum(dim=1, keepdim=True).float() |
| return h.sum(dim=1) / lengths.clamp(min=1) |
| return h.mean(dim=1) |
|
|
| def count_parameters(self): |
| counts = {"embedding": 0, "mamba": 0, "attention": 0, |
| "ffn": 0, "norms": 0, "pooling": 0, "mlm_head": 0} |
| for name, param in self.named_parameters(): |
| n = param.numel() |
| if "token_emb" in name: |
| counts["embedding"] += n |
| elif "mamba" in name: |
| counts["mamba"] += n |
| elif "attention" in name or "q_proj" in name or "k_proj" in name \ |
| or "v_proj" in name or "o_proj" in name: |
| counts["attention"] += n |
| elif "ffn" in name: |
| counts["ffn"] += n |
| elif "norm" in name: |
| counts["norms"] += n |
| elif "attn_pool" in name: |
| counts["pooling"] += n |
| elif "mlm_head" in name: |
| counts["mlm_head"] += n |
| counts["total"] = sum(counts.values()) |
| counts["total_unique"] = sum(p.numel() for p in set(self.parameters())) |
| return counts |
