| """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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| @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), |
| ) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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).""" |
| |
| 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) |
| 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) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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: |
| |
| seqlen = hidden_states.shape[1] |
| BCx = self.in_proj(hidden_states).transpose(-1, -2) |
| B_gate, C_gate, x = BCx.chunk(3, dim=-2) |
| |
| 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 |
|
|
| |
| 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) |
| k = self.k_layernorm(k).transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| cos, sin = position_embeddings |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
|
|
| k = repeat_kv(k, self.num_kv_groups) |
| 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)) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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) |
| 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)) |
|
|
