"""LFM2Small: scaled-down LFM2.5-1.2B backbone for transaction sequences. Reimplements (not subclasses) the core LFM2 architecture at ~8.3M total params. 8 layers in conv-conv-attn-conv-attn-conv-attn-conv order, preserving every structural choice from the full 1.2B model: - Gated short convolution with depthwise causal Conv1d - Grouped query attention (4Q / 2KV, group size 2) with QK RMSNorm - SwiGLU MLP with auto-adjusted intermediate dimension - Pre-norm residual connections - Final RMSNorm (embedding_norm) before LM heads Module naming matches LFM2 conventions exactly: layers[i].self_attn.{q_proj, k_proj, v_proj, out_proj, q_layernorm, k_layernorm} layers[i].conv.{in_proj, out_proj, conv} layers[i].feed_forward.{w1, w2, w3} layers[i].{operator_norm, ffn_norm} embedding_norm Reference: modeling_lfm2.py in HuggingFace transformers (LiquidAI/LFM2-1.2B). """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F import yaml from src.data.schema import SchemaConfig, load_schema from src.model.embedding import StructuredEmbedding from src.model.heads import PerFeatureLMHeads # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @dataclass class ModelConfig: """Typed config for LFM2Small. Loads from configs/model.yaml.""" hidden_size: int = 256 intermediate_size: int = 1024 num_attention_heads: int = 4 num_key_value_heads: int = 2 num_layers: int = 8 layer_order: list[str] = field(default_factory=lambda: [ "conv", "conv", "attn", "conv", "attn", "conv", "attn", "conv", ]) conv_kernel_size: int = 3 block_auto_adjust_ff_dim: bool = True block_multiple_of: int = 256 block_ffn_dim_multiplier: float = 1.0 rms_norm_eps: float = 1e-6 rope_theta: float = 1_000_000.0 max_position_embeddings: int = 4096 initializer_range: float = 0.02 num_transactions: int = 64 num_features: int = 15 @property def head_dim(self) -> int: return self.hidden_size // self.num_attention_heads @property def num_kv_groups(self) -> int: return self.num_attention_heads // self.num_key_value_heads @property def effective_intermediate_size(self) -> int: """MLP dim after LFM2's block_auto_adjust_ff_dim. With hidden=256, intermediate=1024: int(2*1024/3)=682, rounded to 768. """ if not self.block_auto_adjust_ff_dim: return self.intermediate_size size = int(2 * self.intermediate_size / 3) size = int(self.block_ffn_dim_multiplier * size) return self.block_multiple_of * ( (size + self.block_multiple_of - 1) // self.block_multiple_of ) @classmethod def from_yaml(cls, path: str | Path) -> ModelConfig: with open(path) as f: raw = yaml.safe_load(f) bb = raw.get("backbone", {}) seq = raw.get("sequence", {}) return cls( hidden_size=bb.get("hidden_size", 256), intermediate_size=bb.get("intermediate_size", 1024), num_attention_heads=bb.get("num_attention_heads", 4), num_key_value_heads=bb.get("num_key_value_heads", 2), num_layers=bb.get("num_layers", 8), layer_order=bb.get("layer_order", [ "conv", "conv", "attn", "conv", "attn", "conv", "attn", "conv", ]), conv_kernel_size=bb.get("conv_kernel_size", 3), block_auto_adjust_ff_dim=bb.get("block_auto_adjust_ff_dim", True), block_multiple_of=bb.get("block_multiple_of", 256), block_ffn_dim_multiplier=bb.get("block_ffn_dim_multiplier", 1.0), rms_norm_eps=bb.get("rms_norm_eps", 1e-6), rope_theta=bb.get("rope_theta", 1_000_000.0), num_transactions=seq.get("num_transactions", 64), num_features=seq.get("features_per_transaction", 15), ) # --------------------------------------------------------------------------- # Building blocks # --------------------------------------------------------------------------- class RMSNorm(nn.Module): """Root mean square layer normalization (Lfm2RMSNorm).""" def __init__(self, dim: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return (self.weight * x).to(dtype) class RotaryEmbedding(nn.Module): """Rotary position embeddings. Flat token positions 0..S-1.""" def __init__( self, head_dim: int, max_seq_len: int = 4096, theta: float = 1_000_000.0, ) -> None: super().__init__() inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.max_seq_len = max_seq_len def forward( self, x: torch.Tensor, position_ids: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Returns (cos, sin) each shaped (B, S, head_dim).""" # (1, D/2, 1) @ (B, 1, S) -> (B, D/2, S) -> (B, S, D/2) inv_freq = self.inv_freq[None, :, None].float().to(x.device) pos = position_ids[:, None, :].float() freqs = (inv_freq @ pos).transpose(1, 2) emb = torch.cat([freqs, freqs], dim=-1) # (B, S, head_dim) return emb.cos().to(x.dtype), emb.sin().to(x.dtype) def rotate_half(x: torch.Tensor) -> torch.Tensor: x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """cos/sin: (B, S, D) unsqueezed to (B, 1, S, D). q/k: (B, H, S, D).""" cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: """Expand KV heads for GQA: (B, H_kv, S, D) -> (B, H_kv*n_rep, S, D).""" if n_rep == 1: return x B, H, S, D = x.shape return x[:, :, None, :, :].expand(B, H, n_rep, S, D).reshape(B, H * n_rep, S, D) # --------------------------------------------------------------------------- # Layers # --------------------------------------------------------------------------- class SwiGLU(nn.Module): """SwiGLU MLP (Lfm2MLP): w2(silu(w1(x)) * w3(x)).""" def __init__(self, config: ModelConfig) -> None: super().__init__() intermediate = config.effective_intermediate_size self.w1 = nn.Linear(config.hidden_size, intermediate, bias=False) self.w3 = nn.Linear(config.hidden_size, intermediate, bias=False) self.w2 = nn.Linear(intermediate, config.hidden_size, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w2(F.silu(self.w1(x)) * self.w3(x)) class ShortConv(nn.Module): """Gated short convolution (Lfm2ShortConv). in_proj splits hidden -> (B_gate, C_gate, x). B_gate * x feeds a causal depthwise Conv1d, output gated by C_gate, then out_proj. Left-padding (padding=kernel-1, truncated to seqlen) ensures no future token leakage. """ def __init__(self, config: ModelConfig) -> None: super().__init__() h, k = config.hidden_size, config.conv_kernel_size self.in_proj = nn.Linear(h, 3 * h, bias=False) self.conv = nn.Conv1d(h, h, kernel_size=k, groups=h, bias=False, padding=k - 1) self.out_proj = nn.Linear(h, h, bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # (B, S, D) -> project -> (B, 3D, S) -> chunk into three (B, D, S) tensors seqlen = hidden_states.shape[1] BCx = self.in_proj(hidden_states).transpose(-1, -2) B_gate, C_gate, x = BCx.chunk(3, dim=-2) # Causal depthwise conv: left-padded, truncate right to preserve causality conv_out = self.conv(B_gate * x)[..., :seqlen] y = C_gate * conv_out return self.out_proj(y.transpose(-1, -2).contiguous()) class Attention(nn.Module): """Grouped query attention with QK RMSNorm (Lfm2Attention). QK norms after projection and before rotary stabilize deep training. Present in LFM2 but absent from LLaMA-family models. """ def __init__(self, config: ModelConfig) -> None: super().__init__() self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.num_kv_groups = config.num_kv_groups self.head_dim = config.head_dim self.scaling = self.head_dim ** -0.5 self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) self.q_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: B, S, _ = hidden_states.shape # Project -> reshape to heads -> QK norm -> transpose to (B, H, S, D) q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim) k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim) v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim) q = self.q_layernorm(q).transpose(1, 2) # (B, H, S, D) k = self.k_layernorm(k).transpose(1, 2) # (B, H_kv, S, D) v = v.transpose(1, 2) # (B, H_kv, S, D) cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) k = repeat_kv(k, self.num_kv_groups) # (B, H, S, D) v = repeat_kv(v, self.num_kv_groups) attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=self.scaling) return self.out_proj(attn_out.transpose(1, 2).reshape(B, S, -1).contiguous()) class DecoderLayer(nn.Module): """Pre-norm residual: conv or attention + SwiGLU (Lfm2DecoderLayer). x = x + op(operator_norm(x)) # op = conv or self_attn x = x + feed_forward(ffn_norm(x)) """ def __init__(self, config: ModelConfig, layer_idx: int) -> None: super().__init__() self.is_attention_layer = config.layer_order[layer_idx] == "attn" if self.is_attention_layer: self.self_attn = Attention(config) else: self.conv = ShortConv(config) self.feed_forward = SwiGLU(config) self.operator_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: residual = hidden_states if self.is_attention_layer: hidden_states = self.self_attn(self.operator_norm(hidden_states), position_embeddings) else: hidden_states = self.conv(self.operator_norm(hidden_states)) hidden_states = hidden_states + residual return hidden_states + self.feed_forward(self.ffn_norm(hidden_states)) # --------------------------------------------------------------------------- # Full model # --------------------------------------------------------------------------- class LFM2Small(nn.Module): """LFM2-small: structured embedding + interleaved backbone + tied LM heads. ~8.3M params at hidden=256 (embedding ~1.7M, backbone ~6.6M, heads tied). """ def __init__(self, config: ModelConfig, schema: SchemaConfig) -> None: super().__init__() self.config = config assert len(config.layer_order) == config.num_layers assert config.hidden_size % config.num_attention_heads == 0 assert config.num_attention_heads % config.num_key_value_heads == 0 self.embedding = StructuredEmbedding(schema, config.hidden_size) self.layers = nn.ModuleList([ DecoderLayer(config, i) for i in range(config.num_layers) ]) self.rotary_emb = RotaryEmbedding( config.head_dim, config.max_position_embeddings, config.rope_theta, ) self.embedding_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.lm_heads = PerFeatureLMHeads(self.embedding) self._init_weights() def _init_weights(self) -> None: """Initialize weights following LFM2 conventions. Skips lm_heads (tied).""" for name, module in self.named_modules(): if name.startswith("lm_heads"): continue if isinstance(module, (nn.Linear, nn.Conv1d)): nn.init.normal_(module.weight, std=self.config.initializer_range) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, std=self.config.initializer_range) def backbone_forward(self, token_ids: torch.Tensor) -> torch.Tensor: """Embedding + backbone + final norm. Returns (B, S, D). Use for downstream heads (fraud prediction) that skip LM logits. """ hidden_states = self.embedding(token_ids) # (B, T*F, D) position_ids = torch.arange( hidden_states.shape[1], device=hidden_states.device, ).unsqueeze(0) position_embeddings = self.rotary_emb(hidden_states, position_ids) for layer in self.layers: hidden_states = layer(hidden_states, position_embeddings) return self.embedding_norm(hidden_states) def forward(self, token_ids: torch.Tensor) -> list[torch.Tensor]: """Token IDs -> per-feature logits for causal LM pretraining. Args: token_ids: (B, T, F) int tensor. Returns: 15 tensors, each (B, T*F, vocab_size_f). Position p predicts position p+1; the training loop selects head[(p+1) % num_features]. """ return self.lm_heads(self.backbone_forward(token_ids)) def param_count(self) -> dict[str, int]: """Parameter counts by component. Accounts for weight tying.""" emb = sum(p.numel() for p in self.embedding.parameters()) backbone = sum(p.numel() for p in self.layers.parameters()) backbone += sum(p.numel() for p in self.embedding_norm.parameters()) total = sum(p.numel() for p in self.parameters()) return {"embedding": emb, "backbone": backbone, "lm_heads_tied": 0, "total_unique": total} @classmethod def from_config_files(cls, model_yaml: str | Path, schema_yaml: str | Path) -> LFM2Small: """Construct from YAML config files.""" return cls(ModelConfig.from_yaml(model_yaml), load_schema(schema_yaml))