add multi-head demo as 4th-6th tabs; restore Why Liquid + Integration

README.md

@@ -1,5 +1,5 @@
 ---
-title: Transaction Co-Pilot — Dispute · Collections · Fraud
+title: Transaction Encoder — LFM2.5 multi-head + Co-Pilot surfaces
 emoji: 🪙
 colorFrom: gray
 colorTo: indigo
@@ -9,7 +9,7 @@ app_file: app.py
 pinned: false
 ---
 
-# Transaction Co-Pilot — One Backbone, Three Surfaces
+# Transaction Encoder — One Backbone, Many Surfaces
 
 Reference implementation of the multimodal-encoder recipe applied to
 discrete-feature transaction sequences, demonstrating cross-surface
@@ -22,8 +22,14 @@ per-surface LoRA adapts the attention layers.
 
 ## What this demo shows
 
-Three surfaces, three independent predictions on the same backbone:
+Six tabs over one shared LFM2.5-350M backbone:
 
+- **Multi-Head Demo** — the original 4-task-head encoder demo. Live
+  inference on curated customer archetypes: fraud probability,
+  next-merchant prediction, amount-bucket prediction, MCC prediction.
+- **Why Liquid** — architectural pitch for the encoder-on-frozen-backbone
+  recipe.
+- **Integration** — build-it-yourself guide.
 - **Dispute Co-Pilot** — given a customer's transaction history + a
   dispute complaint, classifies the dispute as likely / ambiguous /
   unlikely friendly fraud, with per-transaction attribution.
@@ -34,9 +40,10 @@ Three surfaces, three independent predictions on the same backbone:
   classifies the attack stage (5-way) AND the underlying type (4-way)
   in two independent categorical outputs.
 
-Each surface ships its own slim checkpoint (~24 MB), but all three
-share the LFM2.5-350M base weights — the demonstration that the
-backbone is the product, and per-surface plumbing rides on top.
+Each Co-Pilot surface ships its own slim checkpoint (~24 MB) alongside
+the original multi-head checkpoint (~24 MB), but all four share the
+LFM2.5-350M base weights — the demonstration that the backbone is the
+product, and per-surface plumbing rides on top.
 
 ## How it works
 
@@ -70,6 +77,7 @@ cited as performance on real payment data.
 
 ## Cast performance (held-out)
 
+- Multi-Head V3 (step 4999): per-head accuracy reported on the test split.
 - Dispute v7 (step 2000): 5/6 cast band-match, 6/6 attribution top-5
 - Collections v3 (step 4500): 5/6 cast band-match
 - Fraud v1 (step 3000): 6/6 stage AND type at 1.00 confidence

app.py

@@ -1,8 +1,16 @@
-"""Hugging Face Spaces entrypoint for the unified Transaction Co-Pilot.
+"""Hugging Face Spaces entrypoint for the unified Transaction Encoder demo.
 
 Wrapper that invokes encoder.src.demo.copilot_app_unified.main with
 paths relative to the Space repo root. The full app code lives at
 encoder/src/demo/copilot_app_unified.py.
+
+Six tabs:
+  - Multi-Head Demo       (original V3 encoder + 4 task heads)
+  - Why Liquid            (architectural pitch)
+  - Integration           (build-it-yourself guide)
+  - Dispute Co-Pilot      (friendly-fraud classifier + attribution)
+  - Collections Co-Pilot  (treatment-response scoreboard)
+  - Fraud Co-Pilot        (pattern stage × type classifier)
 """
 
 import sys
@@ -15,18 +23,25 @@ from encoder.src.demo.copilot_app_unified import main
 
 sys.argv = [
     "app",
+    # Multi-head (original) demo
+    "--multihead-checkpoint", "checkpoints/multihead_v3.pt",
+    "--multihead-config", "encoder/configs/model_nocompress.yaml",
+    "--multihead-data-dir", "data/synthetic",
+    # Co-Pilot surfaces
     "--dispute-checkpoint", "checkpoints/dispute_legitimacy_v7.pt",
     "--dispute-config", "encoder/configs/model_dispute_legitimacy.yaml",
+    "--dispute-cast", "encoder/data/demo_cast.json",
     "--collections-checkpoint", "checkpoints/collections_v3.pt",
     "--collections-config", "encoder/configs/model_collections.yaml",
+    "--collections-cast", "encoder/data/collections_cast.json",
     "--fraud-checkpoint", "checkpoints/fraud_pattern_v1.pt",
     "--fraud-config", "encoder/configs/model_fraud_pattern.yaml",
-    "--schema", "data/schema.yaml",
-    "--histories", "data/synthetic/token_ids.npy",
-    "--dispute-cast", "encoder/data/demo_cast.json",
-    "--collections-cast", "encoder/data/collections_cast.json",
     "--fraud-cast", "encoder/data/fraud_pattern_cast.json",
+    # Shared
+    "--schema", "data/schema.yaml",
+    "--cast-histories", "data/synthetic/cast_token_ids.npy",
     "--device", "cpu",
+    "--dtype", "float32",
     "--port", "7860",
 ]
 

checkpoints/multihead_v3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0385a423ae688ae0b0d6f672bb4230f47d1808462fecb29bd619985070007dfe
+size 188887205

data/synthetic/cast_token_ids.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96fa931533be3b36bce7b6c7db369a9bdd6d39f06373cead7138efa0bd186f57
+size 34688

data/synthetic/sequence_labels.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25ac01df0191c2090f1fec782b46a93b3560dd2d42215c929bc18f44519cc443
+size 20128

data/synthetic/split_indices.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9402025636780e3fa0a2ca086a2756fe2f7bb67befcae10e2eca4ffcede91987
-size 886
+oid sha256:fab90d8398882dfac2d9660de977b111732deb31ad925f01d29073324b71fa96
+size 160742

data/synthetic/token_ids.npy

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:96fa931533be3b36bce7b6c7db369a9bdd6d39f06373cead7138efa0bd186f57
-size 34688
+oid sha256:a1c06df0488e17ebec73db4584b5fc49bcb6e02bfc2aba0f49a561443d036cf0
+size 38400128

data/synthetic/tokenizer_state.json

@@ -0,0 +1,424 @@
+{
+  "features": [
+    {
+      "name": "hour",
+      "num_values": 24,
+      "type": "categorical",
+      "vocab_size": 27
+    },
+    {
+      "name": "dow",
+      "num_values": 7,
+      "type": "categorical",
+      "vocab_size": 10
+    },
+    {
+      "boundaries": [
+        0.0,
+        12.166666666666666,
+        24.333333333333332,
+        36.5,
+        48.666666666666664,
+        60.83333333333333,
+        73.0,
+        85.16666666666666,
+        97.33333333333333,
+        109.5,
+        121.66666666666666,
+        133.83333333333331,
+        146.0,
+        158.16666666666666,
+        170.33333333333331,
+        182.5,
+        194.66666666666666,
+        206.83333333333331,
+        219.0,
+        231.16666666666666,
+        243.33333333333331,
+        255.5,
+        267.66666666666663,
+        279.8333333333333,
+        292.0,
+        304.16666666666663,
+        316.3333333333333,
+        328.5,
+        340.66666666666663,
+        352.8333333333333,
+        365.0
+      ],
+      "name": "days_since_last",
+      "num_values": 30,
+      "type": "bucketed",
+      "vocab_size": 33
+    },
+    {
+      "name": "is_recurring",
+      "num_values": 2,
+      "type": "binary",
+      "vocab_size": 5
+    },
+    {
+      "name": "mcc",
+      "num_values": 100,
+      "type": "categorical",
+      "vocab_size": 103
+    },
+    {
+      "name": "merchant_id",
+      "num_values": 10000,
+      "type": "categorical",
+      "vocab_size": 10003
+    },
+    {
+      "boundaries": [
+        0.0,
+        25.0,
+        50.0,
+        75.0,
+        100.0,
+        125.0,
+        150.0,
+        175.0,
+        200.0,
+        225.0,
+        250.0,
+        275.0,
+        300.0,
+        325.0,
+        350.0,
+        375.0,
+        400.0,
+        425.0,
+        450.0,
+        475.0,
+        500.0
+      ],
+      "name": "customer_merchant_count",
+      "num_values": 20,
+      "type": "bucketed",
+      "vocab_size": 23
+    },
+    {
+      "name": "entry_mode",
+      "num_values": 5,
+      "type": "categorical",
+      "vocab_size": 8
+    },
+    {
+      "boundaries": [
+        0.01,
+        97.66621093750001,
+        195.322421875,
+        292.9786328125,
+        390.63484375,
+        488.29105468750004,
+        585.947265625,
+        683.6034765625001,
+        781.2596875,
+        878.9158984375,
+        976.5721093750001,
+        1074.2283203125,
+        1171.88453125,
+        1269.5407421875,
+        1367.1969531250002,
+        1464.8531640625001,
+        1562.509375,
+        1660.1655859375,
+        1757.821796875,
+        1855.4780078125002,
+        1953.1342187500002,
+        2050.7904296875004,
+        2148.4466406250003,
+        2246.1028515625003,
+        2343.7590625000003,
+        2441.4152734375,
+        2539.071484375,
+        2636.7276953125006,
+        2734.3839062500006,
+        2832.0401171875005,
+        2929.6963281250005,
+        3027.3525390625005,
+        3125.0087500000004,
+        3222.6649609375004,
+        3320.3211718750003,
+        3417.9773828125003,
+        3515.6335937500003,
+        3613.2898046875002,
+        3710.9460156250007,
+        3808.6022265625006,
+        3906.2584375000006,
+        4003.9146484375005,
+        4101.5708593750005,
+        4199.2270703125005,
+        4296.88328125,
+        4394.5394921875,
+        4492.195703125,
+        4589.8519140625,
+        4687.508125,
+        4785.1643359375,
+        4882.820546875,
+        4980.4767578125,
+        5078.13296875,
+        5175.7891796875,
+        5273.445390625001,
+        5371.101601562501,
+        5468.757812500001,
+        5566.414023437501,
+        5664.070234375001,
+        5761.726445312501,
+        5859.382656250001,
+        5957.038867187501,
+        6054.695078125001,
+        6152.351289062501,
+        6250.007500000001,
+        6347.663710937501,
+        6445.319921875001,
+        6542.9761328125005,
+        6640.6323437500005,
+        6738.2885546875,
+        6835.944765625,
+        6933.6009765625,
+        7031.2571875,
+        7128.9133984375,
+        7226.569609375,
+        7324.225820312501,
+        7421.882031250001,
+        7519.538242187501,
+        7617.194453125001,
+        7714.850664062501,
+        7812.506875000001,
+        7910.163085937501,
+        8007.819296875001,
+        8105.475507812501,
+        8203.13171875,
+        8300.7879296875,
+        8398.444140625,
+        8496.100351562502,
+        8593.7565625,
+        8691.412773437502,
+        8789.068984375,
+        8886.725195312501,
+        8984.38140625,
+        9082.037617187501,
+        9179.693828125,
+        9277.350039062501,
+        9375.00625,
+        9472.662460937501,
+        9570.318671875,
+        9667.974882812501,
+        9765.63109375,
+        9863.287304687501,
+        9960.943515625,
+        10058.599726562501,
+        10156.2559375,
+        10253.912148437501,
+        10351.568359375,
+        10449.2245703125,
+        10546.880781250002,
+        10644.5369921875,
+        10742.193203125002,
+        10839.8494140625,
+        10937.505625000002,
+        11035.1618359375,
+        11132.818046875002,
+        11230.4742578125,
+        11328.130468750001,
+        11425.7866796875,
+        11523.442890625001,
+        11621.0991015625,
+        11718.755312500001,
+        11816.4115234375,
+        11914.067734375001,
+        12011.7239453125,
+        12109.380156250001,
+        12207.0363671875,
+        12304.692578125001,
+        12402.3487890625,
+        12500.005000000001,
+        12597.661210937502,
+        12695.317421875001,
+        12792.973632812502,
+        12890.62984375,
+        12988.286054687502,
+        13085.942265625,
+        13183.598476562502,
+        13281.2546875,
+        13378.910898437502,
+        13476.567109375,
+        13574.223320312502,
+        13671.87953125,
+        13769.535742187501,
+        13867.191953125,
+        13964.848164062501,
+        14062.504375,
+        14160.160585937501,
+        14257.816796875,
+        14355.473007812501,
+        14453.12921875,
+        14550.785429687501,
+        14648.441640625002,
+        14746.097851562501,
+        14843.754062500002,
+        14941.410273437501,
+        15039.066484375002,
+        15136.722695312501,
+        15234.378906250002,
+        15332.0351171875,
+        15429.691328125002,
+        15527.3475390625,
+        15625.003750000002,
+        15722.6599609375,
+        15820.316171875002,
+        15917.9723828125,
+        16015.628593750002,
+        16113.2848046875,
+        16210.941015625001,
+        16308.5972265625,
+        16406.2534375,
+        16503.9096484375,
+        16601.565859374998,
+        16699.2220703125,
+        16796.87828125,
+        16894.5344921875,
+        16992.190703125,
+        17089.8469140625,
+        17187.503125,
+        17285.1593359375,
+        17382.815546875,
+        17480.4717578125,
+        17578.12796875,
+        17675.7841796875,
+        17773.440390625,
+        17871.096601562498,
+        17968.7528125,
+        18066.4090234375,
+        18164.065234375,
+        18261.721445312498,
+        18359.37765625,
+        18457.0338671875,
+        18554.690078125,
+        18652.346289062498,
+        18750.0025,
+        18847.6587109375,
+        18945.314921875,
+        19042.9711328125,
+        19140.62734375,
+        19238.2835546875,
+        19335.939765625,
+        19433.5959765625,
+        19531.2521875,
+        19628.9083984375,
+        19726.564609375,
+        19824.2208203125,
+        19921.87703125,
+        20019.5332421875,
+        20117.189453125,
+        20214.8456640625,
+        20312.501874999998,
+        20410.1580859375,
+        20507.814296875,
+        20605.4705078125,
+        20703.126718749998,
+        20800.7829296875,
+        20898.439140625,
+        20996.0953515625,
+        21093.7515625,
+        21191.4077734375,
+        21289.063984375,
+        21386.7201953125,
+        21484.37640625,
+        21582.0326171875,
+        21679.688828125,
+        21777.3450390625,
+        21875.00125,
+        21972.6574609375,
+        22070.313671875,
+        22167.9698828125,
+        22265.62609375,
+        22363.2823046875,
+        22460.938515625,
+        22558.5947265625,
+        22656.2509375,
+        22753.907148437498,
+        22851.563359375,
+        22949.2195703125,
+        23046.87578125,
+        23144.5319921875,
+        23242.188203125,
+        23339.8444140625,
+        23437.500625,
+        23535.1568359375,
+        23632.813046875,
+        23730.4692578125,
+        23828.12546875,
+        23925.7816796875,
+        24023.437890625,
+        24121.0941015625,
+        24218.7503125,
+        24316.4065234375,
+        24414.062734375,
+        24511.7189453125,
+        24609.37515625,
+        24707.0313671875,
+        24804.687578125,
+        24902.3437890625,
+        25000.0
+      ],
+      "name": "amount",
+      "num_values": 256,
+      "type": "bucketed",
+      "vocab_size": 259
+    },
+    {
+      "name": "card_product",
+      "num_values": 10,
+      "type": "categorical",
+      "vocab_size": 13
+    },
+    {
+      "name": "country",
+      "num_values": 50,
+      "type": "categorical",
+      "vocab_size": 53
+    },
+    {
+      "name": "avs",
+      "num_values": 5,
+      "type": "categorical",
+      "vocab_size": 8
+    },
+    {
+      "name": "cvv",
+      "num_values": 3,
+      "type": "categorical",
+      "vocab_size": 6
+    },
+    {
+      "name": "device_hash",
+      "num_values": 2000,
+      "type": "categorical",
+      "vocab_size": 2003
+    },
+    {
+      "boundaries": [
+        0.0,
+        12.0,
+        24.0,
+        36.0,
+        48.0,
+        60.0,
+        72.0,
+        84.0,
+        96.0,
+        108.0,
+        120.0
+      ],
+      "name": "customer_tenure",
+      "num_values": 10,
+      "type": "bucketed",
+      "vocab_size": 13
+    }
+  ],
+  "num_features": 15,
+  "num_transactions": 64
+}

encoder/configs/model_nocompress.yaml

@@ -0,0 +1,59 @@
+# Non-compressing encoder configuration (Option A: multi-position pooling).
+#
+# Mirrors parent's StructuredEmbedding pattern at d_lfm=1024: per-feature
+# value tables + feature-type table, summed, expanded to a flat sequence
+# of T_tx * F = 960 tokens. No MLP compression, no separate projector.
+#
+# Categorical heads use `pre_last_tx_mean` to pool the entire tx 62 stripe
+# (positions 930..944) instead of a single position. Restores the
+# holistic per-tx summary that the compress variant got for free.
+#
+# Fraud head keeps `last_tx_mean` — already pools the tx 63 stripe
+# (positions 945..959), no change needed.
+
+architecture:
+  mode: nocompress              # → StructuredEncoder, T=960, no projector
+
+backbone:
+  hf_path: LiquidAI/LFM2.5-350M-Base
+  dtype: bfloat16
+  lora:
+    enabled: true
+    r: 16
+    alpha: 32
+    dropout: 0.05
+    strict_attention_only: false
+
+heads:
+  fraud:
+    output_dim: 1
+    loss: bce
+    pool: last_tx_mean          # mean of positions 945..959 (last tx)
+    target: sequence_label
+    weight: 1.0
+    mlp_hidden: 128
+    dropout: 0.1
+  next_merchant:
+    output_dim: 10003
+    loss: ce
+    pool: pre_last_tx_mean      # mean of positions 930..944 (tx 62 stripe)
+    target: "feature:5"
+    weight: 0.5
+    mlp_hidden: 128
+    dropout: 0.1
+  amount_range:
+    output_dim: 16
+    loss: ce
+    pool: pre_last_tx_mean
+    target: amount_range
+    weight: 0.5
+    mlp_hidden: 128
+    dropout: 0.1
+  mcc:
+    output_dim: 103
+    loss: ce
+    pool: pre_last_tx_mean
+    target: "feature:4"
+    weight: 0.5
+    mlp_hidden: 128
+    dropout: 0.1

encoder/src/data/loader.py

@@ -0,0 +1,139 @@
+"""Data loader for the encoder subproject.
+
+Reuses the parent's `FinetuneDataset` verbatim — the parent's tokenized arrays
+at `data/synthetic/` are already shaped `(N, 64, 15)`, which is exactly what
+the per-transaction encoder needs. This module is a thin orchestrator that
+resolves data paths (via the `encoder/data/synthetic -> ../../data/synthetic`
+symlink), builds train/val/test loaders, and exposes a fingerprint-verification
+helper so accidental data regeneration breaks fast.
+
+Why we don't define a new Dataset class: the encoder's input contract is
+identical to the parent's (`(B, 64, 15) int64` plus fraud + amount_range
+labels). The only thing that changes is what the model does with those
+tokens. Keeping the Dataset shared guarantees apples-to-apples comparison.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from src.training.finetune import FinetuneDataset
+
+
+def load_data_arrays(
+    data_dir: Path | str,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, dict[str, np.ndarray]]:
+    """Load raw token arrays + split indices from `data_dir`.
+
+    Returns:
+        token_ids: (N, 64, 15) int16
+        sequence_labels: (N,) int8 (fraud)
+        ar_targets: (N,) int8 last-transaction amount_range, or None if file absent
+        splits: dict with keys 'train' / 'val' / 'test', each int64 indices
+    """
+    data_dir = Path(data_dir)
+    token_ids = np.load(data_dir / "token_ids.npy")
+    sequence_labels = np.load(data_dir / "sequence_labels.npy")
+    splits = dict(np.load(data_dir / "split_indices.npz"))
+
+    ar_path = data_dir / "amount_range_labels.npy"
+    ar_targets: np.ndarray | None = None
+    if ar_path.exists():
+        # Parent stores per-transaction amount_range as (N, 64). The head
+        # targets the LAST transaction's amount bucket, so we slice [:, -1].
+        ar_all = np.load(ar_path)
+        ar_targets = ar_all[:, -1]
+
+    return token_ids, sequence_labels, ar_targets, splits
+
+
+def verify_fingerprint(data_dir: Path | str, expected: str) -> None:
+    """Raise if data fingerprint differs from `expected`.
+
+    Catches the silent failure where data has been regenerated under us —
+    in which case head-to-head comparison numbers against the parent's
+    already-published eval.md.json would not be apples-to-apples.
+    """
+    fp_path = Path(data_dir) / "fingerprint.txt"
+    if not fp_path.exists():
+        raise FileNotFoundError(
+            f"No fingerprint.txt at {fp_path}. Encoder relies on the parent's "
+            f"data/synthetic/ for head-to-head; regenerate via parent's "
+            f"`python -m scripts.generate` if missing.",
+        )
+    actual = fp_path.read_text().strip()
+    if actual != expected:
+        raise ValueError(
+            f"Data fingerprint mismatch:\n"
+            f"  expected: {expected}\n"
+            f"  actual:   {actual}\n"
+            f"Data has been regenerated since this config was pinned. Head-to-head "
+            f"comparison against the parent's eval.md.json would not be valid.",
+        )
+
+
+def build_loaders(
+    data_dir: Path | str,
+    batch_size: int = 32,
+    label_fraction: float = 1.0,
+    seed: int = 42,
+    num_workers: int = 4,
+) -> tuple[DataLoader, DataLoader, DataLoader]:
+    """Build train/val/test DataLoaders.
+
+    Args:
+        data_dir: path to the tokenized synthetic arrays (symlink to parent OK).
+        batch_size: applied to all three loaders.
+        label_fraction: subsample fraction of `train` indices for the
+            label-scarcity sweep (1.0 = full, 0.10 = 10%, 0.01 = 1%). Val and
+            test are never subsampled.
+        seed: RNG seed for the train-subset selection. Same seed as the parent's
+            scarcity protocol so the head-to-head selects the same training
+            subsets across both architectures.
+        num_workers: DataLoader worker count for train. Val/test use half.
+
+    Returns:
+        (train_loader, val_loader, test_loader)
+    """
+    token_ids, sequence_labels, ar_targets, splits = load_data_arrays(data_dir)
+
+    train_indices = splits["train"]
+    if label_fraction < 1.0:
+        # np.random.RandomState (not Generator) to match the parent's
+        # subsampling RNG exactly. Same seed -> identical train subset.
+        rng = np.random.RandomState(seed)
+        n_keep = max(1, int(len(train_indices) * label_fraction))
+        train_indices = rng.choice(train_indices, n_keep, replace=False)
+
+    train_ds = FinetuneDataset(token_ids, sequence_labels, train_indices, ar_targets)
+    val_ds = FinetuneDataset(token_ids, sequence_labels, splits["val"], ar_targets)
+    test_ds = FinetuneDataset(token_ids, sequence_labels, splits["test"], ar_targets)
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+        drop_last=True,
+    )
+    eval_workers = max(0, num_workers // 2)
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=eval_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+    test_loader = DataLoader(
+        test_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=eval_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+    return train_loader, val_loader, test_loader

encoder/src/demo/app.py

@@ -0,0 +1,490 @@
+"""Interactive Gradio demo for the encoder-on-LFM2.5 transaction model.
+
+Three-tab structure (Demo / Why Liquid / Integration). Demo is the
+landing tab so a customer immediately sees what the model does; Why
+Liquid is the architectural pitch; Integration is the build-it-yourself
+playbook. Same Gradio theme + CSS as the rest of Liquid's customer-
+facing demos for visual consistency.
+
+The demo is intentionally self-contained: no side-by-side comparison
+against an alternative architecture, no per-tab references to other
+work. The argument is the architecture pattern itself — encoder +
+frozen LFM2.5 backbone + LoRA + multi-head — and the reader gets a
+clean read of it.
+
+Usage:
+    python -m encoder.src.demo.app \\
+        --checkpoint encoder/experiments/.../step_004999.pt \\
+        --model-config encoder/configs/model_nocompress.yaml \\
+        --schema data/schema.yaml \\
+        --data-dir data/synthetic \\
+        --port 7860
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import gradio as gr
+import torch
+
+from src.data.schema import load_schema
+from src.demo.app import DemoData  # reuse parent's curated-test-set loader
+from src.demo.decode import TransactionDecoder
+from src.demo.merchant_catalog import DemoMerchantCatalog
+from src.demo.profile_inference import format_profile_html, infer_profile
+
+from encoder.src.demo.inference import EncoderDemoModel
+from encoder.src.demo.render import (
+    format_amount_predictions,
+    format_fraud_score,
+    format_mcc_predictions,
+    format_merchant_predictions,
+    format_timeline,
+    render_encoder_integration,
+    render_why_encoder,
+)
+
+
+# ---------------------------------------------------------------------------
+# Header
+# ---------------------------------------------------------------------------
+
+_HEADER_HTML = """
+<div style="text-align: center; margin-bottom: 16px;">
+    <h1 style="margin: 0; font-size: 26px; font-weight: 700; color: #171717; letter-spacing: -0.02em;">
+        Encoder on LFM2.5 — Transaction Foundation Model
+    </h1>
+    <p style="color: #737373; margin: 6px 0 0 0; font-size: 13px;
+              font-family: JetBrains Mono, ui-monospace, monospace;">
+        Liquid AI &middot; LFM2.5-350M base &middot; Encoder + LoRA + multi-head
+    </p>
+</div>
+"""
+
+
+# Container width applied to every tab's content so the three surfaces
+# read at the same width. Kept in sync with render.py's _CONTAINER_WIDTH.
+_CONTAINER_WIDTH = "1180px"
+
+
+# ---------------------------------------------------------------------------
+# Demo tab intro — standalone, no comparison framing
+# ---------------------------------------------------------------------------
+
+_DEMO_INTRO_HTML = f"""
+<div style="max-width: {_CONTAINER_WIDTH}; margin: 16px auto 8px auto; padding: 14px 18px;
+            background: #ffffff; border: 1px solid rgba(0,0,0,0.1);
+            border-radius: 12px;">
+    <div style="font-size: 14px; font-weight: 600; color: #171717; margin-bottom: 6px;">
+        Live multi-head inference on a frozen LFM2.5-350M backbone
+    </div>
+    <div style="font-size: 12px; color: #525252; line-height: 1.55;">
+        Pick a curated customer archetype, click <b>Run Inference</b>, watch four task
+        heads predict in parallel from one shared backbone — fraud probability,
+        next-merchant, amount bucket, and merchant category (MCC). The encoder turns
+        the 64-transaction history into 960 pseudo-tokens; the frozen LFM2.5-350M
+        backbone with LoRA processes them; per-task heads pool the hidden states.
+        For the architectural rationale see <i>Why Liquid</i>; for the build-it-yourself
+        playbook see <i>Integration</i>.
+    </div>
+</div>
+"""
+
+
+# ---------------------------------------------------------------------------
+# CSS — scoped to the encoder demo
+# ---------------------------------------------------------------------------
+
+_CSS = f"""
+/* Force light mode regardless of system preference */
+:root, .dark {{ color-scheme: light !important; }}
+
+.gradio-container {{
+    background: #f5f5f5 !important;
+    max-width: 1280px !important;
+    margin: auto !important;
+    padding: 1.5rem !important;
+}}
+
+/* Tabs: pill-style matching the rest of Liquid's customer-facing demos */
+.tabs {{ background: transparent !important; }}
+.tab-nav {{
+    background: #f5f5f5 !important;
+    border: none !important;
+    border-bottom: 1px solid rgba(0,0,0,0.1) !important;
+    gap: 4px !important;
+    padding: 4px 0 !important;
+}}
+.tab-nav button {{
+    font-weight: 500 !important;
+    font-size: 14px !important;
+    color: #737373 !important;
+    background: transparent !important;
+    border: none !important;
+    border-bottom: 2px solid transparent !important;
+    padding: 8px 16px !important;
+    border-radius: 0 !important;
+    transition: all 0.15s ease !important;
+}}
+.tab-nav button:hover {{
+    color: #171717 !important;
+    background: rgba(0,0,0,0.03) !important;
+}}
+.tab-nav button.selected {{
+    color: #171717 !important;
+    font-weight: 600 !important;
+    border-bottom: 2px solid #171717 !important;
+    background: transparent !important;
+}}
+
+/* Run button: pill style, dark fill. Scoped by elem_id so it does not
+   bleed into Gradio's internal radio/checkbox <button> wrappers. */
+#run-inference-btn button,
+button#run-inference-btn {{
+    background: #171717 !important;
+    color: #ffffff !important;
+    border: none !important;
+    border-radius: 9999px !important;
+    padding: 10px 24px !important;
+    font-weight: 600 !important;
+    letter-spacing: -0.01em !important;
+}}
+#run-inference-btn button:hover,
+button#run-inference-btn:hover {{
+    background: #404040 !important;
+}}
+
+/* Wrap the Demo tab's inner content at the same width as the HTML tabs
+   so all three tabs read at consistent width. */
+#demo-tab-container {{
+    max-width: {_CONTAINER_WIDTH} !important;
+    margin: 0 auto !important;
+    padding: 0 16px !important;
+}}
+"""
+
+
+# Gradio theme matching parent demo's design system (monochrome neutral
+# with Inter + JetBrains Mono fonts). Kept in this file rather than a
+# shared module to keep encoder/ a self-contained directory.
+def _build_theme() -> gr.themes.Soft:
+    return gr.themes.Soft(
+        primary_hue="neutral",
+        secondary_hue="neutral",
+        neutral_hue="neutral",
+        font=gr.themes.GoogleFont("Inter"),
+        font_mono=gr.themes.GoogleFont("JetBrains Mono"),
+    ).set(
+        body_background_fill="#f5f5f5",
+        body_text_color="#171717",
+        body_text_color_subdued="#737373",
+        block_background_fill="#ffffff",
+        block_border_color="rgba(0,0,0,0.1)",
+        block_label_background_fill="#f5f5f5",
+        block_label_text_color="#525252",
+        block_title_text_color="#171717",
+        block_shadow="0 1px 3px rgba(0,0,0,0.04)",
+        input_background_fill="#ffffff",
+        input_border_color="rgba(0,0,0,0.1)",
+        input_border_color_focus="#171717",
+        input_placeholder_color="#a3a3a3",
+        panel_background_fill="#fafafa",
+        panel_border_color="rgba(0,0,0,0.06)",
+        border_color_primary="rgba(0,0,0,0.1)",
+        button_primary_background_fill="#171717",
+        button_primary_background_fill_hover="#404040",
+        button_primary_text_color="#ffffff",
+        button_secondary_background_fill="#ffffff",
+        button_secondary_text_color="#525252",
+        button_secondary_border_color="rgba(0,0,0,0.1)",
+        slider_color="#171717",
+        table_border_color="rgba(0,0,0,0.06)",
+        table_even_background_fill="#fafafa",
+        table_odd_background_fill="#ffffff",
+        shadow_spread="0px",
+        color_accent_soft="rgba(0,0,0,0.04)",
+    )
+
+
+# ---------------------------------------------------------------------------
+# App builder
+# ---------------------------------------------------------------------------
+
+
+def _build_demo_tab_contents(
+    model: EncoderDemoModel,
+    data: DemoData,
+    decoder: TransactionDecoder,
+    merchant_catalog: DemoMerchantCatalog,
+    app: gr.Blocks,
+) -> None:
+    """Build the original Multi-Head Demo tab into the current Gradio context.
+
+    Caller owns the outer Blocks/Tab. The `app` argument is the host
+    Blocks instance, needed so we can register the auto-run `app.load(...)`
+    hook against the right Blocks.
+    """
+
+    def on_customer_select(curated_name: str) -> tuple[str, str, str, str, str, str, str, str]:
+        if not curated_name or curated_name not in data.curated_names:
+            return ("Select a customer to see predictions.", "", "", "", "", "", "", "")
+
+        idx = data.get_curated_index(curated_name)
+        token_ids = data.token_ids[idx]
+        is_fraud = bool(data.labels[idx])
+
+        summary = decoder.summarize_customer(token_ids, is_fraud)
+
+        t0 = time.perf_counter()
+        preds = model.run_inference(token_ids)
+        latency_ms = (time.perf_counter() - t0) * 1000
+
+        timeline_html = format_timeline(decoder, token_ids)
+        fraud_html = format_fraud_score(float(preds["fraud"][0]))
+        merchant_html = format_merchant_predictions(
+            preds["next_merchant"], merchant_catalog, k=5,
+        )
+        amount_html = format_amount_predictions(preds["amount_range"], k=5)
+        mcc_html = format_mcc_predictions(preds["mcc"], k=5)
+
+        profile_match = infer_profile(token_ids)
+        profile_html = format_profile_html(profile_match)
+
+        latency_html = (
+            f"<div style='font-family: JetBrains Mono, ui-monospace, monospace; "
+            f"font-size: 11px; color: #737373; padding: 6px 10px; "
+            f"background: #fafafa; border-radius: 6px; display: inline-block;'>"
+            f"Inference: <b style='color: #171717;'>{latency_ms:.0f} ms</b> "
+            f"({'CPU' if model.device.type == 'cpu' else 'GPU'}) "
+            f"&middot; ground truth: "
+            f"<b style='color: {'#EF4444' if is_fraud else '#10B981'};'>"
+            f"{'FRAUD' if is_fraud else 'LEGITIMATE'}</b>"
+            f"</div>"
+        )
+
+        return (
+            summary,
+            timeline_html,
+            profile_html,
+            fraud_html,
+            merchant_html,
+            amount_html,
+            mcc_html,
+            latency_html,
+        )
+
+    gr.HTML(_DEMO_INTRO_HTML)
+
+    with gr.Column(elem_id="demo-tab-container"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.HTML(
+                    "<h3 style='margin: 0 0 8px 0; color: #171717;'>"
+                    "Select Customer</h3>"
+                )
+                curated_dropdown = gr.Dropdown(
+                    choices=data.curated_names,
+                    value=data.curated_names[0] if data.curated_names else None,
+                    label="Curated archetypes",
+                    info="5 legitimate profiles + 3 fraud archetypes",
+                )
+                run_btn = gr.Button(
+                    "Run Inference",
+                    variant="primary",
+                    size="lg",
+                    elem_id="run-inference-btn",
+                )
+                latency_html = gr.HTML("")
+            with gr.Column(scale=2):
+                summary_html = gr.HTML("")
+
+        gr.HTML(
+            "<h3 style='margin: 20px 0 8px 0; color: #171717;'>"
+            "Transaction Timeline</h3>"
+        )
+        timeline_html = gr.HTML("")
+
+        with gr.Row():
+            with gr.Column():
+                gr.HTML(
+                    "<h3 style='margin: 16px 0 8px 0; color: #171717;'>"
+                    "Behavioral Profile</h3>"
+                )
+                profile_html = gr.HTML("")
+            with gr.Column():
+                gr.HTML(
+                    "<h3 style='margin: 16px 0 8px 0; color: #171717;'>"
+                    "Fraud Score</h3>"
+                )
+                fraud_html = gr.HTML("")
+
+        gr.HTML(
+            "<h3 style='margin: 20px 0 8px 0; color: #171717;'>"
+            "Next-Transaction Predictions</h3>"
+        )
+        with gr.Row():
+            with gr.Column():
+                gr.HTML(
+                    "<h4 style='margin: 4px 0; color: #525252; font-size: 13px;'>"
+                    "Next merchant</h4>"
+                )
+                merchant_html = gr.HTML("")
+            with gr.Column():
+                gr.HTML(
+                    "<h4 style='margin: 4px 0; color: #525252; font-size: 13px;'>"
+                    "Amount bucket</h4>"
+                )
+                amount_html = gr.HTML("")
+            with gr.Column():
+                gr.HTML(
+                    "<h4 style='margin: 4px 0; color: #525252; font-size: 13px;'>"
+                    "Merchant category (MCC)</h4>"
+                )
+                mcc_html = gr.HTML("")
+
+    outputs = [
+        summary_html,
+        timeline_html,
+        profile_html,
+        fraud_html,
+        merchant_html,
+        amount_html,
+        mcc_html,
+        latency_html,
+    ]
+
+    run_btn.click(
+        fn=on_customer_select,
+        inputs=[curated_dropdown],
+        outputs=outputs,
+    )
+
+    # Auto-run on first load so the user lands on populated outputs.
+    app.load(
+        fn=on_customer_select,
+        inputs=[curated_dropdown],
+        outputs=outputs,
+    )
+
+
+def _build_why_liquid_tab_contents() -> None:
+    """Render the Why Liquid tab content (HTML pitch)."""
+    gr.HTML(render_why_encoder())
+
+
+def _build_integration_tab_contents() -> None:
+    """Render the Integration tab content (HTML build-it-yourself guide)."""
+    gr.HTML(render_encoder_integration())
+
+
+def create_app(
+    model: EncoderDemoModel,
+    data: DemoData,
+    decoder: TransactionDecoder,
+    merchant_catalog: DemoMerchantCatalog,
+) -> gr.Blocks:
+    """Standalone 3-tab Gradio app for the multi-head encoder demo."""
+    with gr.Blocks(
+        title="Encoder on LFM2.5 — Transaction Foundation Model",
+        css=_CSS,
+        theme=_build_theme(),
+    ) as app:
+        gr.HTML(_HEADER_HTML)
+        with gr.Tabs():
+            with gr.Tab("Demo"):
+                _build_demo_tab_contents(model, data, decoder, merchant_catalog, app)
+            with gr.Tab("Why Liquid"):
+                _build_why_liquid_tab_contents()
+            with gr.Tab("Integration"):
+                _build_integration_tab_contents()
+    return app
+
+
+# ---------------------------------------------------------------------------
+# Entrypoint
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Encoder + LFM2.5-350M interactive demo",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=None,
+        help="Trained checkpoint path. Omit to run with random-init (debug only).",
+    )
+    parser.add_argument(
+        "--model-config",
+        type=Path,
+        default=Path("encoder/configs/model_nocompress.yaml"),
+    )
+    parser.add_argument(
+        "--schema",
+        type=Path,
+        default=Path("data/schema.yaml"),
+    )
+    parser.add_argument(
+        "--data-dir",
+        type=Path,
+        default=Path("data/synthetic"),
+        help="Test-set data directory (token_ids.npy, sequence_labels.npy, etc.)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        choices=["cpu", "cuda", "mps"],
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="float32",
+        choices=["float32", "bfloat16"],
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=7860,
+    )
+    args = parser.parse_args()
+
+    dtype = torch.float32 if args.dtype == "float32" else torch.bfloat16
+
+    print("Loading schema, data, model...")
+    schema = load_schema(args.schema)
+    data = DemoData(args.data_dir, schema)
+    merchant_catalog = DemoMerchantCatalog(schema)
+    decoder = TransactionDecoder(schema, merchant_catalog)
+
+    model = EncoderDemoModel(
+        model_config_path=args.model_config,
+        schema_path=args.schema,
+        checkpoint_path=args.checkpoint,
+        dtype=dtype,
+        device=args.device,
+    )
+    print(f"Model loaded: {model.checkpoint_status}")
+    pc = model.num_params()
+    print(f"Params — total: {pc['total']:,} / trainable: {pc['trainable']:,}")
+
+    app = create_app(model, data, decoder, merchant_catalog)
+
+    # Walk a few ports if 7860 is in use.
+    for port in range(args.port, args.port + 10):
+        try:
+            app.launch(server_port=port, server_name="0.0.0.0")
+            break
+        except OSError as e:
+            if "Cannot find empty port" in str(e) or "Address already in use" in str(e):
+                print(f"  port {port} in use, trying {port + 1}...")
+                continue
+            raise
+
+
+if __name__ == "__main__":
+    main()

encoder/src/demo/copilot_app_unified.py

@@ -1,24 +1,29 @@
-"""Unified Co-Pilot: Dispute + Collections + Fraud in one Gradio app.
+"""Unified Co-Pilot: 6 tabs over one Gradio Blocks.
 
-Three Tabs on a single Blocks. Each tab is the per-surface content
-built by the surface's `_build_tab(model)` function. Models are loaded
-once at startup and reused — each surface owns its own checkpoint +
-config + encoder marker flags.
-
-Layout:
-    [shared header]
+Top-level layout:
     Tabs:
-      - "Dispute" → Dispute Co-Pilot tab (cast, complaint, score)
-      - "Collections" → treatment scoreboard tab
-      - "Fraud" → two-distribution stage+type tab
+      - "Multi-Head Demo" — original encoder demo (4 task heads, V3 nocompress
+                            + meanpool checkpoint, curated customer dropdown)
+      - "Why Liquid" — architectural pitch (original demo content)
+      - "Integration" — build-it-yourself guide (original demo content)
+      - "Dispute Co-Pilot" — friendly-fraud classifier + attribution
+      - "Collections Co-Pilot" — treatment-response scoreboard
+      - "Fraud Co-Pilot" — pattern stage + type classifier
 
-This is the deployment target for the Hugging Face Space.
+Each tab's content is a composable `_build_..._tab_contents(...)` helper
+exported from the per-surface app module. Models are loaded once at
+startup (4 model instances total — V3 multi-head + 3 surface-specific
+multi-surface).
 
 CLI:
     python -m encoder.src.demo.copilot_app_unified \\
+        --multihead-checkpoint encoder/experiments/.../step_004999_slim.pt \\
+        --multihead-config encoder/configs/model_nocompress.yaml \\
+        --multihead-data-dir data/synthetic \\
         --dispute-checkpoint encoder/experiments/dispute_legitimacy_v7/demo_checkpoint.pt \\
         --collections-checkpoint encoder/experiments/collections_v3/demo_checkpoint.pt \\
         --fraud-checkpoint encoder/experiments/fraud_pattern_v1/demo_checkpoint.pt \\
+        --cast-histories data/synthetic/cast_token_ids.npy \\
         --port 7860
 """
 
@@ -30,6 +35,20 @@ from pathlib import Path
 import gradio as gr
 import torch
 
+# Original multi-head demo (V3) imports
+from src.data.schema import load_schema
+from src.demo.app import DemoData
+from src.demo.decode import TransactionDecoder
+from src.demo.merchant_catalog import DemoMerchantCatalog
+
+from encoder.src.demo.app import (
+    _build_demo_tab_contents,
+    _build_integration_tab_contents,
+    _build_why_liquid_tab_contents,
+    _build_theme,
+    _CSS,
+    _HEADER_HTML,
+)
 from encoder.src.demo.copilot_app import _build_tab as _build_dispute_tab
 from encoder.src.demo.copilot_app_collections import (
     _build_tab as _build_collections_tab,
@@ -44,71 +63,77 @@ from encoder.src.demo.copilot_inference_collections import (
 from encoder.src.demo.copilot_inference_fraud_pattern import (
     FraudPatternCopilotModel,
 )
-
-
-_CONTAINER_WIDTH_PX = 1200
-
-_INK = "#171717"
-_INK_DIM = "#525252"
-
-
-def _render_unified_header() -> str:
-    """Shared top header for the unified Co-Pilot.
-
-    States: backbone, the three surfaces, the architecture name. Sits
-    above the Tabs.
-    """
-    return f"""
-    <div style="text-align: center; margin-bottom: 18px;">
-        <div style="
-            font-size: 11px;
-            color: {_INK_DIM};
-            font-family: 'JetBrains Mono', ui-monospace, monospace;
-            text-transform: uppercase;
-            letter-spacing: 0.08em;
-            margin-bottom: 6px;
-        ">
-            Liquid AI &middot; LFM2.5-350M backbone &middot; encoder + LoRA per surface
-        </div>
-        <h1 style="
-            margin: 0;
-            font-size: 28px;
-            font-weight: 700;
-            color: {_INK};
-            letter-spacing: -0.02em;
-        ">
-            Transaction Co-Pilot
-        </h1>
-        <div style="
-            margin-top: 4px;
-            font-size: 13px;
-            color: {_INK_DIM};
-        ">
-            One backbone, three surfaces: Dispute &middot; Collections &middot; Fraud
-        </div>
-    </div>
-    """
+from encoder.src.demo.inference import EncoderDemoModel
 
 
 def build_unified_ui(
+    multihead_model: EncoderDemoModel,
+    multihead_data: DemoData,
+    multihead_decoder: TransactionDecoder,
+    multihead_merchant_catalog: DemoMerchantCatalog,
     dispute_model: CopilotModel,
     collections_model: CollectionsCopilotModel,
     fraud_model: FraudPatternCopilotModel,
 ) -> gr.Blocks:
-    """Compose the three surface tabs into one Blocks."""
-    with gr.Blocks(title="Transaction Co-Pilot — Liquid AI") as demo:
-        gr.HTML(_render_unified_header())
+    """Compose the original 3-tab demo + 3 new Co-Pilot tabs into one Blocks.
+
+    Uses the original demo's theme + CSS so the visual vocabulary stays
+    consistent across tabs.
+    """
+    # Gradio 6.0 moved `css` and `theme` from Blocks() to launch(); the
+    # caller passes them in via demo.queue().launch(theme=..., css=...).
+    # We stash them on the returned Blocks instance for the caller.
+    with gr.Blocks(title="Transaction Encoder — Liquid AI") as app:
+        gr.HTML(_HEADER_HTML)
         with gr.Tabs():
-            with gr.Tab("Dispute"):
+            with gr.Tab("Multi-Head Demo"):
+                _build_demo_tab_contents(
+                    multihead_model,
+                    multihead_data,
+                    multihead_decoder,
+                    multihead_merchant_catalog,
+                    app,
+                )
+            with gr.Tab("Why Liquid"):
+                _build_why_liquid_tab_contents()
+            with gr.Tab("Integration"):
+                _build_integration_tab_contents()
+            with gr.Tab("Dispute Co-Pilot"):
                 _build_dispute_tab(dispute_model)
-            with gr.Tab("Collections"):
+            with gr.Tab("Collections Co-Pilot"):
                 _build_collections_tab(collections_model)
-            with gr.Tab("Fraud"):
+            with gr.Tab("Fraud Co-Pilot"):
                 _build_fraud_tab(fraud_model)
-    return demo
+    return app
+
+
+def _load_multihead(
+    checkpoint: Path,
+    config: Path,
+    schema: Path,
+    data_dir: Path,
+    dtype: torch.dtype,
+    device: str,
+) -> tuple[EncoderDemoModel, DemoData, TransactionDecoder, DemoMerchantCatalog]:
+    """Load the original multi-head V3 model + curated test data."""
+    print(f"[multihead] schema + data ...")
+    schema_cfg = load_schema(schema)
+    data = DemoData(data_dir, schema_cfg)
+    merchant_catalog = DemoMerchantCatalog(schema_cfg)
+    decoder = TransactionDecoder(schema_cfg, merchant_catalog)
+    print(f"[multihead] loading EncoderDemoModel ...")
+    model = EncoderDemoModel(
+        model_config_path=config,
+        schema_path=schema,
+        checkpoint_path=checkpoint,
+        dtype=dtype,
+        device=device,
+    )
+    print(f"[multihead] checkpoint: {model.checkpoint_status}")
+    return model, data, decoder, merchant_catalog
 
 
-def _load_all(
+def _load_copilots(
     dispute_checkpoint: Path,
     dispute_config: Path,
     collections_checkpoint: Path,
@@ -116,39 +141,37 @@ def _load_all(
     fraud_checkpoint: Path,
     fraud_config: Path,
     schema: Path,
-    histories: Path,
+    cast_histories: Path,
     dispute_cast: Path,
     collections_cast: Path,
     fraud_cast: Path,
     device: torch.device,
 ) -> tuple[CopilotModel, CollectionsCopilotModel, FraudPatternCopilotModel]:
-    """Load all three multi-surface models. Each surface has its own
-    backbone copy in memory; for CPU deployments this is acceptable
-    (3 × ~700MB fp32 = ~2.1GB)."""
-    print(f"[1/3] Loading Dispute on {device} ...")
+    """Load the three Co-Pilot surfaces. Each has its own backbone copy."""
+    print(f"[copilot 1/3] loading Dispute ...")
     dispute_model = CopilotModel.from_paths(
         checkpoint_path=dispute_checkpoint,
         model_config_path=dispute_config,
         schema_path=schema,
-        histories_path=histories,
+        histories_path=cast_histories,
         cast_path=dispute_cast,
         device=device,
     )
-    print(f"[2/3] Loading Collections on {device} ...")
+    print(f"[copilot 2/3] loading Collections ...")
     collections_model = CollectionsCopilotModel.from_paths(
         checkpoint_path=collections_checkpoint,
         model_config_path=collections_config,
         schema_path=schema,
-        histories_path=histories,
+        histories_path=cast_histories,
         cast_path=collections_cast,
         device=device,
     )
-    print(f"[3/3] Loading Fraud on {device} ...")
+    print(f"[copilot 3/3] loading Fraud ...")
     fraud_model = FraudPatternCopilotModel.from_paths(
         checkpoint_path=fraud_checkpoint,
         model_config_path=fraud_config,
         schema_path=schema,
-        histories_path=histories,
+        histories_path=cast_histories,
         cast_path=fraud_cast,
         device=device,
     )
@@ -157,8 +180,26 @@ def _load_all(
 
 def main() -> None:
     parser = argparse.ArgumentParser(
-        description="Unified Transaction Co-Pilot Gradio app",
+        description="Unified Transaction Encoder Gradio app (6 tabs)",
+    )
+    # --- multi-head V3 ---
+    parser.add_argument(
+        "--multihead-checkpoint",
+        type=Path,
+        default=Path("encoder/experiments/nocompress_meanpool/"
+                     "encoder_sft_20260519_144916/checkpoints/step_004999_slim.pt"),
     )
+    parser.add_argument(
+        "--multihead-config",
+        type=Path,
+        default=Path("encoder/configs/model_nocompress.yaml"),
+    )
+    parser.add_argument(
+        "--multihead-data-dir",
+        type=Path,
+        default=Path("data/synthetic"),
+    )
+    # --- dispute ---
     parser.add_argument(
         "--dispute-checkpoint",
         type=Path,
@@ -169,6 +210,12 @@ def main() -> None:
         type=Path,
         default=Path("encoder/configs/model_dispute_legitimacy.yaml"),
     )
+    parser.add_argument(
+        "--dispute-cast",
+        type=Path,
+        default=Path("encoder/data/demo_cast.json"),
+    )
+    # --- collections ---
     parser.add_argument(
         "--collections-checkpoint",
         type=Path,
@@ -179,6 +226,12 @@ def main() -> None:
         type=Path,
         default=Path("encoder/configs/model_collections.yaml"),
     )
+    parser.add_argument(
+        "--collections-cast",
+        type=Path,
+        default=Path("encoder/data/collections_cast.json"),
+    )
+    # --- fraud ---
     parser.add_argument(
         "--fraud-checkpoint",
         type=Path,
@@ -190,29 +243,21 @@ def main() -> None:
         default=Path("encoder/configs/model_fraud_pattern.yaml"),
     )
     parser.add_argument(
-        "--schema",
-        type=Path,
-        default=Path("data/schema.yaml"),
-    )
-    parser.add_argument(
-        "--histories",
-        type=Path,
-        default=Path("data/synthetic/token_ids.npy"),
-    )
-    parser.add_argument(
-        "--dispute-cast",
+        "--fraud-cast",
         type=Path,
-        default=Path("encoder/data/demo_cast.json"),
+        default=Path("encoder/data/fraud_pattern_cast.json"),
     )
+    # --- shared ---
     parser.add_argument(
-        "--collections-cast",
+        "--schema",
         type=Path,
-        default=Path("encoder/data/collections_cast.json"),
+        default=Path("data/schema.yaml"),
     )
     parser.add_argument(
-        "--fraud-cast",
+        "--cast-histories",
         type=Path,
-        default=Path("encoder/data/fraud_pattern_cast.json"),
+        default=Path("data/synthetic/token_ids.npy"),
+        help="Histories file for the Co-Pilot tabs (subset of 18 cast customers).",
     )
     parser.add_argument(
         "--device",
@@ -220,12 +265,33 @@ def main() -> None:
         default="cpu",
         choices=["cpu", "cuda", "mps"],
     )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="float32",
+        choices=["float32", "bfloat16"],
+    )
     parser.add_argument("--port", type=int, default=7860)
     parser.add_argument("--share", action="store_true")
     args = parser.parse_args()
 
     device = torch.device(args.device)
-    dispute_model, collections_model, fraud_model = _load_all(
+    multihead_dtype = (
+        torch.float32 if args.dtype == "float32" else torch.bfloat16
+    )
+
+    multihead_model, multihead_data, multihead_decoder, multihead_merchant = (
+        _load_multihead(
+            checkpoint=args.multihead_checkpoint,
+            config=args.multihead_config,
+            schema=args.schema,
+            data_dir=args.multihead_data_dir,
+            dtype=multihead_dtype,
+            device=args.device,
+        )
+    )
+
+    dispute_model, collections_model, fraud_model = _load_copilots(
         dispute_checkpoint=args.dispute_checkpoint,
         dispute_config=args.dispute_config,
         collections_checkpoint=args.collections_checkpoint,
@@ -233,29 +299,30 @@ def main() -> None:
         fraud_checkpoint=args.fraud_checkpoint,
         fraud_config=args.fraud_config,
         schema=args.schema,
-        histories=args.histories,
+        cast_histories=args.cast_histories,
         dispute_cast=args.dispute_cast,
         collections_cast=args.collections_cast,
         fraud_cast=args.fraud_cast,
         device=device,
     )
-    print("  all three surfaces loaded.")
 
-    demo = build_unified_ui(dispute_model, collections_model, fraud_model)
+    print("all four models loaded.")
+
+    demo = build_unified_ui(
+        multihead_model=multihead_model,
+        multihead_data=multihead_data,
+        multihead_decoder=multihead_decoder,
+        multihead_merchant_catalog=multihead_merchant,
+        dispute_model=dispute_model,
+        collections_model=collections_model,
+        fraud_model=fraud_model,
+    )
     demo.queue().launch(
         server_name="0.0.0.0",
         server_port=args.port,
         share=args.share,
-        theme=gr.themes.Default(
-            font=["Inter", "system-ui", "sans-serif"],
-            font_mono=["JetBrains Mono", "ui-monospace", "monospace"],
-        ),
-        css=f"""
-        .gradio-container {{
-            max-width: {_CONTAINER_WIDTH_PX}px !important;
-            background: #fafafa !important;
-        }}
-        """,
+        theme=_build_theme(),
+        css=_CSS,
     )
 
 

encoder/src/demo/inference.py

@@ -0,0 +1,152 @@
+"""Inference path for the encoder demo.
+
+Loads V3 (Option A: nocompress + mean-pool + attn-LoRA) and exposes a
+single `run_inference(token_ids)` function returning multi-head
+predictions in a UI-friendly shape.
+
+Single model, no side-by-side — different from the parent's
+pretrained-vs-random comparison. The encoder demo's narrative is about
+the architecture itself, not the value of pretraining.
+
+Runs on CPU (fp32) for HF Spaces basic tier and local Mac smoke tests.
+On the production H100 (192.222.55.165 / demos.liquid.ai surface), the
+caller can pass `dtype=torch.bfloat16, device="cuda"` for the bf16 fast
+path.
+
+Performance note: even on CPU, ~64-tx inference is dominated by the
+backbone's 16-layer forward pass over 960 pseudo-tokens. Expect
+~3-8 seconds per inference on a Mac M-series CPU. For an interactive
+demo we cache the model in memory and accept the latency; the user
+selects a customer, clicks Run, waits ~5 seconds, sees outputs. GPU
+deployment cuts this to <100ms.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import yaml
+
+from src.data.schema import SchemaConfig, load_schema
+from src.training.trainer_utils import load_checkpoint
+
+from encoder.src.model.transaction_fm import build_transaction_fm
+
+
+class EncoderDemoModel:
+    """Wraps the V3 stack with cached inference for the demo.
+
+    Construction:
+        model = EncoderDemoModel(
+            model_config_path="encoder/configs/model_nocompress.yaml",
+            schema_path="data/schema.yaml",
+            checkpoint_path="encoder/experiments/.../step_004999.pt",
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    Inference:
+        results = model.run_inference(token_ids)  # (64, 15) int64
+        # results: dict[str, np.ndarray]
+        #   fraud:         shape (1,)  — sigmoid probability in [0, 1]
+        #   next_merchant: shape (10003,) — softmax over merchant_id vocab
+        #   amount_range:  shape (16,)    — softmax over 16 amount buckets
+        #   mcc:           shape (103,)   — softmax over MCC vocab
+    """
+
+    def __init__(
+        self,
+        model_config_path: str | Path,
+        schema_path: str | Path,
+        checkpoint_path: str | Path | None = None,
+        dtype: torch.dtype = torch.float32,
+        device: str = "cpu",
+    ) -> None:
+        self.device = torch.device(device)
+        self.dtype = dtype
+
+        with open(model_config_path) as f:
+            mcfg = yaml.safe_load(f)
+        self.schema: SchemaConfig = load_schema(schema_path)
+        self.head_configs = mcfg["heads"]
+
+        # device_map=None + manual .to(device) below. We deliberately don't
+        # use device_map="auto" because the encoder/projector/heads are
+        # constructed separately and would otherwise stay on CPU.
+        self.model = build_transaction_fm(
+            schema=self.schema,
+            head_configs=self.head_configs,
+            model_path=mcfg["backbone"]["hf_path"],
+            architecture_cfg=mcfg.get("architecture"),
+            encoder_cfg=mcfg.get("encoder"),
+            projector_cfg=mcfg.get("projector"),
+            lora_cfg=mcfg["backbone"]["lora"],
+            dtype=dtype,
+            device_map=None,
+        ).to(self.device)
+        self.model.eval()
+
+        self.checkpoint_status = "no checkpoint (random-init)"
+        if checkpoint_path is not None:
+            ckpt_path = Path(checkpoint_path)
+            if not ckpt_path.exists():
+                raise FileNotFoundError(
+                    f"Checkpoint not found: {ckpt_path}",
+                )
+            # Peek at the checkpoint to detect slim variant (LFM base
+            # stripped). Slim checkpoints load with strict=False because
+            # the freshly-loaded LFM base keys are missing from the
+            # state_dict — but the runtime values are identical to what
+            # was originally in the checkpoint, so the math is unchanged.
+            ckpt_peek = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+            is_slim = ckpt_peek.get("model_state_dict_slim", False)
+            del ckpt_peek  # release memory before load_checkpoint re-reads
+
+            ckpt = load_checkpoint(ckpt_path, self.model, strict=not is_slim)
+            step = ckpt.get("step", "?")
+            slim_note = " (slim — LFM base from HF cache)" if is_slim else ""
+            self.checkpoint_status = f"step {step}{slim_note}"
+
+    @torch.no_grad()
+    def run_inference(self, token_ids: np.ndarray) -> dict[str, np.ndarray]:
+        """Run all heads on one customer sequence.
+
+        Args:
+            token_ids: (64, 15) int64 numpy array — one customer's
+                64-transaction history.
+
+        Returns:
+            dict mapping head name to its prediction array. Fraud is a
+            sigmoid probability; the rest are softmax-normalized class
+            distributions.
+        """
+        tensor = torch.from_numpy(token_ids).unsqueeze(0).long().to(self.device)
+        # tensor: (1, 64, 15)
+        predictions = self.model(tensor)
+        # predictions: dict[head_name, logits]
+        #   fraud:         (1, 1)
+        #   next_merchant: (1, 10003)
+        #   amount_range:  (1, 16)
+        #   mcc:           (1, 103)
+
+        results: dict[str, np.ndarray] = {}
+        for name, logits in predictions.items():
+            if name == "fraud":
+                prob = torch.sigmoid(logits).squeeze().cpu().numpy()
+                # Cast to scalar then back to 1-element array for uniform downstream handling
+                results[name] = np.array([float(prob)])
+            else:
+                probs = F.softmax(logits, dim=-1).squeeze(0).cpu().numpy()
+                results[name] = probs
+
+        return results
+
+    def num_params(self) -> dict[str, int]:
+        """Param breakdown for the architecture-display card in the UI."""
+        total = sum(p.numel() for p in self.model.parameters())
+        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        return {"total": total, "trainable": trainable}

encoder/src/demo/render.py

@@ -0,0 +1,728 @@
+"""Encoder-specific render module.
+
+Reuses parent's format functions (fraud bar, top-k predictions, timeline)
+directly via import — those are model-agnostic. The encoder demo only
+adds two pieces of new content:
+
+1. `render_why_encoder()` — the architectural pitch for the encoder pattern
+2. `render_encoder_integration()` — build-it-yourself integration guide
+
+The fraud / merchant / amount / mcc / timeline / profile formatters are
+unchanged from parent. We import them so the encoder demo's prediction
+cards are visually consistent.
+"""
+
+from __future__ import annotations
+
+# Re-export parent's format functions so encoder app.py can do a single
+# import from this module.
+from src.demo.render import (  # noqa: F401
+    format_amount_predictions,
+    format_fraud_score,
+    format_mcc_predictions,
+    format_merchant_predictions,
+    format_timeline,
+    format_topk_predictions,
+)
+
+# Liquid design tokens
+_TEXT = "#171717"
+_TEXT_MUTED = "#525252"
+_TEXT_DIM = "#737373"
+_BG_CARD = "#ffffff"
+_BG_CARD_ALT = "#fafafa"
+_BORDER = "rgba(0,0,0,0.1)"
+_BORDER_SUBTLE = "rgba(0,0,0,0.05)"
+_ACCENT_GREEN = "#10B981"
+_ACCENT_BLUE = "#3B82F6"
+_ACCENT_AMBER = "#F59E0B"
+_ACCENT_PURPLE = "#7c3aed"
+_RADIUS_CARD = "16px"
+_RADIUS_SM = "8px"
+_FONT_MONO = "JetBrains Mono, ui-monospace, SFMono-Regular, monospace"
+
+# Single max-width applied to every tab's content so the three surfaces
+# read at the same width. Picked to fit the Gradio container (1280px)
+# with a small inset on either side.
+_CONTAINER_WIDTH = "1180px"
+
+
+def render_why_encoder() -> str:
+    """Why Liquid tab content.
+
+    Opens with the buyer's problem (multi-customer / multi-task
+    transaction-FM economics), then the published precedent
+    (LFM2.5-Audio / LFM2.5-VL), then the architectural and operational
+    properties, then scope-of-claim. Written for an external audience —
+    no internal codenames, no design-log register, no "we claim" framing.
+    """
+
+    def _table_header(cols: list[str]) -> str:
+        ths = ""
+        for i, c in enumerate(cols):
+            align = "right" if i > 0 else "left"
+            ths += (
+                f'<th style="padding: 6px 10px; font-size: 10px; color: {_TEXT_DIM};'
+                f' text-transform: uppercase; letter-spacing: 0.05em; text-align: {align};'
+                f' font-weight: 600;">{c}</th>'
+            )
+        return f"<tr style='border-bottom: 1px solid {_BORDER};'>{ths}</tr>"
+
+    def _table_row(cells: list[str], highlight: bool = False) -> str:
+        bg = "background: rgba(16,185,129,0.05);" if highlight else ""
+        tds = ""
+        for i, c in enumerate(cells):
+            align = "right" if i > 0 else "left"
+            tds += (
+                f'<td style="padding: 6px 10px; font-family: {_FONT_MONO}; font-size: 11px;'
+                f' color: {_TEXT}; text-align: {align};">{c}</td>'
+            )
+        return f"<tr style='border-bottom: 1px solid {_BORDER_SUBTLE}; {bg}'>{tds}</tr>"
+
+    return f"""
+    <div style="max-width: {_CONTAINER_WIDTH}; margin: 0 auto; padding: 16px;
+                font-family: -apple-system, BlinkMacSystemFont, Segoe UI, Roboto, sans-serif;">
+
+      <!-- Lead -->
+      <h2 style="margin: 0 0 4px 0; color: {_TEXT}; font-size: 22px; font-weight: 700;
+                 letter-spacing: -0.02em;">
+        Why LFM2.5 for Your Transaction Foundation Model
+      </h2>
+      <p style="color: {_TEXT_DIM}; font-size: 13px; margin: 0 0 24px 0; line-height: 1.5;">
+        If you are putting a transaction foundation model into production &mdash;
+        especially across more than one business unit, customer, or downstream task &mdash;
+        the architectural choice determines per-customer training cost,
+        time-to-first-production-task, and the marginal cost of adding the second
+        and third tasks. The encoder-on-pretrained-backbone architecture applies
+        a recipe Liquid AI already ships in LFM2.5-Audio and LFM2.5-VL to
+        discrete-feature payment sequences. Three properties, each with a
+        different kind of evidence.
+      </p>
+
+      <!-- 1. Recipe is validated -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          1. The Recipe Is Already Shipping for Two Other Modalities
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 10px 0;">
+          A small per-modality encoder produces continuous embeddings; a projection
+          adapter (when needed) maps them into the LFM2.5 text backbone's hidden space;
+          LoRA adapts the attention layers per customer. LFM2.5-Audio ingests waveforms
+          this way. LFM2.5-VL ingests vision patches this way. This demo applies the
+          same shape to discrete transaction tokens.
+        </p>
+        <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 10px;">
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_BLUE};
+                        font-weight: 600; margin-bottom: 4px;">LFM2.5-AUDIO</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Audio encoder → projection → LFM2.5 backbone. Ships in production.
+            </div>
+          </div>
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_BLUE};
+                        font-weight: 600; margin-bottom: 4px;">LFM2.5-VL</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Vision encoder → projection → LFM2.5 backbone. Ships at multiple sizes.
+            </div>
+          </div>
+          <div style="padding: 10px; background: rgba(16,185,129,0.06);
+                      border: 1px solid rgba(16,185,129,0.2); border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_GREEN};
+                        font-weight: 600; margin-bottom: 4px;">TRANSACTIONS (THIS DEMO)</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Structured encoder → frozen LFM2.5-350M + LoRA → multi-head outputs.
+            </div>
+          </div>
+        </div>
+      </div>
+
+      <!-- 2. Serving cost / backbone speed -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          2. The Backbone Serves at Production Latency
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 12px 0;">
+          LFM2.5's conv-dominant layer stack gives O(N) prefill scaling on most layers
+          where a pure-attention model pays O(N&sup2;). Published hardware-in-the-loop
+          benchmarks from the LFM2 technical report, S25 / 4K context:
+        </p>
+        <table style="width: 100%; border-collapse: collapse; margin-bottom: 8px;">
+          {_table_header(["Model", "Prefill (tok/s)", "Decode (tok/s)"])}
+          {_table_row(["LFM2-2.6B", "<b>116</b>", "<b>30.0</b>"], highlight=True)}
+          {_table_row(["Qwen3-4B", "35", "11.4"])}
+          {_table_row(["Llama-3.2-3B", "51", "15.8"])}
+        </table>
+        <p style="font-size: 11px; color: {_TEXT_DIM}; margin: 0; line-height: 1.5;">
+          Your serving path is the published LFM2.5 backbone unchanged — only
+          the input side differs from a text deployment. The published latency
+          advantage transfers directly.
+          Source: <a href="https://arxiv.org/abs/2511.23404" style="color: {_TEXT_DIM};
+          text-decoration: underline;">LFM2 Technical Report, arXiv 2511.23404</a>.
+        </p>
+      </div>
+
+      <!-- 3. Frozen base + LoRA is the local maximum -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          3. Frozen Backbone + LoRA Is the Higher-Quality Configuration at Typical Label Budgets
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 10px 0;">
+          Freezing the LFM2.5 backbone and adapting it with LoRA produces higher
+          quality than unfreezing the full backbone end-to-end at typical finserv
+          label budgets. LoRA&rsquo;s low-rank update structure acts as effective
+          regularization; lifting that constraint lets the backbone memorize the
+          training labels rather than generalize. <b>The frozen-backbone commitment
+          is not a quality compromise &mdash; it is the higher-quality operating
+          point.</b>
+        </p>
+        <table style="width: 100%; border-collapse: collapse; margin-bottom: 6px;">
+          {_table_header(["Configuration", "Trainable", "Fraud ROC-AUC", "MCC top-1"])}
+          {_table_row(["Frozen backbone + LoRA (this demo)", "~16M", "<b>0.951</b>", "<b>40.5%</b>"],
+                      highlight=True)}
+          {_table_row(["Full backbone unfreeze", "~370M", "0.900", "38.1%"])}
+        </table>
+        <p style="font-size: 11px; color: {_TEXT_DIM}; margin: 0; line-height: 1.5;">
+          Measured on 200K synthetic sequences (64 transactions × 15 features each).
+          At ~16M trainable parameters (encoder + LoRA + heads), per-customer
+          adaptation is small relative to the deployed footprint and completes
+          in hours, not days.
+        </p>
+      </div>
+
+      <!-- 4. Multi-head, multi-customer -->
+      <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-bottom: 12px;">
+        <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+            One Backbone, Many Heads
+          </h3>
+          <p style="color: {_TEXT_MUTED}; font-size: 12px; line-height: 1.6; margin: 0 0 8px 0;">
+            A single forward pass through the backbone produces hidden states that
+            four task heads pool independently — fraud detection, next-merchant
+            prediction, amount-bucket forecasting, MCC classification. New
+            use-cases (disputes, authorization optimization, AML) add a head, not
+            a foundation model.
+          </p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;">
+            Per-head MLP: ~0.5M params. Add a new task in hours.
+          </div>
+        </div>
+
+        <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+            One Backbone, Many Customers
+          </h3>
+          <p style="color: {_TEXT_MUTED}; font-size: 12px; line-height: 1.6; margin: 0 0 8px 0;">
+            The pretrained LFM2.5 weights ship once. Per-customer training is the
+            encoder + LoRA + heads — under 5% of base size in bf16 slim format.
+            Adding a customer is loading new artifacts on top of the cached
+            backbone, not retraining from scratch.
+          </p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;">
+            Slim per-customer artifact: ~30 MB bf16 at LFM2.5-1.2B scale.
+          </div>
+        </div>
+      </div>
+
+      <!-- 5. Architectural fit -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          The Architecture Matches Transaction Data Structure
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 8px 0;">
+          Transaction data is information-dense locally (within-transaction
+          feature correlations, adjacent-transaction continuity) with sparse
+          long-range signal (behavioral baselines across the full history).
+          LFM2.5 allocates O(N) conv to the dense local patterns and O(N&sup2;)
+          attention to the sparse global ones. A pure transformer would spend
+          O(N&sup2;) compute uniformly.
+        </p>
+        <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 10px;">
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 4px;">Within Transaction</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Merchant determines MCC. Entry mode correlates with amount. Dense,
+              local, often deterministic. A 3-wide conv kernel captures this.
+            </div>
+          </div>
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 4px;">Adjacent Transactions</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Strong temporal continuity. A customer at Starbucks at 8am is
+              likely at a similar merchant tomorrow. Local conv handles it.
+            </div>
+          </div>
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 4px;">Distant Transactions</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Weak but non-zero signal. Behavioral profile matters for fraud
+              baseline. This is where attention earns its quadratic cost.
+            </div>
+          </div>
+        </div>
+      </div>
+
+      <!-- 6. Data sovereignty -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          Your Data, Your Model, Your Infrastructure
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0;">
+          LFM2.5 base weights are open. Liquid licenses the architecture, training
+          recipe, and engineering support. Customers train on their proprietary
+          data behind their firewall. No data leaves customer infrastructure.
+          No dependency on external model APIs. The result is a foundation model
+          the customer owns, adapted to their transaction distribution.
+        </p>
+      </div>
+
+      <!-- Scope: what this demo validates / what a POC would establish -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 16px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 8px 0; font-size: 14px; font-weight: 600;">
+          Scope of Claim
+        </h3>
+        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; font-size: 12px;">
+          <div>
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_GREEN};
+                        font-weight: 600; margin-bottom: 6px; text-transform: uppercase;
+                        letter-spacing: 0.05em;">What this demo validates</div>
+            <ul style="margin: 0; padding-left: 14px; color: {_TEXT_MUTED}; line-height: 1.6;">
+              <li>The encoder-on-pretrained-backbone architecture used by LFM2.5-Audio
+                  and LFM2.5-VL applies to discrete-feature transaction sequences
+                  without modifying the transformers library.</li>
+              <li>Per-customer training touches ~2&ndash;5% of the deployed footprint
+                  and trains in hours rather than days.</li>
+              <li>On synthetic data, frozen-backbone-plus-LoRA outperforms
+                  full-backbone unfreezing on every measured head.</li>
+              <li>One pretrained backbone serves all task heads and is identical
+                  across customer deployments.</li>
+            </ul>
+          </div>
+          <div>
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_AMBER};
+                        font-weight: 600; margin-bottom: 6px; text-transform: uppercase;
+                        letter-spacing: 0.05em;">What a POC on your data would establish</div>
+            <ul style="margin: 0; padding-left: 14px; color: {_TEXT_MUTED}; line-height: 1.6;">
+              <li>Whether synthetic-data quality numbers reproduce on your
+                  transaction distribution.</li>
+              <li>Production-scale quality at LFM2.5-1.2B on your hardware and
+                  sequence lengths (this demo runs at LFM2.5-350M).</li>
+              <li>Inference latency against your authorization-decision budget
+                  at your concurrency.</li>
+              <li>Cross-customer or cross-business-unit transfer of the encoder
+                  and LoRA artifacts.</li>
+            </ul>
+          </div>
+        </div>
+      </div>
+
+      <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM}; text-align: center;">
+        Architecture: <a href="https://arxiv.org/abs/2511.23404" style="color: {_TEXT_DIM};
+        text-decoration: underline;">arXiv 2511.23404</a> &middot;
+        Weights: <a href="https://huggingface.co/LiquidAI" style="color: {_TEXT_DIM};
+        text-decoration: underline;">huggingface.co/LiquidAI</a>
+      </div>
+    </div>
+    """
+
+
+def render_encoder_integration() -> str:
+    """Build-it-yourself integration guide.
+
+    Walks the reader through every component a customer team would build
+    to reproduce this demo on their own data: preprocessing, encoder,
+    backbone wiring, heads, postprocessing, training, deployment.
+    Includes hyperparameter cards, gotchas, and an engagement timeline.
+    """
+
+    def _phase_card(num: str, title: str, body: str, detail: str) -> str:
+        return f"""
+        <div style="padding: 14px 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="display: flex; align-items: baseline; gap: 8px; margin-bottom: 6px;">
+            <span style="font-family: {_FONT_MONO}; font-size: 11px; color: {_TEXT_DIM};
+                         font-weight: 600;">{num}</span>
+            <span style="font-size: 14px; font-weight: 600; color: {_TEXT};">{title}</span>
+          </div>
+          <p style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5; margin: 0 0 8px 0;">
+            {body}</p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;
+                      line-height: 1.5;">
+            {detail}</div>
+        </div>"""
+
+    def _gotcha(num: str, title: str, desc: str) -> str:
+        return f"""
+        <div style="display: flex; gap: 8px; padding: 5px 0;
+                    border-bottom: 1px solid {_BORDER_SUBTLE};">
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      font-weight: 600; min-width: 18px;">{num}.</div>
+          <div>
+            <span style="font-size: 12px; font-weight: 600; color: {_TEXT};">{title}</span>
+            <span style="font-size: 12px; color: {_TEXT_MUTED};"> — {desc}</span>
+          </div>
+        </div>"""
+
+    def _pill(text: str) -> str:
+        return (
+            f'<span style="padding: 5px 12px; background: {_TEXT}; color: #fff;'
+            f' border-radius: 9999px; font-family: {_FONT_MONO};'
+            f' font-size: 10px; font-weight: 600;">{text}</span>'
+        )
+
+    arrow = f'<span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>'
+
+    return f"""
+    <div style="max-width: {_CONTAINER_WIDTH}; margin: 0 auto; padding: 16px;
+                font-family: -apple-system, BlinkMacSystemFont, Segoe UI, Roboto, sans-serif;">
+
+      <h2 style="margin: 0 0 4px 0; color: {_TEXT}; font-size: 22px; font-weight: 700;
+                 letter-spacing: -0.02em;">
+        Integration Architecture
+      </h2>
+      <p style="color: {_TEXT_DIM}; font-size: 13px; margin: 0 0 20px 0; line-height: 1.5;">
+        How a customer team builds this stack end to end. Six components, three
+        ship from Liquid (LFM2.5 base weights, training recipes, architecture
+        support); three are customer-bespoke (schema, encoder, task heads).
+        Per-customer adaptation is one ML engineer for a few weeks, not a
+        research project.
+      </p>
+
+      <!-- Pipeline flow -->
+      <div style="display: flex; align-items: center; justify-content: center; gap: 6px;
+                  margin-bottom: 24px; padding: 10px 0; flex-wrap: wrap;">
+        {_pill("Preprocess")}{arrow}
+        {_pill("Encode")}{arrow}
+        {_pill("Backbone + LoRA")}{arrow}
+        {_pill("Heads")}{arrow}
+        {_pill("Postprocess")}{arrow}
+        {_pill("Deploy")}
+      </div>
+
+      <!-- Phase cards -->
+      <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-bottom: 20px;">
+
+        {_phase_card(
+            "1",
+            "Schema & Preprocessing",
+            "Define the discrete feature schema first — features, vocab sizes, ordering. "
+            "Categorical features (merchant_id, MCC, country) map directly to integer IDs. "
+            "Continuous features (amount, days-since-last) get quantile-bucketed into N bins. "
+            "High-cardinality features (10K+ merchants) have their long tail bucketed or "
+            "factored. Reserve 3 token IDs per feature for MASK / OOV / NULL. "
+            "Final per-customer batch shape: (B, T_tx, F).",
+            "Sequence: 64 tx &times; 15 feat = 960 tokens &nbsp;|&nbsp; "
+            "amount → 16 quantile bins &nbsp;|&nbsp; "
+            "merchant_id top-10K + frequency bucketing for the tail &nbsp;|&nbsp; "
+            "unseen values at inference → OOV (ID 1)"
+        )}
+
+        {_phase_card(
+            "2",
+            "Structured Encoder",
+            "One embedding table per feature (sized to its vocab) plus a shared "
+            "feature-type table. Value + type embeddings are summed to identify "
+            "which feature each token represents. The 15 per-tx feature embeddings "
+            "are kept as separate positions in the sequence — compressing them to "
+            "one token per transaction collapses fraud quality (fraud ROC-AUC drops "
+            "to 0.535 on this demo's data), because the per-tx MLP averages away the "
+            "intra-tx feature combinations fraud depends on. This is the same shape "
+            "as the audio and vision encoders' input embedding step.",
+            "Output shape: (B, T_tx*F, d_lfm) = (B, 960, 1024) at LFM2.5-350M &nbsp;|&nbsp; "
+            "value_tables[f](token) + type_table(f) &nbsp;|&nbsp; "
+            "Encoder params dominated by high-cardinality value tables (~14M at 350M)"
+        )}
+
+        {_phase_card(
+            "3",
+            "Projection Adapter (When Needed)",
+            "When the encoder's output dimension matches d_lfm directly, no adapter "
+            "is needed — the encoder outputs flow straight into the backbone. When "
+            "d_encoder &lt; d_lfm (typical at LFM2.5-1.2B where d_lfm=2048), a single "
+            "linear projection lifts the encoder output into the backbone hidden space, "
+            "exactly mirroring the audio/VL projection adapter. Layer init: identity "
+            "for d_encoder=d_lfm, Xavier for the projection case.",
+            "350M: d_lfm=1024, d_encoder=1024, no adapter (identity) &nbsp;|&nbsp; "
+            "1.2B: d_lfm=2048, project from d_encoder=512-1024 → 2048 &nbsp;|&nbsp; "
+            "Adds ~2M params at 1.2B scale"
+        )}
+
+        {_phase_card(
+            "4",
+            "Backbone + LoRA",
+            "Load the pretrained LFM2.5 base from Hugging Face. The backbone&rsquo;s "
+            "parameters are excluded from the optimizer&rsquo;s parameter set during "
+            "training — gradients flow through the backbone to update the upstream "
+            "encoder and downstream heads, but the backbone&rsquo;s own weights are "
+            "never modified. Forward pass executes through all 354M backbone "
+            "parameters at full capacity, at both training and inference time. "
+            "Customer-distribution adaptation enters through (i) LoRA&rsquo;s low-rank "
+            "delta on the attention projections (q_proj / k_proj / v_proj / out_proj) "
+            "and (ii) the per-feature encoder, both trained from scratch on customer "
+            "labels. Encoder outputs are injected via the published "
+            "<code>inputs_embeds</code> hook in <code>Lfm2Model.forward</code>. "
+            "Adding LoRA to the conv layers does not improve quality enough to justify the ~50% increase in training cost; attention-only LoRA is the recommended starting configuration.",
+            "Backbone params excluded from optimizer; backbone forward at full capacity &nbsp;|&nbsp; "
+            "LoRA r=16, &alpha;=32, dropout 0.05 on q_proj / k_proj / v_proj / out_proj &nbsp;|&nbsp; "
+            "PEFT wraps the leaf modules &nbsp;|&nbsp; "
+            "~1M LoRA params at 350M, ~2M at 1.2B"
+        )}
+
+        {_phase_card(
+            "5",
+            "Task Heads",
+            "Per-task downstream heads pool backbone hidden states and predict via "
+            "small MLPs. Fraud (BCE loss) pools the last-transaction stripe — "
+            "mean of positions T-F..T (positions 945..959). Categorical heads "
+            "(next-merchant, amount-bucket, MCC) use cross-entropy and pool the "
+            "<i>pre-last</i> transaction stripe (positions 930..944) to avoid "
+            "leaking the prediction target. New tasks add a head, backbone "
+            "untouched.",
+            "Per-head MLP: 128 hidden, dropout 0.1 &nbsp;|&nbsp; "
+            "Pool: <code>last_tx_mean</code> for sequence tasks &nbsp;|&nbsp; "
+            "Pool: <code>pre_last_tx_mean</code> for next-tx tasks &nbsp;|&nbsp; "
+            "~0.5M params per head"
+        )}
+
+        {_phase_card(
+            "6",
+            "Postprocessing",
+            "Fraud logits → sigmoid → probability in [0, 1]; calibrate against the "
+            "customer's operational threshold (typical: 70% precision @ 60% recall "
+            "for review-queue handoff). Categorical logits → softmax → top-k "
+            "distribution. Use the predicted distribution for downstream "
+            "decisioning, not just argmax — the runner-up matters when the top-1 "
+            "is uncertain. Behavioral attribution: gradient-based saliency on the "
+            "per-feature embeddings identifies which input features drove the score.",
+            "Fraud: sigmoid(logits) → operational threshold &nbsp;|&nbsp; "
+            "Categorical: softmax(logits) → top-k + calibration &nbsp;|&nbsp; "
+            "Saliency: ∂loss/∂value_embed identifies driving features"
+        )}
+      </div>
+
+      <!-- Training recipe (full-width card) -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 16px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 8px 0; font-size: 15px; font-weight: 600;">
+          Training Recipe
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 12px 0;">
+          Single-stage supervised fine-tune on the customer&rsquo;s labelled data — no
+          separate pretraining stage. Three trainable parameter groups (LoRA delta,
+          per-feature encoder, task heads), three learning rates, because each group
+          differs in initialization, parameter scale, and gradient-norm profile.
+        </p>
+        <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 10px;
+                    margin-bottom: 12px;">
+          <div style="padding: 12px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_BLUE};
+                        font-weight: 600; margin-bottom: 6px; letter-spacing: 0.05em;">
+              LORA GROUP</div>
+            <div style="font-family: {_FONT_MONO}; font-size: 11px; color: {_TEXT};
+                        margin-bottom: 4px;">
+              lr = 1e-3 &middot; ~1M params
+            </div>
+            <div style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5;">
+              Low-rank adapters on the backbone&rsquo;s attention projections.
+              Initialized so the LoRA path contributes zero at step 0, then steers
+              attention behavior toward the customer&rsquo;s distribution. Higher LR
+              than the encoder group is fine — the low-rank constraint regularizes
+              the update by construction.
+            </div>
+          </div>
+          <div style="padding: 12px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_BLUE};
+                        font-weight: 600; margin-bottom: 6px; letter-spacing: 0.05em;">
+              ENCODER GROUP</div>
+            <div style="font-family: {_FONT_MONO}; font-size: 11px; color: {_TEXT};
+                        margin-bottom: 4px;">
+              lr = 3e-4 &middot; ~14M params
+            </div>
+            <div style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5;">
+              Per-feature value tables + feature-type table, from random init on the
+              customer&rsquo;s tokenized vocabulary. Lower LR than LoRA because
+              random-init embedding matrices destabilize at higher rates;
+              high-cardinality tables (10K-vocab merchant) dominate gradient norm if
+              not damped.
+            </div>
+          </div>
+          <div style="padding: 12px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_ACCENT_BLUE};
+                        font-weight: 600; margin-bottom: 6px; letter-spacing: 0.05em;">
+              HEADS GROUP</div>
+            <div style="font-family: {_FONT_MONO}; font-size: 11px; color: {_TEXT};
+                        margin-bottom: 4px;">
+              lr = 1e-3 &middot; ~2M params
+            </div>
+            <div style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5;">
+              Per-task MLPs (fraud, next-merchant, amount-bucket, MCC), from random
+              init. Higher LR is fine — small per-head parameter count, well-conditioned
+              loss surface. New downstream tasks attach as additional heads without
+              retraining the backbone or the encoder.
+            </div>
+          </div>
+        </div>
+        <div style="display: grid; grid-template-columns: auto 1fr; gap: 6px 16px;
+                    font-size: 12px; padding: 10px 12px; background: {_BG_CARD_ALT};
+                    border-radius: 8px;">
+          <div style="font-family: {_FONT_MONO}; color: {_TEXT_DIM};">Optimizer</div>
+          <div style="color: {_TEXT_MUTED};">
+            AdamW, &beta; = (0.9, 0.95), weight decay 0.1
+          </div>
+          <div style="font-family: {_FONT_MONO}; color: {_TEXT_DIM};">Schedule</div>
+          <div style="color: {_TEXT_MUTED};">
+            200-step linear warmup, cosine decay to 10% of peak over ~5K steps
+          </div>
+          <div style="font-family: {_FONT_MONO}; color: {_TEXT_DIM};">Precision</div>
+          <div style="color: {_TEXT_MUTED};">
+            bf16 forward and backward, fp32 loss accumulation
+          </div>
+          <div style="font-family: {_FONT_MONO}; color: {_TEXT_DIM};">Multi-task</div>
+          <div style="color: {_TEXT_MUTED};">
+            fraud 1.0, categorical heads 0.5 each — chosen to match per-task
+            gradient norm in the first 200 warmup steps
+          </div>
+          <div style="font-family: {_FONT_MONO}; color: {_TEXT_DIM};">Compute</div>
+          <div style="color: {_TEXT_MUTED};">
+            ~2 hours end-to-end on a single A100 at LFM2.5-350M scale
+          </div>
+        </div>
+      </div>
+
+      <!-- Deployment card (full-width) -->
+      <div style="margin-bottom: 16px;">
+        {_phase_card(
+            "Deploy",
+            "Per-Customer Adapter on a Shared Backbone",
+            "The deployable per-customer artifact is the trained LoRA delta + "
+            "per-feature encoder + task heads. The LFM2.5 base is not included; "
+            "it is loaded once per serving GPU from the public weights. At "
+            "LFM2.5-350M the artifact is ~190 MB in bf16; the unstripped version "
+            "including the base would be ~900 MB. Multi-tenant serving keeps "
+            "one backbone resident on the GPU and switches the active LoRA "
+            "delta + encoder per request, so a new customer adds a small "
+            "adapter rather than a second foundation model. The conv-dominant "
+            "backbone quantizes cleanly to INT8 for cost-sensitive deployments. "
+            "Same code path runs CPU or GPU.",
+            "Per-customer artifact: ~190 MB bf16 at LFM2.5-350M (LoRA + encoder + heads) &nbsp;|&nbsp; "
+            "Multi-tenant: shared backbone forward + per-request LoRA switching &nbsp;|&nbsp; "
+            "CPU (~5s/inference) or H100 (&lt;100ms), same code &nbsp;|&nbsp; "
+            "INT8 quantization clean for conv layers"
+        )}
+      </div>
+
+      <!-- Gotchas -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 20px;">
+        <div style="font-size: 14px; font-weight: 600; color: {_TEXT}; margin-bottom: 8px;">
+          Configuration Choices That Look Right and Aren&rsquo;t
+        </div>
+        {_gotcha("1", "Do not compress each transaction to one token",
+                 "an MLP that averages the 15 feature embeddings into a single "
+                 "vector destroys the intra-tx fraud signal. Fraud ROC-AUC "
+                 "collapsed to 0.535. Keep the full T_tx*F stripe.")}
+        {_gotcha("2", "Pool the pre-last transaction, not the last, for next-tx heads",
+                 "pooling the last-tx stripe for next-merchant prediction "
+                 "leaks the prediction target into the input. Use the prior tx.")}
+        {_gotcha("3", "Don't unfreeze the backbone at typical label budgets",
+                 "full-backbone unfreezing produces lower quality than frozen-plus-LoRA "
+                 "(fraud ROC-AUC 0.951 → 0.900 on this demo's data). LoRA acts as "
+                 "effective regularization; lifting it forces overfitting.")}
+        {_gotcha("4", "Tied embedding heads need SSL-pretrained value tables",
+                 "tying the next-merchant head to the encoder's merchant value table "
+                 "without self-supervised pretraining of those tables reduces "
+                 "next-merchant top-1 from 7.78% to 3.74%. Use a fresh MLP head "
+                 "until SSL pretraining anchors the value tables.")}
+        {_gotcha("5", "Average per-feature losses; do not sum",
+                 "summing CE losses across features makes high-cardinality "
+                 "features (10K-vocab merchant) dominate the gradient. The "
+                 "low-cardinality features stop training.")}
+        {_gotcha("6", "Match the schema fingerprint between training and inference",
+                 "if the tokenizer's vocab changes between training and "
+                 "deployment, the encoder's value tables index into a different "
+                 "semantic space. Embed the fingerprint in checkpoint metadata.")}
+      </div>
+
+      <!-- Engagement model -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 16px;">
+        <div style="font-size: 14px; font-weight: 600; color: {_TEXT}; margin-bottom: 10px;">
+          Typical Engagement
+        </div>
+        <table style="width: 100%; border-collapse: collapse;">
+          <tr style="border-bottom: 1px solid {_BORDER};">
+            <th style="padding: 6px 10px; text-align: left; font-size: 10px;
+                       color: {_TEXT_DIM}; text-transform: uppercase;
+                       letter-spacing: 0.05em; font-weight: 600;">Phase</th>
+            <th style="padding: 6px 10px; text-align: left; font-size: 10px;
+                       color: {_TEXT_DIM}; text-transform: uppercase; font-weight: 600;">
+              Duration</th>
+            <th style="padding: 6px 10px; text-align: left; font-size: 10px;
+                       color: {_TEXT_DIM}; text-transform: uppercase; font-weight: 600;">
+              What Happens</th>
+          </tr>
+          <tr style="border-bottom: 1px solid {_BORDER_SUBTLE};">
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">Discovery</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">1-2 weeks</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Schema design, data sample (~100K-1M sequences), compliance review,
+              architectural fit assessment.
+            </td>
+          </tr>
+          <tr style="border-bottom: 1px solid {_BORDER_SUBTLE};">
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">POC</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">1 week</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Fine-tune encoder + LoRA + heads on customer sample, measurement
+              report, go/no-go recommendation.
+            </td>
+          </tr>
+          <tr style="border-bottom: 1px solid {_BORDER_SUBTLE};">
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">Production</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">2-3 months</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Customer engineering team builds with Liquid architectural support,
+              weekly design review, scale-up to LFM2.5-1.2B.
+            </td>
+          </tr>
+          <tr>
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">Scale</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">Ongoing</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Multi-task expansion (add heads), multi-tenant serving, retraining
+              cadence, architecture evolution.
+            </td>
+          </tr>
+        </table>
+      </div>
+
+      <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM}; text-align: center;">
+        Architecture: <a href="https://arxiv.org/abs/2511.23404" style="color: {_TEXT_DIM};
+        text-decoration: underline;">arXiv 2511.23404</a> &middot;
+        Base weights: <a href="https://huggingface.co/LiquidAI" style="color: {_TEXT_DIM};
+        text-decoration: underline;">huggingface.co/LiquidAI</a>
+      </div>
+    </div>
+    """

encoder/src/model/encoder_heads.py

@@ -0,0 +1,131 @@
+"""Encoder-specific downstream head subclass with multi-position pooling.
+
+Extends parent's `DownstreamHead` with one new pool strategy:
+
+    "pre_last_tx_mean" — pool the entire tx_(T-1) stripe (last num_features
+                         positions BEFORE tx_T), then mean across positions.
+
+Why this exists:
+
+    The encoder non-compress mode emits 15 pseudo-tokens per transaction.
+    The parent's `pre_last_tx` strategy reads a SINGLE position
+    `hidden[:, -(num_features+1), :]` — for our T=960, num_features=15
+    layout that's position 944 = the last feature of tx 62
+    (`customer_tenure` per schema ordering).
+
+    That single position has:
+      - Full causal context through tx 62 (good — attention has propagated
+        everything relevant).
+      - Direct semantics around `customer_tenure` (its value table and
+        type-embedding offset).
+
+    For predicting tx 63's mcc / merchant / amount, the
+    `customer_tenure`-anchored representation is suboptimal. The first
+    nocompress run showed exactly this: fraud ROC-AUC 0.96 (great), but
+    auxiliary head top-1 noticeably below the compress run's holistic-tx
+    summary. The information is in the stripe; we're just reading it
+    through the wrong feature.
+
+    `pre_last_tx_mean` pools across all 15 positions of tx 62 — a
+    holistic tx-62 summary, much closer to what the compress run's
+    single per-tx pseudo-token represented.
+
+When `pre_last_tx` collapses to the same thing:
+
+    For `num_features=1` (compress mode), `pre_last_tx` reads
+    `hidden[:, -2, :]` and `pre_last_tx_mean` reads
+    `mean(hidden[:, -2:-1, :])` = same single position. So compress mode
+    can use either strategy without behavior change; we still default to
+    `pre_last_tx` there for clarity.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from src.model.task_heads import DownstreamHead, HeadConfig, TiedEmbeddingHead
+
+
+class EncoderDownstreamHead(DownstreamHead):
+    """DownstreamHead with `pre_last_tx_mean` pool strategy added.
+
+    Forward + extract_targets + compute_loss are all inherited unchanged —
+    only `pool()` is overridden to handle the new strategy.
+    """
+
+    def __init__(
+        self,
+        config: HeadConfig,
+        hidden_dim: int,
+        num_features: int,
+    ) -> None:
+        super().__init__(config, hidden_dim, num_features)
+
+    def pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """(B, S, D) → (B, D) via head-specific pooling.
+
+        Adds `pre_last_tx_mean` to parent's strategies. Falls through to
+        parent for `last_tx_mean` and `pre_last_tx`.
+        """
+        nf = self.num_features
+        if self.config.pool_strategy == "pre_last_tx_mean":
+            # tx_(T-1) stripe = positions [-2*nf, -nf). For T=960 nf=15:
+            # positions 930..944 inclusive (the 15 features of tx 62).
+            # For T=64 nf=1: position 62 only (single-element mean).
+            stripe = hidden_states[:, -(2 * nf):-nf, :]
+            return stripe.mean(dim=1)
+        return super().pool(hidden_states)
+
+
+class EncoderTiedEmbeddingHead(TiedEmbeddingHead):
+    """TiedEmbeddingHead with `pre_last_tx_mean` pool strategy added.
+
+    Why this exists for the encoder:
+
+        The encoder's per-feature value tables (e.g. merchant_id at
+        `vocab_size=10003, dim=d_lfm=1024`) are exactly the right
+        projection matrix for a high-cardinality classifier head. Instead
+        of learning a fresh `Linear(128, 10003)` from a 128-dim
+        bottleneck — which gave only 7.78% top-1 — we share weights with
+        the encoder's merchant_id table. This recovers the pattern parent
+        observed: a tied embedding head on merchant_id gave 20.8% top-1
+        vs 13.8% for the non-tied head.
+
+    What this changes mechanically:
+
+        - The classifier matrix `weight` is the encoder's value table
+          (shape `(vocab_size, d_lfm)`).
+        - Forward pass: pool hidden → adapter MLP (d_lfm → d_lfm → d_lfm)
+          → matmul through the value table → logits over that feature's
+          vocab.
+        - Gradients flow back into the value table from BOTH the input-
+          embedding side (during forward of every transaction) AND this
+          head (during loss), keeping the table consistent.
+
+    Only differs from parent's TiedEmbeddingHead in the supported pool
+    strategies — adds `pre_last_tx_mean` for the encoder's 15-token-per-tx
+    layout.
+
+    Constraints inherited from parent:
+        - `target_type` must start with "feature:N" (the value-table
+          dimension this head is tied to).
+        - `output_dim` is implicitly the vocab_size of that feature; the
+          head config's `output_dim` is unused.
+    """
+
+    def __init__(
+        self,
+        config: HeadConfig,
+        hidden_dim: int,
+        num_features: int,
+        value_tables: nn.ModuleList,
+    ) -> None:
+        super().__init__(config, hidden_dim, num_features, value_tables)
+
+    def pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        nf = self.num_features
+        if self.config.pool_strategy == "pre_last_tx_mean":
+            stripe = hidden_states[:, -(2 * nf):-nf, :]
+            return stripe.mean(dim=1)
+        return super().pool(hidden_states)

encoder/src/model/lfm_pseudo_token_wrapper.py

@@ -0,0 +1,213 @@
+"""LFM2.5 backbone wrapper that consumes pre-computed pseudo-token embeddings.
+
+This is the load-bearing module. It wraps `transformers.Lfm2Model` so the
+pretrained LFM2.5 backbone processes our 64 per-transaction pseudo-tokens as
+if they were 64 text tokens — the same hook LFM2.5-Audio and LFM2.5-VL use
+to feed their continuous-embedding streams into the LM.
+
+Why `Lfm2Model` and not `Lfm2ForCausalLM`:
+    We don't use the LM head — downstream task heads pool the hidden states
+    instead. `Lfm2Model` returns `last_hidden_state` without the vocab
+    projection, which is exactly what we want and skips ~67M unused params
+    (vocab_size 65536 * d_lfm 1024).
+
+The injection hook (verified against transformers/models/lfm2/modeling_lfm2.py
+lines 523-558):
+
+    Lfm2Model.forward(input_ids=None, inputs_embeds=<our_tensor>) skips
+    embed_tokens entirely, applies RoPE to whatever hidden states arrive,
+    and runs the full conv+attention stack. The XOR guard at line 533
+    requires us to pass exactly one of input_ids / inputs_embeds, so we
+    explicitly pass input_ids=None.
+
+Frozen-base invariant:
+    `freeze_base()` sets requires_grad=False on every parameter at load
+    time. LoRA layers, attached AFTER freezing, have requires_grad=True
+    by default. The unit tests verify this invariant — if any base
+    parameter accumulates a gradient, the test fails.
+
+LoRA target modules:
+    LFM2's attention modules use `q_proj`, `k_proj`, `v_proj`, `out_proj`
+    (NOT LLaMA's `o_proj`). The conv block uses its own `in_proj` and
+    `out_proj`; the SwiGLU MLP uses `w1`, `w2`, `w3`. POC default is
+    attention-only ("q_proj", "k_proj", "v_proj", "out_proj") at r=16,
+    α=32. Escalate to the Spotify-full set at r=64, α=128 if quality is
+    capacity-limited.
+
+Note on `out_proj` collision:
+    `target_modules=["out_proj"]` will match BOTH attention's `out_proj`
+    AND the conv block's `out_proj`. For attention-only LoRA, this is
+    not what we want. We use a regex or explicit attention-scoped names
+    via `target_modules=["q_proj", "k_proj", "v_proj"]` and add
+    `out_proj` only when conv-LoRA is intentional. The current default
+    is `q_proj/k_proj/v_proj/out_proj`; the test asserts that
+    attention-only LoRA does NOT touch conv layers.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from peft import LoraConfig, get_peft_model
+from transformers import Lfm2Model
+
+# Canonical LFM2 LoRA target sets. The attention-only POC default targets only
+# the attention proj modules; the Spotify-full set (from v8 production) adds
+# conv `in_proj` and SwiGLU `w1/w2/w3`. Both sets include `out_proj`, which
+# matches both attention `out_proj` AND conv `out_proj` — see Lfm2Attention
+# vs Lfm2ShortConv. We constrain target_modules with the LAYER_PREFIX_REGEX
+# fallback below if attention-only-strict is required.
+ATTENTION_ONLY_TARGETS = ["q_proj", "k_proj", "v_proj"]  # + out_proj added below
+SPOTIFY_FULL_TARGETS = ["q_proj", "k_proj", "v_proj", "out_proj", "in_proj", "w1", "w2", "w3"]
+
+
+class LfmPseudoTokenBackbone(nn.Module):
+    """Wraps Lfm2Model to consume pre-computed pseudo-token embeddings.
+
+    Forward:
+        pseudo_tokens: (B, T, d_lfm) float — output of the projection adapter
+        attention_mask: (B, T) int or None — padding mask. For our use case
+            of fully-populated 64-tx sequences with no padding, None is the
+            normal call. Pass attention_mask explicitly if any pseudo-token
+            position is logically a "padding" slot.
+
+        Returns:
+            (B, T, d_lfm) float — `last_hidden_state` from the LFM2.5 stack
+            after 16 layers (10 conv + 6 attention).
+
+    Args:
+        model_path: HF-format directory containing config.json + safetensors.
+            For local-only POC we use ~/Projects/_models/LFM25-350M-Base; for
+            HF-hosted models pass a repo ID string.
+        lora: LoraConfig instance, or None for frozen-base-only (no LoRA).
+        dtype: bfloat16 for training (matches LFM2's pretraining dtype) or
+            float32 for CPU smoke tests.
+        device_map: "auto" for GPU runs, None for CPU/MPS local smoke runs.
+        trust_remote_code: required True for LiquidAI checkpoints that ship
+            modeling code alongside weights. Safe for LiquidAI/* paths.
+    """
+
+    def __init__(
+        self,
+        model_path: str | Path,
+        lora: LoraConfig | None = None,
+        dtype: torch.dtype = torch.bfloat16,
+        device_map: str | None = "auto",
+        trust_remote_code: bool = True,
+        freeze_base: bool = True,
+    ) -> None:
+        super().__init__()
+        load_kwargs: dict = {"torch_dtype": dtype, "trust_remote_code": trust_remote_code}
+        if device_map is not None:
+            load_kwargs["device_map"] = device_map
+        self.base = Lfm2Model.from_pretrained(str(model_path), **load_kwargs)
+
+        # Capture d_lfm before freezing for downstream consumers.
+        self.d_lfm: int = self.base.config.hidden_size
+
+        # Freeze controls. The default (freeze_base=True) is the "encoder
+        # pattern with a frozen base" — what makes the shared-base-across-
+        # customers story possible. Setting freeze_base=False is the
+        # stage-2 VL recipe ("unfreeze the base and fine-tune end-to-end"),
+        # used as a diagnostic upper-bound experiment.
+        if freeze_base:
+            self.freeze_base()
+
+        if lora is not None:
+            # get_peft_model wraps `self.base` and registers LoRA layers as
+            # trainable. The XOR guard inside Lfm2Model.forward still
+            # accepts inputs_embeds + input_ids=None — PEFT does not modify
+            # the forward signature, only the attention/MLP module weights.
+            self.base = get_peft_model(self.base, lora)
+
+    def freeze_base(self) -> None:
+        """Set requires_grad=False on every base parameter."""
+        for p in self.base.parameters():
+            p.requires_grad = False
+
+    def forward(
+        self,
+        pseudo_tokens: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # pseudo_tokens: (B, T, d_lfm) — must already be in the base's dtype
+        # (caller's responsibility; the projector should produce bf16 on GPU).
+        outputs = self.base(
+            input_ids=None,
+            inputs_embeds=pseudo_tokens,
+            attention_mask=attention_mask,
+            use_cache=False,
+        )
+        return outputs.last_hidden_state
+        # → (B, T, d_lfm)
+
+    def trainable_parameters(self) -> int:
+        """Number of trainable params (should be LoRA-only when LoRA is on)."""
+        return sum(p.numel() for p in self.base.parameters() if p.requires_grad)
+
+    def total_parameters(self) -> int:
+        return sum(p.numel() for p in self.base.parameters())
+
+
+# Regex matching attention modules only (strict mode). LFM2 attention modules
+# live under `layers.{i}.self_attn.{q_proj,k_proj,v_proj,out_proj}`; conv
+# `out_proj` lives under `layers.{i}.conv.out_proj`. The regex includes
+# `self_attn` in the path to disambiguate.
+STRICT_ATTENTION_REGEX = r".*self_attn\.(q_proj|k_proj|v_proj|out_proj)$"
+
+
+def build_lora_config(
+    r: int = 16,
+    alpha: int = 32,
+    dropout: float = 0.05,
+    target_modules: list[str] | str | None = None,
+    strict_attention_only: bool = False,
+) -> LoraConfig:
+    """Build a LoraConfig tuned for the LFM2 attention-LoRA POC default.
+
+    Default behavior (`target_modules=None`, `strict_attention_only=False`):
+        target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"]
+        PEFT matches by leaf-module name, so `out_proj` matches BOTH
+        attention.out_proj (6 occurrences) AND conv.out_proj (10 occurrences).
+        Net LoRA count for LFM2.5-350M: ~1.01M params.
+
+        We accept this conv.out_proj collision because:
+          - The conv layers do local sequence mixing — having them slightly
+            adaptable helps with our pseudo-token input distribution.
+          - It's 0.22M extra trainable params, dwarfed by the projector (~2.6M)
+            and downstream heads.
+          - Spotify production uses an even larger set (`SPOTIFY_FULL_TARGETS`).
+
+    To opt out of conv adaptation, pass `strict_attention_only=True` — this
+    selects target modules by regex so only attention paths match.
+
+    Args:
+        r: LoRA rank. 16 for POC, escalate to 64 for production.
+        alpha: LoRA scaling. Conventionally 2*r.
+        dropout: LoRA dropout. 0.05 keeps regularization light at POC scale.
+        target_modules: override the default. Pass `SPOTIFY_FULL_TARGETS`
+            (list) to escalate, or a regex string for custom matching.
+        strict_attention_only: if True and `target_modules` is None, use
+            the `STRICT_ATTENTION_REGEX` so only attention modules get LoRA.
+
+    Returns:
+        LoraConfig ready for `get_peft_model(base, config)`.
+    """
+    if target_modules is None:
+        if strict_attention_only:
+            target_modules = STRICT_ATTENTION_REGEX
+        else:
+            target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"]
+    # FEATURE_EXTRACTION task type — we use Lfm2Model (no LM head), and
+    # downstream task heads handle the loss. CAUSAL_LM would try to save
+    # an lm_head that does not exist on Lfm2Model.
+    return LoraConfig(
+        r=r,
+        lora_alpha=alpha,
+        lora_dropout=dropout,
+        target_modules=target_modules,
+        bias="none",
+        task_type="FEATURE_EXTRACTION",
+    )

encoder/src/model/transaction_encoder_nocompress.py

@@ -0,0 +1,118 @@
+"""Non-compressing transaction encoder.
+
+Mirrors the parent repo's `StructuredEmbedding` pattern exactly:
+per-feature value tables + feature-type table, summed, expanded to a
+flat sequence of length T_tx * F. The only difference is the output
+dimension — we project directly to `d_lfm` (1024 for LFM2.5-350M)
+instead of parent's `hidden_dim=256`.
+
+Why this exists:
+    The compressing encoder (TransactionEncoder) outputs 1 pseudo-token
+    per transaction (sequence length 64 from 64 transactions). That is
+    "encoder-style" — a small MLP collapses 15 features per tx into a
+    single 256-dim vector, then a projector lifts to 1024-dim. Fraud
+    quality is poor (ROC-AUC 0.53 at 100% labels) — likely because
+    intra-transaction feature combinations get averaged into the
+    256-dim bottleneck.
+
+    This non-compressing variant keeps the full 64*15 = 960-token
+    stream. No MLP-level compression. The LFM2.5 base sees the same
+    sequence shape parent's structured-feature backbone sees — but on
+    its own 1024-dim hidden space, with frozen text-pretrained weights
+    plus LoRA. Tests whether the compression was the binding constraint
+    on fraud quality.
+
+    What we GIVE UP relative to the compressing variant:
+    - The 15× sequence-length latency advantage
+    - The "modality-token" framing (now we're really doing per-feature-
+      token, like parent — the recipe is closer to "rebuild parent's
+      embedding layer on a bigger backbone" than to LFM2-VL).
+
+    What we keep:
+    - Frozen base (cross-customer shared base story)
+    - LoRA on attention
+    - The pretrained text backbone
+
+Shape contract:
+    (B, T_tx, F) int64  →  (B, T_tx * F, d_lfm) float
+    For default (T_tx=64, F=15): (B, 64, 15) → (B, 960, 1024)
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from src.data.schema import SchemaConfig
+
+
+class StructuredEncoder(nn.Module):
+    """Per-feature value embeddings + feature-type embeddings, summed.
+
+    This is the parent repo's `StructuredEmbedding` at `hidden_dim=d_lfm`.
+    Drop-in input-side replacement for the LFM2's `embed_tokens` table —
+    same shape contract.
+
+    Args:
+        schema: parent's SchemaConfig.
+        d_lfm: LFM2 backbone hidden size. 1024 for LFM2.5-350M, 2048 for
+            LFM2.5-1.2B. Must match the LFM wrapper's d_lfm.
+
+    Forward:
+        token_ids: (B, T_tx, F) int64
+        returns:   (B, T_tx * F, d_lfm) float, dtype matching the
+                   embedding tables (fp32 by default).
+    """
+
+    def __init__(self, schema: SchemaConfig, d_lfm: int) -> None:
+        super().__init__()
+        self.num_features = schema.num_features
+        self.num_transactions = schema.num_transactions
+        self.d_lfm = d_lfm
+
+        # Per-feature value tables, each sized to that feature's full vocab.
+        # No padding_idx — reserved tokens (MASK/OOV/NULL) get learned
+        # embeddings like everything else. Total params dominated by
+        # merchant_id (10003 * d_lfm = 10.2M for d_lfm=1024).
+        self.value_tables = nn.ModuleList(
+            [nn.Embedding(f.vocab_size, d_lfm) for f in schema.features],
+        )
+
+        # Feature-type table. Distinguishes "the 5th token is an mcc"
+        # from "the 5th token is a merchant_id" — RoPE alone can't carry
+        # this since intra-transaction position is arbitrary. Mirrors
+        # parent's pattern exactly.
+        self.type_table = nn.Embedding(schema.num_features, d_lfm)
+
+    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+        # token_ids: (B, T_tx, F) int64
+        B, T_tx, F = token_ids.shape
+        assert F == self.num_features, (
+            f"Expected {self.num_features} features, got {F}"
+        )
+
+        # Type embeddings: (F, d_lfm). Same per-feature offsets added to
+        # every transaction's value embedding for that feature.
+        type_indices = torch.arange(F, device=token_ids.device)
+        type_emb = self.type_table(type_indices)
+        # type_emb: (F, d_lfm)
+
+        # Embed each feature column with its own table, add the type
+        # offset, stack into (B, T_tx, F, d_lfm).
+        feature_embeddings: list[torch.Tensor] = []
+        for f_idx in range(F):
+            feat_tokens = token_ids[:, :, f_idx]
+            val_emb = self.value_tables[f_idx](feat_tokens)
+            # val_emb: (B, T_tx, d_lfm)
+            feature_embeddings.append(val_emb + type_emb[f_idx])
+
+        stacked = torch.stack(feature_embeddings, dim=2)
+        # stacked: (B, T_tx, F, d_lfm)
+        return stacked.reshape(B, T_tx * F, self.d_lfm)
+        # → (B, T_tx * F, d_lfm) — flat sequence the LFM can consume
+
+    def num_embedding_params(self) -> int:
+        """Total params in value + type tables (sanity check)."""
+        val = sum(e.weight.numel() for e in self.value_tables)
+        typ = self.type_table.weight.numel()
+        return val + typ

encoder/src/model/transaction_fm.py

@@ -0,0 +1,268 @@
+"""Orchestrator: encoder + projector + LFM2.5 wrapper → unified backbone.
+
+This module produces a `MultiHeadModel`-compatible object so the parent
+repo's `validate()`, `compute_losses`, and downstream-head infrastructure
+work without any rewriting. The encoder approach plugs into the parent
+harness as a different "backbone" — same downstream-head contract.
+
+The naming dance:
+- `EncoderBackbone` is what `MultiHeadModel.backbone` expects. It exposes
+  `backbone_forward(token_ids) → (B, T, d_lfm)`.
+- `build_transaction_fm()` constructs the full
+  `MultiHeadModel(backbone=EncoderBackbone, heads=…)` stack from configs.
+
+Head pool semantics with T=64, num_features=1:
+- `last_tx_mean` (num_features=1) ⇒ `hidden[:, -1:, :].mean(1)` = last
+  pseudo-token = "tx 63's hidden state". Used by fraud (sequence-level).
+- `pre_last_tx` (num_features=1) ⇒ `hidden[:, -2, :]` = "tx 62's hidden
+  state". Used by next_merchant / amount_range / mcc — predicts tx 63's
+  features from the tx 62 representation (which has causally seen tx 0..62).
+
+The semantics line up with the parent's `T=960, num_features=15` layout:
+- Parent last_tx_mean reads positions 945..959 (tx 63's 15 features).
+- Parent pre_last_tx reads position 944 (end of tx 62).
+Encoder version reads the same logical positions at compressed sequence
+length.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from src.data.schema import SchemaConfig
+from src.model.task_heads import (
+    AnyHead,
+    DownstreamHead,
+    HeadConfig,
+    MultiHeadModel,
+)
+from encoder.src.model.encoder_heads import EncoderDownstreamHead, EncoderTiedEmbeddingHead
+from encoder.src.model.lfm_pseudo_token_wrapper import (
+    LfmPseudoTokenBackbone,
+    build_lora_config,
+)
+from encoder.src.model.projection_adapter import ProjectionAdapter
+from encoder.src.model.transaction_encoder import TransactionEncoder
+from encoder.src.model.transaction_encoder_nocompress import StructuredEncoder
+
+
+class EncoderBackbone(nn.Module):
+    """Encoder (+ optional projector) + LFM2.5 wrapper as a single backbone.
+
+    Implements the `backbone_forward(token_ids: Tensor) -> Tensor` contract
+    that the parent's `MultiHeadModel` expects.
+
+    Two modes, dispatched by the encoder type:
+      - **Compress** (TransactionEncoder, default): encoder outputs (B, 64,
+        d_encoder=256). Projector lifts to (B, 64, d_lfm=1024). Sequence
+        length 64.
+      - **Non-compress** (StructuredEncoder): encoder outputs (B, 960,
+        d_lfm) directly. No projector. Sequence length 960. Matches the
+        parent's structured-feature input shape.
+
+    The projector is None in non-compress mode. The branch is selected at
+    construction time by `build_transaction_fm`.
+    """
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        projector: ProjectionAdapter | None,
+        lfm_wrapper: LfmPseudoTokenBackbone,
+    ) -> None:
+        super().__init__()
+        self.encoder = encoder
+        # nn.Module attribute, may be None in non-compress mode.
+        self.projector = projector
+        self.lfm = lfm_wrapper
+        self.d_lfm = lfm_wrapper.d_lfm
+
+    def backbone_forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+        # token_ids: (B, 64, 15) int64
+        # encoder output shape depends on mode:
+        #   compress    → (B, 64, d_encoder)
+        #   nocompress  → (B, 960, d_lfm)
+        x = self.encoder(token_ids)
+
+        if self.projector is not None:
+            # Compress mode: lift d_encoder → d_lfm.
+            x = self.projector(x)
+        # x: (B, T, d_lfm) where T = 64 (compress) or 960 (nocompress)
+
+        # Cast to the LFM base's dtype before injection. The LFM was loaded
+        # in bf16 (GPU) or fp32 (CPU); upstream modules produce whatever
+        # they want (default fp32). This cast bridges the boundary.
+        target_dtype = next(self.lfm.base.parameters()).dtype
+        if x.dtype != target_dtype:
+            x = x.to(target_dtype)
+
+        hidden = self.lfm(x)
+        # → (B, T, d_lfm) in target_dtype
+        return hidden
+
+    # `MultiHeadModel.forward` calls `self.backbone.backbone_forward`, but
+    # for debugging / standalone use we also expose `forward` as an alias.
+    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.backbone_forward(token_ids)
+
+
+def build_heads_encoder(
+    head_configs: dict[str, dict[str, Any]],
+    hidden_dim: int,
+    num_features: int,
+    value_tables: nn.ModuleList | None = None,
+) -> dict[str, AnyHead]:
+    """Instantiate downstream heads from a config dict.
+
+    `num_features` is the per-tx token count and shapes the pool strategy:
+      - **compress mode**: `num_features=1` (one pseudo-token per tx)
+        → `last_tx_mean` pools position -1, `pre_last_tx` pools position -2.
+      - **nocompress mode**: `num_features=15` (matches parent's structured
+        layout) → `last_tx_mean` pools last 15 positions, `pre_last_tx`
+        pools position -(15+1) = -16 = end of tx 62.
+
+    Two head implementations are dispatched:
+      - **EncoderDownstreamHead** (default): fresh MLP from pool → output.
+      - **EncoderTiedEmbeddingHead**: shares its classifier matrix with one
+        of the encoder's per-feature value tables. Selected by
+        `tied: true` on a head's config. Requires `value_tables` to be
+        passed (only available in `nocompress` mode where the encoder is a
+        `StructuredEncoder`); `compress` mode's TransactionEncoder uses
+        d_feat=32 feature embeddings that can't be tied at d_lfm=1024.
+
+    Heads with `tied: true` MUST have `target_type: "feature:N"` matching
+    the value table they share weights with. The head's `output_dim` is
+    ignored in tied mode (implicitly the feature's vocab_size).
+    """
+    heads: dict[str, AnyHead] = {}
+    for name, hcfg in head_configs.items():
+        hc = HeadConfig(
+            name=name,
+            output_dim=hcfg["output_dim"],
+            loss_type=hcfg["loss"],
+            pool_strategy=hcfg["pool"],
+            target_type=hcfg["target"],
+            weight=hcfg.get("weight", 1.0),
+            mlp_hidden=hcfg.get("mlp_hidden", 128),
+            dropout=hcfg.get("dropout", 0.1),
+        )
+        if hcfg.get("tied", False):
+            if value_tables is None:
+                raise ValueError(
+                    f"Head '{name}' has tied=true but no value_tables were "
+                    f"passed. Tied heads are only supported in nocompress "
+                    f"mode (where the encoder is StructuredEncoder).",
+                )
+            heads[name] = EncoderTiedEmbeddingHead(
+                hc, hidden_dim, num_features=num_features,
+                value_tables=value_tables,
+            )
+        else:
+            heads[name] = EncoderDownstreamHead(
+                hc, hidden_dim, num_features=num_features,
+            )
+    return heads
+
+
+def build_transaction_fm(
+    schema: SchemaConfig,
+    head_configs: dict[str, dict[str, Any]],
+    model_path: str | Path,
+    architecture_cfg: dict[str, Any] | None = None,
+    encoder_cfg: dict[str, Any] | None = None,
+    projector_cfg: dict[str, Any] | None = None,
+    lora_cfg: dict[str, Any] | None = None,
+    dtype: torch.dtype = torch.bfloat16,
+    device_map: str | None = "auto",
+) -> MultiHeadModel:
+    """Construct the full encoder + LFM + heads stack.
+
+    The architecture mode is selected by
+    `architecture_cfg["mode"] ∈ {"compress", "nocompress"}`:
+
+      - **compress** (default): TransactionEncoder → ProjectionAdapter →
+        LFM. One pseudo-token per transaction; sequence length 64.
+      - **nocompress**: StructuredEncoder → LFM directly. 15 pseudo-tokens
+        per transaction (one per feature); sequence length 960. Mirrors
+        parent's structured-feature input shape.
+
+    Returns a `MultiHeadModel` so it's drop-in compatible with the parent's
+    `validate()`, `compute_losses`, and training utilities regardless of mode.
+    """
+    architecture_cfg = architecture_cfg or {}
+    encoder_cfg = encoder_cfg or {}
+    projector_cfg = projector_cfg or {}
+    lora_cfg = lora_cfg or {}
+    mode = architecture_cfg.get("mode", "compress")
+
+    # Build the LFM wrapper first so we know d_lfm before constructing the
+    # encoder/projector. (Compress mode needs d_lfm for the projector;
+    # nocompress mode needs it for the StructuredEncoder.)
+    lora = None
+    if lora_cfg.get("enabled", True):
+        # Allow config to override the default target module list. Passing
+        # an explicit list (e.g. ["q_proj", ..., "in_proj"]) opts into
+        # conv-LoRA. None falls through to build_lora_config's default
+        # (attention names with conv.out_proj collision).
+        lora = build_lora_config(
+            r=lora_cfg.get("r", 16),
+            alpha=lora_cfg.get("alpha", 32),
+            dropout=lora_cfg.get("dropout", 0.05),
+            target_modules=lora_cfg.get("target_modules"),
+            strict_attention_only=lora_cfg.get("strict_attention_only", False),
+        )
+    lfm_wrapper = LfmPseudoTokenBackbone(
+        model_path,
+        lora=lora,
+        dtype=dtype,
+        device_map=device_map,
+        # `freeze_base=True` is the encoder-pattern default. The stage-2
+        # diagnostic (architecture.unfreeze_backbone=true) flips this so
+        # the full 354M backbone is end-to-end trainable.
+        freeze_base=not architecture_cfg.get("unfreeze_backbone", False),
+    )
+
+    if mode == "compress":
+        encoder = TransactionEncoder(
+            schema,
+            d_feat=encoder_cfg.get("d_feat", 32),
+            d_encoder=encoder_cfg.get("d_encoder", 256),
+            mlp_hidden=encoder_cfg.get("mlp_hidden", 384),
+        )
+        projector: ProjectionAdapter | None = ProjectionAdapter(
+            d_encoder=encoder_cfg.get("d_encoder", 256),
+            d_lfm=lfm_wrapper.d_lfm,
+            hidden=projector_cfg.get("hidden", 2 * lfm_wrapper.d_lfm),
+            use_layernorm=projector_cfg.get("use_layernorm", True),
+        )
+        num_features_for_heads = 1
+    elif mode == "nocompress":
+        encoder = StructuredEncoder(schema, d_lfm=lfm_wrapper.d_lfm)
+        projector = None
+        # Matches parent: heads pool over 15-feature transaction stripes.
+        num_features_for_heads = schema.num_features
+    else:
+        raise ValueError(
+            f"Unknown architecture mode: {mode!r}. Expected 'compress' or 'nocompress'.",
+        )
+
+    backbone = EncoderBackbone(encoder, projector, lfm_wrapper)
+    # Tied heads need access to the encoder's per-feature value tables.
+    # Only StructuredEncoder (nocompress mode) exposes them at d_lfm; the
+    # compress mode's TransactionEncoder has d_feat=32 feature embeddings
+    # that can't share weights with a d_lfm-dim classifier.
+    value_tables_for_heads = (
+        encoder.value_tables if mode == "nocompress" else None
+    )
+    heads = build_heads_encoder(
+        head_configs,
+        hidden_dim=lfm_wrapper.d_lfm,
+        num_features=num_features_for_heads,
+        value_tables=value_tables_for_heads,
+    )
+
+    return MultiHeadModel(backbone=backbone, heads=heads)

requirements.txt

@@ -5,3 +5,4 @@ peft>=0.13.0
 accelerate>=1.0.0
 numpy>=1.26
 pyyaml>=6.0
+scikit-learn>=1.5

src/data/generator.py

@@ -0,0 +1,1022 @@
+"""Synthetic transaction data generator.
+
+Design features:
+  1. Merchant catalog: 10K merchants with fixed MCC assignments and Zipf
+     popularity. Creates realistic merchant-MCC correlation the model can learn.
+  2. MCC-conditional amounts: transaction amounts are sampled from MCC-specific
+     distributions. Restaurants produce $15-80, electronics $200-2000, etc.
+  3. Profile blending: each customer is a mixture of primary + secondary
+     profile, not a pure archetype. Continuous behavioral diversity.
+  4. MCC-conditional entry_mode: online merchants -> CNP, restaurants -> tap/chip.
+  5. MCC-conditional hour: restaurants peak at lunch/dinner, transit at commute.
+  6. MCC-conditional is_recurring: subscriptions/streaming are high, restaurants low.
+  7. Customer_tenure-card_product correlation: new customers get basic cards.
+  8. Country-amount correlation: international transactions tend higher value.
+
+Output format: token_ids.npy (N,T,F), sequence_labels.npy,
+transaction_labels.npy, split_indices.npz, fingerprint.txt. Drives the
+training pipeline via schema.yaml and model.yaml.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import yaml
+
+from src.data.schema import SchemaConfig, load_schema, NULL_TOKEN, VALUES_START
+from src.data.tokenizer import TransactionTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Amount range bins: 16 coarse ranges over the 256 fine-grained amount buckets.
+# Each range covers 16 consecutive buckets. Used as the prediction target for
+# the amount_range fine-tuning head (replaces the 259-class amount head).
+# Dollar labels are approximate (quantile buckets are non-linear).
+# ---------------------------------------------------------------------------
+
+NUM_AMOUNT_RANGES: int = 16
+AMOUNT_RANGE_WIDTH: int = 16    # 256 buckets / 16 ranges
+
+AMOUNT_RANGE_LABELS: dict[int, str] = {
+    0: "$0-5",
+    1: "$5-15",
+    2: "$15-30",
+    3: "$30-50",
+    4: "$50-80",
+    5: "$80-120",
+    6: "$120-175",
+    7: "$175-250",
+    8: "$250-375",
+    9: "$375-550",
+    10: "$550-800",
+    11: "$800-1.2K",
+    12: "$1.2K-2K",
+    13: "$2K-3.5K",
+    14: "$3.5K-7.5K",
+    15: "$7.5K+",
+}
+
+
+def amount_bucket_to_range(bucket: np.ndarray) -> np.ndarray:
+    """Map 256 fine-grained amount buckets (0-255) to 16 coarse ranges (0-15)."""
+    return np.clip(bucket // AMOUNT_RANGE_WIDTH, 0, NUM_AMOUNT_RANGES - 1)
+
+
+# ---------------------------------------------------------------------------
+# MCC-conditional amount distributions (mean_bucket, std_bucket).
+# 256 amount buckets spanning $0.01-$25K in quantile space. These bucket-index
+# ranges are calibrated for behavioral plausibility across MCC categories.
+# ---------------------------------------------------------------------------
+
+MCC_AMOUNT_MAP: dict[tuple[int, int], tuple[float, float]] = {
+    (0, 3): (25.0, 12.0),       # transit, coffee, gas: $3-20
+    (4, 9): (45.0, 20.0),       # convenience, misc low: $8-50
+    (10, 14): (160.0, 55.0),    # airlines, hotels, car rental: $80-800
+    (15, 19): (50.0, 30.0),     # other transport: $10-100
+    (20, 24): (80.0, 45.0),     # online retail general: $15-200
+    (25, 29): (55.0, 30.0),     # online retail niche: $10-120
+    (30, 34): (55.0, 25.0),     # restaurants, bars: $15-100
+    (35, 39): (40.0, 20.0),     # other food service: $8-60
+    (40, 44): (185.0, 55.0),    # luxury, dept stores, jewelry: $80-2000
+    (45, 49): (70.0, 35.0),     # fashion retail: $20-200
+    (50, 54): (110.0, 50.0),    # office supplies, wholesale: $30-500
+    (55, 59): (90.0, 50.0),     # professional services: $30-400
+    (60, 64): (50.0, 22.0),     # grocery, pharmacy: $12-120
+    (65, 69): (55.0, 30.0),     # healthcare, fitness: $15-150
+    (70, 74): (22.0, 12.0),     # fast food, entertainment: $4-25
+    (75, 79): (35.0, 20.0),     # streaming, digital goods: $5-50
+    (80, 84): (30.0, 22.0),     # subscriptions, SaaS: $5-80
+    (85, 89): (100.0, 55.0),    # automotive: $25-500
+    (90, 94): (95.0, 55.0),     # home improvement: $20-600
+    (95, 99): (75.0, 50.0),     # education, government: $15-500
+}
+
+
+def _mcc_to_amount_params(mcc_value: int) -> tuple[float, float]:
+    """Look up amount distribution params for an MCC value."""
+    for (lo, hi), params in MCC_AMOUNT_MAP.items():
+        if lo <= mcc_value <= hi:
+            return params
+    return (60.0, 40.0)
+
+
+# ---------------------------------------------------------------------------
+# MCC-conditional entry_mode weights.
+# Indices: 0=card_present, 1=card_not_present, 2=contactless, 3=chip, 4=manual
+# ---------------------------------------------------------------------------
+
+MCC_ENTRY_MODE: dict[tuple[int, int], list[float]] = {
+    (0, 3): [2, 1, 5, 2, 0],       # transit/coffee: contactless heavy
+    (10, 14): [1, 6, 1, 1, 0],     # travel: mostly online bookings
+    (20, 29): [0, 8, 0, 0, 1],     # online retail: card-not-present
+    (30, 39): [3, 0, 4, 3, 0],     # restaurants: contactless/chip
+    (40, 49): [3, 2, 3, 2, 0],     # fashion/luxury: mixed
+    (50, 59): [2, 4, 1, 1, 1],     # business: mix of online + in-person
+    (60, 69): [3, 1, 4, 3, 0],     # grocery/pharmacy: in-person
+    (70, 74): [2, 1, 5, 2, 0],     # fast food: contactless
+    (75, 84): [0, 9, 0, 0, 0],     # digital/subscriptions: all online
+    (85, 99): [3, 2, 2, 3, 0],     # auto/home/edu: mixed
+}
+
+
+def _mcc_to_entry_mode_weights(mcc_value: int) -> list[float]:
+    """Look up entry_mode distribution for an MCC value."""
+    for (lo, hi), weights in MCC_ENTRY_MODE.items():
+        if lo <= mcc_value <= hi:
+            return weights
+    return [2, 3, 2, 2, 1]
+
+
+# ---------------------------------------------------------------------------
+# MCC-conditional hour-of-day distributions.
+# Peak hours where a category has higher probability. Applied as a 50% blend
+# with the profile-sampled hour to keep per-customer variation.
+# Format: (peak_hours, peak_weight) -- applied on a uniform-24 base.
+# ---------------------------------------------------------------------------
+
+MCC_HOUR_PEAKS: dict[tuple[int, int], tuple[list[int], float]] = {
+    (0, 3): ([7, 8, 9, 17, 18], 5.0),       # transit/coffee: commute hours
+    (10, 14): ([9, 10, 14, 15, 20], 2.5),    # travel: business hours + evening
+    (20, 29): ([10, 11, 20, 21, 22], 3.0),   # online: late morning + evening
+    (30, 31): ([12, 13, 18, 19, 20], 6.0),   # casual dining: lunch + dinner
+    (32, 34): ([19, 20, 21, 22, 23], 5.0),   # fine dining/bars: evening + late
+    (35, 39): ([11, 12, 13, 17, 18], 4.0),   # other food: lunch + early dinner
+    (60, 64): ([10, 11, 16, 17, 18], 4.0),   # grocery/pharmacy: mid-morning + after work
+    (70, 71): ([12, 13, 18, 19, 21], 4.5),   # fast food: lunch + dinner + late
+    (75, 84): ([0, 1, 6, 7], 2.0),           # digital/subs: auto-renew any hour
+}
+
+
+# ---------------------------------------------------------------------------
+# MCC-conditional is_recurring probability.
+# ---------------------------------------------------------------------------
+
+MCC_RECURRING_RATE: dict[tuple[int, int], float] = {
+    (0, 3): 0.20,       # transit passes
+    (10, 14): 0.05,     # travel: rarely recurring
+    (20, 29): 0.12,     # online retail: occasional subscriptions
+    (30, 39): 0.03,     # restaurants: very rarely recurring
+    (40, 49): 0.02,     # luxury/fashion: almost never
+    (50, 54): 0.15,     # office supplies: recurring orders
+    (60, 64): 0.08,     # grocery: delivery subscriptions
+    (65, 69): 0.25,     # fitness/health: gym memberships
+    (70, 74): 0.05,     # fast food: rarely
+    (75, 79): 0.60,     # streaming: almost always recurring
+    (80, 84): 0.70,     # SaaS/subscriptions: by definition
+    (85, 94): 0.05,     # auto/home: rarely
+    (95, 99): 0.10,     # education/gov: tuition installments
+}
+
+
+# ---------------------------------------------------------------------------
+# Customer tenure -> card_product correlation.
+# New customers (low tenure bucket) skew toward basic cards.
+# Established customers (high tenure bucket) skew toward premium/rewards.
+# Format: tenure_bucket_threshold -> card_product weights (10 products).
+# ---------------------------------------------------------------------------
+
+TENURE_CARD_PRODUCT: list[tuple[int, list[float]]] = [
+    (2, [6, 3, 1, 0, 0, 0, 0, 2, 1, 0]),    # 0-2: new -> basic/prepaid/virtual
+    (5, [3, 4, 3, 1, 0, 1, 1, 1, 1, 0]),     # 3-5: moderate -> basic/rewards
+    (7, [2, 3, 4, 2, 1, 1, 2, 0, 1, 1]),     # 6-7: established -> rewards/business
+    (99, [1, 2, 3, 3, 2, 1, 2, 0, 0, 1]),    # 8+: long-tenure -> premium/platinum
+]
+
+
+# ---------------------------------------------------------------------------
+# Merchant catalog
+# ---------------------------------------------------------------------------
+
+@dataclass
+class MerchantCatalog:
+    """10K merchants with MCC, popularity, and characteristic amount range."""
+
+    merchant_ids: np.ndarray       # (num_merchants,) int32
+    mcc_assignments: np.ndarray    # (num_merchants,) int32
+    popularity: np.ndarray         # (num_merchants,) float64, sums to 1
+    amount_mean: np.ndarray        # (num_merchants,) float64, amount bucket mean
+    amount_std: np.ndarray         # (num_merchants,) float64, amount bucket std
+    mcc_to_merchants: dict[int, np.ndarray]  # MCC -> array of merchant indices
+
+    @staticmethod
+    def build(
+        num_merchants: int,
+        num_mccs: int,
+        rng: np.random.Generator,
+        zipf_exponent: float = 1.2,
+    ) -> MerchantCatalog:
+        """Build a merchant catalog with MCC assignments, Zipf popularity,
+        and per-merchant amount distributions.
+
+        Each merchant gets a characteristic amount (mean, std) derived from
+        its MCC's range with per-merchant variation. Std is kept tight (4-8
+        buckets) so the model can learn merchant -> amount.
+        """
+        merchant_ids = np.arange(num_merchants, dtype=np.int32)
+
+        mcc_merchant_counts = _sample_mcc_merchant_counts(
+            num_merchants, num_mccs, rng,
+        )
+
+        mcc_assignments = np.zeros(num_merchants, dtype=np.int32)
+        offset = 0
+        for mcc, count in enumerate(mcc_merchant_counts):
+            mcc_assignments[offset : offset + count] = mcc
+            offset += count
+
+        rng.shuffle(mcc_assignments)
+
+        raw_pop = 1.0 / np.arange(1, num_merchants + 1, dtype=np.float64) ** zipf_exponent
+        rng.shuffle(raw_pop)
+        popularity = raw_pop / raw_pop.sum()
+
+        amount_mean = np.zeros(num_merchants, dtype=np.float64)
+        amount_std = np.zeros(num_merchants, dtype=np.float64)
+        for i in range(num_merchants):
+            mcc_mean, mcc_std = _mcc_to_amount_params(int(mcc_assignments[i]))
+            amount_mean[i] = np.clip(
+                mcc_mean + rng.normal(0, mcc_std * 0.3), 1.0, 254.0,
+            )
+            amount_std[i] = np.clip(rng.uniform(3.0, 7.0), 2.0, 10.0)
+
+        mcc_to_merchants: dict[int, np.ndarray] = {}
+        for mcc in range(num_mccs):
+            mask = mcc_assignments == mcc
+            mcc_to_merchants[mcc] = np.where(mask)[0].astype(np.int32)
+
+        return MerchantCatalog(
+            merchant_ids=merchant_ids,
+            mcc_assignments=mcc_assignments,
+            popularity=popularity,
+            amount_mean=amount_mean,
+            amount_std=amount_std,
+            mcc_to_merchants=mcc_to_merchants,
+        )
+
+    def sample_merchants_for_mccs(
+        self,
+        mcc_values: np.ndarray,
+        rng: np.random.Generator,
+    ) -> np.ndarray:
+        """Given MCC values (any shape), return merchant IDs from that MCC.
+
+        Falls back to popularity-weighted global sampling when an MCC has
+        no assigned merchants (shouldn't happen with proper catalog).
+        """
+        flat_mccs = mcc_values.ravel()
+        result = np.zeros(len(flat_mccs), dtype=np.int32)
+
+        unique_mccs = np.unique(flat_mccs)
+        for mcc in unique_mccs:
+            positions = np.where(flat_mccs == mcc)[0]
+            merchants = self.mcc_to_merchants.get(int(mcc))
+            if merchants is None or len(merchants) == 0:
+                result[positions] = rng.integers(0, len(self.merchant_ids), size=len(positions))
+                continue
+            pop = self.popularity[merchants]
+            pop = pop / pop.sum()
+            chosen = rng.choice(merchants, size=len(positions), p=pop)
+            result[positions] = chosen
+
+        return result.reshape(mcc_values.shape)
+
+
+def _sample_mcc_merchant_counts(
+    num_merchants: int, num_mccs: int, rng: np.random.Generator,
+) -> np.ndarray:
+    """Distribute merchants across MCCs with realistic category sizes.
+
+    Restaurants/retail get more merchants (fragmented industries).
+    Airlines/utilities get fewer (concentrated industries).
+    """
+    base_weights = np.ones(num_mccs, dtype=np.float64)
+
+    category_weights = {
+        (0, 3): 0.6,       # transit/coffee: moderate
+        (4, 9): 0.8,
+        (10, 14): 0.3,     # airlines/hotels: concentrated
+        (15, 19): 0.5,
+        (20, 24): 1.5,     # online retail: many merchants
+        (25, 29): 1.0,
+        (30, 34): 2.0,     # restaurants: very fragmented
+        (35, 39): 1.2,
+        (40, 44): 0.5,     # luxury/dept: fewer
+        (45, 49): 1.3,     # fashion: many
+        (50, 54): 0.8,
+        (55, 59): 0.7,
+        (60, 64): 1.0,     # grocery: moderate
+        (65, 69): 0.8,
+        (70, 74): 1.5,     # fast food: many
+        (75, 79): 0.6,
+        (80, 84): 0.5,     # subscriptions: concentrated
+        (85, 89): 0.7,
+        (90, 94): 0.9,
+        (95, 99): 0.4,     # education/gov: concentrated
+    }
+
+    for (lo, hi), w in category_weights.items():
+        lo_clamped = min(lo, num_mccs - 1)
+        hi_clamped = min(hi, num_mccs - 1)
+        base_weights[lo_clamped : hi_clamped + 1] = w
+
+    noise = rng.uniform(0.8, 1.2, size=num_mccs)
+    weights = base_weights * noise
+    weights /= weights.sum()
+
+    counts = (weights * num_merchants).astype(np.int32)
+    remainder = num_merchants - counts.sum()
+    if remainder > 0:
+        top_idx = np.argsort(weights)[-remainder:]
+        counts[top_idx] += 1
+    elif remainder < 0:
+        top_idx = np.argsort(counts)[remainder:]
+        counts[top_idx] -= 1
+
+    assert counts.sum() == num_merchants
+    return counts
+
+
+# ---------------------------------------------------------------------------
+# Generated dataset
+# ---------------------------------------------------------------------------
+
+@dataclass
+class GeneratedDataset:
+    """Output of the generation pipeline."""
+
+    token_ids: np.ndarray
+    sequence_labels: np.ndarray
+    transaction_labels: np.ndarray
+    amount_range_labels: np.ndarray
+    split_indices: dict[str, np.ndarray]
+    fingerprint: str
+
+
+# ---------------------------------------------------------------------------
+# Generator
+# ---------------------------------------------------------------------------
+
+class DataGenerator:
+    """Enhanced data generator with merchant catalog and correlated features."""
+
+    def __init__(
+        self,
+        schema_path: str | Path = "data/schema.yaml",
+        profiles_path: str | Path = "data/profiles.yaml",
+        fraud_path: str | Path = "data/fraud_patterns.yaml",
+        null_rates_path: str | Path = "data/null_rates.yaml",
+        seed: int = 42,
+        blend_fraction: float = 0.25,
+    ) -> None:
+        self.schema = load_schema(schema_path)
+        self._profiles_raw = self._load_yaml(profiles_path)
+        self._fraud_raw = self._load_yaml(fraud_path)
+        self._null_raw = self._load_yaml(null_rates_path)
+        self._config_paths = [
+            Path(p) for p in [schema_path, profiles_path, fraud_path, null_rates_path]
+        ]
+        self.rng = np.random.default_rng(seed)
+        self._seed = seed
+        self._blend_fraction = blend_fraction
+
+        em = self.schema.get_feature("entry_mode")
+        self._entry_mode_name_to_val: dict[str, int] = {}
+        if em.values:
+            self._entry_mode_name_to_val = {v: k for k, v in em.values.items()}
+
+        n_merchants = self.schema.get_feature("merchant_id").num_values
+        n_mccs = self.schema.get_feature("mcc").num_values
+        self.catalog = MerchantCatalog.build(n_merchants, n_mccs, self.rng)
+        logger.info(
+            "Built merchant catalog: %d merchants across %d MCCs",
+            n_merchants, n_mccs,
+        )
+
+    @staticmethod
+    def _load_yaml(path: str | Path) -> dict[str, Any]:
+        with open(path) as fh:
+            return yaml.safe_load(fh)
+
+    # --- Distribution sampling ---
+
+    def _sample_feature(
+        self, config: dict[str, Any], num_values: int, shape: tuple[int, int],
+    ) -> np.ndarray:
+        dist = config["distribution"]
+
+        if dist == "uniform":
+            return self.rng.integers(0, num_values, size=shape, dtype=np.int32)
+
+        if dist == "weights":
+            w = np.array(config["weights"], dtype=np.float64)
+            w /= w.sum()
+            return self.rng.choice(len(w), size=shape, p=w).astype(np.int32)
+
+        if dist == "peaks":
+            w = np.ones(num_values, dtype=np.float64)
+            for p in config["peaks"]:
+                if 0 <= p < num_values:
+                    w[p] = config["peak_weight"]
+            w /= w.sum()
+            return self.rng.choice(num_values, size=shape, p=w).astype(np.int32)
+
+        if dist == "normal":
+            mean = float(config["mean_bucket"])
+            std = max(float(config["std_bucket"]), 0.1)
+            vals = self.rng.normal(mean, std, size=shape)
+            return np.clip(np.round(vals), 0, num_values - 1).astype(np.int32)
+
+        if dist == "bernoulli":
+            return (self.rng.random(size=shape) < config["p_true"]).astype(np.int32)
+
+        if dist == "concentrated":
+            return self._sample_concentrated(config, num_values, shape)
+
+        raise ValueError(f"Unknown distribution type: {dist}")
+
+    def _sample_concentrated(
+        self, config: dict[str, Any], num_values: int, shape: tuple[int, int],
+    ) -> np.ndarray:
+        """Sample from a per-customer concentrated distribution with Zipf skew.
+
+        Each customer has a fixed set of preferred values. Within the preferred
+        set, frequency follows Zipf: the top item gets ~30% of traffic, the
+        second ~18%, etc. This creates a learnable frequency ranking the model
+        can exploit from sequence history.
+        """
+        top_n = min(config["top_n"], num_values)
+        concentration = config["concentration"]
+        zipf_exp = config.get("zipf_exponent", 1.0)
+        n_seq, n_tx = shape
+
+        preferred = np.stack([
+            self.rng.choice(num_values, size=top_n, replace=False)
+            for _ in range(n_seq)
+        ])
+
+        zipf_weights = 1.0 / np.arange(1, top_n + 1, dtype=np.float64) ** zipf_exp
+        zipf_weights /= zipf_weights.sum()
+
+        use_preferred = self.rng.random(shape) < concentration
+        pref_idx = self.rng.choice(top_n, size=shape, p=zipf_weights).astype(np.int32)
+        pref_values = preferred[np.arange(n_seq)[:, None], pref_idx]
+        unif_values = self.rng.integers(0, num_values, size=shape, dtype=np.int32)
+
+        return np.where(use_preferred, pref_values, unif_values).astype(np.int32)
+
+    # --- Core generation ---
+
+    def generate(self, num_sequences: int = 100_000) -> GeneratedDataset:
+        """Generate the full dataset with correlated features."""
+        N = num_sequences
+        T = self.schema.num_transactions
+        F = self.schema.num_features
+
+        raw_values = np.zeros((N, T, F), dtype=np.int32)
+        profiles = self._profiles_raw["profiles"]
+        profile_weights = np.array([p["weight"] for p in profiles])
+        profile_weights /= profile_weights.sum()
+
+        primary_indices = self.rng.choice(len(profiles), size=N, p=profile_weights)
+        secondary_indices = self.rng.choice(len(profiles), size=N, p=profile_weights)
+        blend_mask = self.rng.random((N, T)) < self._blend_fraction
+
+        mcc_idx = self.schema.feature_index("mcc")
+        merchant_idx = self.schema.feature_index("merchant_id")
+        amount_idx = self.schema.feature_index("amount")
+        entry_mode_idx = self.schema.feature_index("entry_mode")
+
+        for prof_idx, profile in enumerate(profiles):
+            primary_mask = primary_indices == prof_idx
+            n_primary = int(primary_mask.sum())
+            if n_primary == 0:
+                continue
+
+            logger.debug(
+                "Profile '%s': %d primary sequences", profile["name"], n_primary,
+            )
+
+            for feat_idx, feature in enumerate(self.schema.features):
+                feat_dist = profile["features"][feature.name]
+                values = self._sample_feature(feat_dist, feature.num_values, (n_primary, T))
+                raw_values[primary_mask, :, feat_idx] = values
+
+        for prof_idx, profile in enumerate(profiles):
+            secondary_mask = secondary_indices == prof_idx
+            n_secondary = int(secondary_mask.sum())
+            if n_secondary == 0:
+                continue
+
+            seq_indices = np.where(secondary_mask)[0]
+            for feat_idx, feature in enumerate(self.schema.features):
+                feat_dist = profile["features"][feature.name]
+                alt_values = self._sample_feature(
+                    feat_dist, feature.num_values, (n_secondary, T),
+                )
+                for local_i, seq_i in enumerate(seq_indices):
+                    tx_blend = blend_mask[seq_i]
+                    raw_values[seq_i, tx_blend, feat_idx] = alt_values[local_i, tx_blend]
+
+        logger.info("Base features sampled with profile blending (%.0f%%)", self._blend_fraction * 100)
+
+        # --- Per-customer preferred merchant set + MCC-aligned assignment ---
+        mcc_values = raw_values[:, :, mcc_idx]
+        merch_params = np.zeros((N, 3), dtype=np.float64)
+        for i in range(N):
+            prof = profiles[primary_indices[i]]
+            md = prof["features"]["merchant_id"]
+            merch_params[i, 0] = md.get("top_n", 8)
+            merch_params[i, 1] = md.get("concentration", 0.88)
+            merch_params[i, 2] = md.get("zipf_exponent", 1.0)
+        self._assign_customer_merchants(
+            raw_values, mcc_values, merchant_idx, merch_params,
+        )
+
+        # --- Per-merchant amount distributions (tight std 4-7 buckets) ---
+        self._apply_merchant_amounts(raw_values, merchant_idx, amount_idx)
+
+        # --- MCC-conditional entry_mode (50% blend with profile) ---
+        self._apply_mcc_conditional_entry_mode(raw_values, mcc_values, entry_mode_idx)
+
+        # --- MCC-conditional hour-of-day ---
+        hour_idx = self.schema.feature_index("hour")
+        self._apply_mcc_conditional_hour(raw_values, mcc_values, hour_idx)
+
+        # --- MCC-conditional is_recurring ---
+        recurring_idx = self.schema.feature_index("is_recurring")
+        self._apply_mcc_conditional_recurring(raw_values, mcc_values, recurring_idx)
+
+        # --- Customer_tenure -> card_product correlation ---
+        tenure_idx = self.schema.feature_index("customer_tenure")
+        card_idx = self.schema.feature_index("card_product")
+        self._apply_tenure_card_correlation(raw_values, tenure_idx, card_idx)
+
+        # --- Country -> amount boost for international ---
+        country_idx = self.schema.feature_index("country")
+        self._apply_international_amount_boost(raw_values, country_idx, amount_idx)
+
+        # --- Temporal autocorrelation (MCC and merchant repeat) ---
+        self._apply_temporal_autocorrelation(
+            raw_values, mcc_idx, merchant_idx, amount_idx,
+        )
+
+        logger.info(
+            "Applied feature correlations: merchant-amount, MCC-entry_mode, "
+            "MCC-hour, MCC-recurring, tenure-card, country-amount, temporal-repeat",
+        )
+
+        # --- Tokenize ---
+        token_ids = (raw_values + VALUES_START).astype(np.int16)
+
+        # --- Fraud injection ---
+        tx_labels = np.zeros((N, T), dtype=np.int8)
+        self._inject_fraud(token_ids, tx_labels, primary_indices)
+
+        # --- Amount range labels (derived after fraud, before NULL) ---
+        amount_buckets = token_ids[:, :, amount_idx].astype(np.int32) - VALUES_START
+        amount_range_labels = amount_bucket_to_range(amount_buckets).astype(np.int8)
+
+        # --- NULL injection ---
+        self._apply_nulls(token_ids)
+
+        seq_labels = tx_labels.any(axis=1).astype(np.int8)
+        fraud_rate = seq_labels.mean()
+        logger.info("Generated %d sequences, fraud rate: %.3f", N, fraud_rate)
+
+        splits = self._compute_splits(N)
+        tokenizer = self._build_tokenizer()
+        fingerprint = self._compute_fingerprint(tokenizer, splits)
+
+        return GeneratedDataset(
+            token_ids=token_ids,
+            sequence_labels=seq_labels,
+            transaction_labels=tx_labels,
+            amount_range_labels=amount_range_labels,
+            split_indices=splits,
+            fingerprint=fingerprint,
+        )
+
+    def _assign_customer_merchants(
+        self,
+        raw_values: np.ndarray,
+        mcc_values: np.ndarray,
+        merchant_idx: int,
+        per_customer_params: np.ndarray,
+    ) -> None:
+        """Assign per-customer preferred merchant sets aligned with their MCCs.
+
+        Each customer gets a profile-specific number of preferred merchants
+        drawn from their most-visited MCCs. Concentration and Zipf exponent
+        also come from the customer's profile, so a retiree (8 merchants,
+        0.92 concentration) behaves differently from a digital nomad
+        (18 merchants, 0.72 concentration).
+
+        Args:
+            per_customer_params: (N, 3) array of [preferred_count, concentration, zipf_exp]
+        """
+        N, T, _ = raw_values.shape
+        mcc_idx_col = self.schema.feature_index("mcc")
+
+        for i in range(N):
+            pref_count = int(per_customer_params[i, 0])
+            concentration = float(per_customer_params[i, 1])
+            zipf_exp = float(per_customer_params[i, 2])
+
+            zipf_w = 1.0 / np.arange(1, pref_count + 1, dtype=np.float64) ** zipf_exp
+            zipf_w /= zipf_w.sum()
+
+            customer_mccs = mcc_values[i]
+            mcc_counts = np.bincount(customer_mccs[customer_mccs >= 0], minlength=100)
+            top_mccs = np.argsort(mcc_counts)[-pref_count:][::-1]
+            top_mccs = top_mccs[mcc_counts[top_mccs] > 0]
+
+            preferred_merchants: list[int] = []
+            for mcc_val in top_mccs:
+                candidates = self.catalog.mcc_to_merchants.get(int(mcc_val), np.array([]))
+                if len(candidates) == 0:
+                    continue
+                pop = self.catalog.popularity[candidates]
+                pop = pop / pop.sum()
+                n_from_mcc = max(1, pref_count // len(top_mccs))
+                n_from_mcc = min(n_from_mcc, len(candidates))
+                chosen = self.rng.choice(candidates, size=n_from_mcc, replace=False, p=pop)
+                preferred_merchants.extend(chosen.tolist())
+
+            if len(preferred_merchants) == 0:
+                raw_values[i, :, merchant_idx] = self.catalog.sample_merchants_for_mccs(
+                    customer_mccs, self.rng,
+                )
+                continue
+
+            pref_arr = np.array(preferred_merchants[:pref_count], dtype=np.int32)
+            n_pref = len(pref_arr)
+            pref_mccs = self.catalog.mcc_assignments[pref_arr]
+
+            pref_zipf = zipf_w[:n_pref].copy()
+            pref_zipf /= pref_zipf.sum()
+
+            for t in range(T):
+                tx_mcc = int(customer_mccs[t])
+                if self.rng.random() < concentration:
+                    same_mcc_mask = pref_mccs == tx_mcc
+                    if same_mcc_mask.any():
+                        candidates_idx = np.where(same_mcc_mask)[0]
+                        w = pref_zipf[candidates_idx]
+                        w /= w.sum()
+                        chosen = self.rng.choice(candidates_idx, p=w)
+                        raw_values[i, t, merchant_idx] = pref_arr[chosen]
+                    else:
+                        chosen = self.rng.choice(n_pref, p=pref_zipf)
+                        raw_values[i, t, merchant_idx] = pref_arr[chosen]
+                        raw_values[i, t, mcc_idx_col] = int(pref_mccs[chosen])
+                else:
+                    candidates = self.catalog.mcc_to_merchants.get(tx_mcc, np.array([]))
+                    if len(candidates) > 0:
+                        pop = self.catalog.popularity[candidates]
+                        pop = pop / pop.sum()
+                        raw_values[i, t, merchant_idx] = self.rng.choice(candidates, p=pop)
+                    else:
+                        raw_values[i, t, merchant_idx] = self.rng.integers(0, len(self.catalog.merchant_ids))
+
+    def _apply_merchant_amounts(
+        self,
+        raw_values: np.ndarray,
+        merchant_idx: int,
+        amount_idx: int,
+    ) -> None:
+        """Sample amounts from per-merchant distributions (tight std 4-7 buckets).
+
+        Each merchant has a characteristic (mean, std) derived from its MCC
+        category with per-merchant variation. This creates a strong
+        merchant->amount signal the model can learn.
+        """
+        N, T, _ = raw_values.shape
+        num_amount = self.schema.features[amount_idx].num_values
+        flat_merchants = raw_values[:, :, merchant_idx].ravel()
+
+        means = self.catalog.amount_mean[flat_merchants]
+        stds = self.catalog.amount_std[flat_merchants]
+        vals = self.rng.normal(means, stds)
+        vals = np.clip(np.round(vals), 0, num_amount - 1).astype(np.int32)
+        raw_values[:, :, amount_idx] = vals.reshape(N, T)
+
+    def _apply_mcc_conditional_entry_mode(
+        self,
+        raw_values: np.ndarray,
+        mcc_values: np.ndarray,
+        entry_mode_idx: int,
+    ) -> None:
+        """Blend MCC-conditional entry_mode with profile-sampled values (50/50)."""
+        N, T, _ = raw_values.shape
+        num_em = self.schema.features[entry_mode_idx].num_values
+        use_mcc = self.rng.random((N, T)) < 0.5
+
+        for (lo, hi), weights in MCC_ENTRY_MODE.items():
+            mask = (mcc_values >= lo) & (mcc_values <= hi) & use_mcc
+            n_matching = int(mask.sum())
+            if n_matching == 0:
+                continue
+            w = np.array(weights, dtype=np.float64)
+            w /= w.sum()
+            vals = self.rng.choice(len(w), size=n_matching, p=w).astype(np.int32)
+            raw_values[:, :, entry_mode_idx][mask] = vals
+
+    def _apply_mcc_conditional_hour(
+        self,
+        raw_values: np.ndarray,
+        mcc_values: np.ndarray,
+        hour_idx: int,
+    ) -> None:
+        """Blend MCC-conditional hour peaks with profile-sampled hours (50/50)."""
+        N, T, _ = raw_values.shape
+        use_mcc = self.rng.random((N, T)) < 0.5
+
+        for (lo, hi), (peak_hours, peak_weight) in MCC_HOUR_PEAKS.items():
+            mask = (mcc_values >= lo) & (mcc_values <= hi) & use_mcc
+            n_matching = int(mask.sum())
+            if n_matching == 0:
+                continue
+            w = np.ones(24, dtype=np.float64)
+            for h in peak_hours:
+                w[h] = peak_weight
+            w /= w.sum()
+            vals = self.rng.choice(24, size=n_matching, p=w).astype(np.int32)
+            raw_values[:, :, hour_idx][mask] = vals
+
+    def _apply_mcc_conditional_recurring(
+        self,
+        raw_values: np.ndarray,
+        mcc_values: np.ndarray,
+        recurring_idx: int,
+    ) -> None:
+        """Override is_recurring based on MCC category (70% weight)."""
+        N, T, _ = raw_values.shape
+        use_mcc = self.rng.random((N, T)) < 0.7
+
+        for (lo, hi), rate in MCC_RECURRING_RATE.items():
+            mask = (mcc_values >= lo) & (mcc_values <= hi) & use_mcc
+            n_matching = int(mask.sum())
+            if n_matching == 0:
+                continue
+            vals = (self.rng.random(n_matching) < rate).astype(np.int32)
+            raw_values[:, :, recurring_idx][mask] = vals
+
+    def _apply_tenure_card_correlation(
+        self,
+        raw_values: np.ndarray,
+        tenure_idx: int,
+        card_idx: int,
+    ) -> None:
+        """Correlate card_product with customer_tenure (60% blend)."""
+        N, T, _ = raw_values.shape
+        tenure_values = raw_values[:, :, tenure_idx]
+        use_corr = self.rng.random((N, T)) < 0.6
+        num_cards = self.schema.features[card_idx].num_values
+
+        for threshold, weights in TENURE_CARD_PRODUCT:
+            prev_threshold = 0
+            for prev_t, _ in TENURE_CARD_PRODUCT:
+                if prev_t < threshold:
+                    prev_threshold = prev_t + 1
+                    break
+
+            if threshold == TENURE_CARD_PRODUCT[0][0]:
+                mask = (tenure_values <= threshold) & use_corr
+            else:
+                lower = 0
+                for i, (t, _) in enumerate(TENURE_CARD_PRODUCT):
+                    if t == threshold and i > 0:
+                        lower = TENURE_CARD_PRODUCT[i - 1][0] + 1
+                        break
+                mask = (tenure_values >= lower) & (tenure_values <= threshold) & use_corr
+
+            n_matching = int(mask.sum())
+            if n_matching == 0:
+                continue
+            w = np.array(weights[:num_cards], dtype=np.float64)
+            w /= w.sum()
+            vals = self.rng.choice(num_cards, size=n_matching, p=w).astype(np.int32)
+            raw_values[:, :, card_idx][mask] = vals
+
+    def _apply_international_amount_boost(
+        self,
+        raw_values: np.ndarray,
+        country_idx: int,
+        amount_idx: int,
+    ) -> None:
+        """International transactions (country > 0) get a 30% amount boost."""
+        country_values = raw_values[:, :, country_idx]
+        num_amount = self.schema.features[amount_idx].num_values
+        international = country_values > 0
+        n_intl = int(international.sum())
+        if n_intl == 0:
+            return
+        current = raw_values[:, :, amount_idx][international].astype(np.float64)
+        boosted = current * 1.3 + self.rng.normal(0, 5, size=n_intl)
+        boosted = np.clip(np.round(boosted), 0, num_amount - 1).astype(np.int32)
+        raw_values[:, :, amount_idx][international] = boosted
+
+    def _apply_temporal_autocorrelation(
+        self,
+        raw_values: np.ndarray,
+        mcc_idx: int,
+        merchant_idx: int,
+        amount_idx: int,
+    ) -> None:
+        """Apply Markov-style temporal dependencies within each sequence.
+
+        For each transaction t > 0:
+          - P(repeat merchant+MCC from t-1) = 0.20 (same store again)
+          - P(repeat MCC only from t-1) = 0.10 (same category, maybe same store)
+          - If repeating, amount stays close to previous (std 3 buckets)
+        """
+        N, T, _ = raw_values.shape
+        num_amount = self.schema.features[amount_idx].num_values
+
+        for t in range(1, T):
+            merchant_repeat = self.rng.random(N) < 0.20
+            raw_values[merchant_repeat, t, merchant_idx] = raw_values[merchant_repeat, t - 1, merchant_idx]
+            raw_values[merchant_repeat, t, mcc_idx] = raw_values[merchant_repeat, t - 1, mcc_idx]
+
+            mcc_only = (self.rng.random(N) < 0.10) & ~merchant_repeat
+            raw_values[mcc_only, t, mcc_idx] = raw_values[mcc_only, t - 1, mcc_idx]
+
+            repeat_mask = merchant_repeat | mcc_only
+            n_repeating = int(repeat_mask.sum())
+            if n_repeating > 0:
+                prev_amt = raw_values[repeat_mask, t - 1, amount_idx].astype(np.float64)
+                new_amt = prev_amt + self.rng.normal(0, 3.0, size=n_repeating)
+                new_amt = np.clip(np.round(new_amt), 0, num_amount - 1).astype(np.int32)
+                raw_values[repeat_mask, t, amount_idx] = new_amt
+
+    # --- Fraud injection ---
+
+    def _inject_fraud(
+        self,
+        token_ids: np.ndarray,
+        tx_labels: np.ndarray,
+        profile_indices: np.ndarray,
+    ) -> None:
+        N = token_ids.shape[0]
+        fraud_rate = self._fraud_raw["overall_fraud_rate"]
+        patterns = self._fraud_raw["patterns"]
+        pattern_weights = np.array([p["weight"] for p in patterns])
+        pattern_weights /= pattern_weights.sum()
+
+        fraud_seq_mask = self.rng.random(N) < fraud_rate
+        fraud_indices = np.where(fraud_seq_mask)[0]
+
+        if len(fraud_indices) == 0:
+            return
+
+        pattern_assignments = self.rng.choice(
+            len(patterns), size=len(fraud_indices), p=pattern_weights,
+        )
+
+        for i, seq_idx in enumerate(fraud_indices):
+            pattern = patterns[pattern_assignments[i]]
+            self._inject_single_fraud(
+                token_ids[seq_idx], tx_labels[seq_idx], pattern,
+            )
+
+    def _inject_single_fraud(
+        self,
+        seq_tokens: np.ndarray,
+        seq_tx_labels: np.ndarray,
+        pattern: dict[str, Any],
+    ) -> None:
+        T = seq_tokens.shape[0]
+        inj = pattern["injection"]
+        lo, hi = inj["num_transactions"]
+        num_fraud_tx = self.rng.integers(lo, hi + 1)
+        num_fraud_tx = min(num_fraud_tx, T)
+
+        if inj["position"] == "late":
+            max_start = max(T // 2, T - num_fraud_tx)
+            start = self.rng.integers(T // 2, max_start + 1)
+        else:
+            start = self.rng.integers(0, max(T - num_fraud_tx + 1, 1))
+
+        end = min(start + num_fraud_tx, T)
+        seq_tx_labels[start:end] = 1
+
+        overrides = pattern.get("feature_overrides", {})
+        blend = overrides.pop("blend_with_profile", False) if "blend_with_profile" in overrides else False
+        span_len = end - start
+
+        for feat_name, feat_dist in overrides.items():
+            feat_idx = self.schema.feature_index(feat_name)
+            feat = self.schema.features[feat_idx]
+            fraud_values = self._sample_feature(feat_dist, feat.num_values, (1, span_len))
+            fraud_tokens = (fraud_values[0] + VALUES_START).astype(np.int16)
+
+            if blend:
+                keep_mask = self.rng.random(span_len) < 0.5
+                fraud_tokens[keep_mask] = seq_tokens[start:end, feat_idx][keep_mask]
+
+            seq_tokens[start:end, feat_idx] = fraud_tokens
+
+    # --- NULL injection ---
+
+    def _apply_nulls(self, token_ids: np.ndarray) -> None:
+        N, T, _ = token_ids.shape
+        null_rates = self._null_raw["null_rates"]
+        entry_mode_idx = self.schema.feature_index("entry_mode")
+        entry_mode_vals = token_ids[:, :, entry_mode_idx].astype(np.int32) - VALUES_START
+
+        for feat_idx, feature in enumerate(self.schema.features):
+            rate_config = null_rates.get(feature.name, 0.0)
+
+            if isinstance(rate_config, (int, float)):
+                if rate_config > 0:
+                    null_mask = self.rng.random((N, T)) < rate_config
+                    token_ids[:, :, feat_idx][null_mask] = NULL_TOKEN
+            elif isinstance(rate_config, dict):
+                null_mask = self.rng.random((N, T)) < rate_config.get("default", 0.0)
+
+                for cond in rate_config.get("conditional", []):
+                    when_modes = cond["when"]["entry_mode"]
+                    if isinstance(when_modes, str):
+                        when_modes = [when_modes]
+                    mode_vals = [self._entry_mode_name_to_val[m] for m in when_modes]
+                    cond_positions = np.isin(entry_mode_vals, mode_vals)
+                    cond_rand = self.rng.random((N, T))
+                    null_mask[cond_positions] = cond_rand[cond_positions] < cond["rate"]
+
+                token_ids[:, :, feat_idx][null_mask] = NULL_TOKEN
+
+    # --- Splits and fingerprint ---
+
+    def _compute_splits(self, num_sequences: int) -> dict[str, np.ndarray]:
+        perm = self.rng.permutation(num_sequences)
+        n_test = int(num_sequences * 0.10)
+        n_val = int(num_sequences * 0.05)
+
+        test_idx = np.sort(perm[:n_test])
+        val_idx = np.sort(perm[n_test : n_test + n_val])
+        train_idx = np.sort(perm[n_test + n_val :])
+
+        assert len(test_idx) + len(val_idx) + len(train_idx) == num_sequences
+        return {"train": train_idx, "val": val_idx, "test": test_idx}
+
+    def _build_tokenizer(self) -> TransactionTokenizer:
+        tokenizer = TransactionTokenizer(self.schema)
+        for feature in self.schema.features:
+            if feature.type == "bucketed":
+                tokenizer.get_feature_tokenizer(feature.name).fit_uniform_from_range()
+        return tokenizer
+
+    def _compute_fingerprint(
+        self, tokenizer: TransactionTokenizer, splits: dict[str, np.ndarray],
+    ) -> str:
+        hasher = hashlib.sha256()
+        for path in sorted(self._config_paths, key=str):
+            with open(path, "rb") as fh:
+                hasher.update(fh.read())
+        state_bytes = json.dumps(tokenizer.get_state(), sort_keys=True).encode("utf-8")
+        hasher.update(state_bytes)
+        for key in sorted(splits.keys()):
+            hasher.update(splits[key].tobytes())
+        hasher.update(str(self._seed).encode("utf-8"))
+        return hasher.hexdigest()
+
+    # --- Persistence ---
+
+    def save_dataset(
+        self, dataset: GeneratedDataset, output_dir: str | Path = "data/synthetic",
+    ) -> None:
+        out = Path(output_dir)
+        out.mkdir(parents=True, exist_ok=True)
+
+        np.save(out / "token_ids.npy", dataset.token_ids)
+        np.save(out / "sequence_labels.npy", dataset.sequence_labels)
+        np.save(out / "transaction_labels.npy", dataset.transaction_labels)
+        np.save(out / "amount_range_labels.npy", dataset.amount_range_labels)
+
+        np.savez(
+            out / "split_indices.npz",
+            train=dataset.split_indices["train"],
+            val=dataset.split_indices["val"],
+            test=dataset.split_indices["test"],
+        )
+
+        tokenizer = self._build_tokenizer()
+        tokenizer.save_state(out / "tokenizer_state.json")
+
+        with open(out / "fingerprint.txt", "w") as fh:
+            fh.write(dataset.fingerprint + "\n")
+
+        token_mb = dataset.token_ids.nbytes / 1024 / 1024
+        logger.info(
+            "Saved dataset to %s: token_ids=%.1f MB, fingerprint=%s",
+            out, token_mb, dataset.fingerprint[:16],
+        )

src/data/tokenizer.py

@@ -0,0 +1,316 @@
+"""Deterministic value-to-token mapping for structured transaction features.
+
+Each feature has its own tokenizer that maps raw values to token IDs using
+the reserved token convention (D6): 0=MASK, 1=OOV, 2=NULL, 3+=real values.
+Categorical features map directly; bucketed features use quantile or uniform
+boundaries computed during fit().
+
+OOV handling (D3): values outside the known range map to OOV token (1) and
+increment a per-feature counter logged during eval.
+"""
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from src.data.schema import (
+    FeatureSchema,
+    SchemaConfig,
+    MASK_TOKEN,
+    NULL_TOKEN,
+    OOV_TOKEN,
+    VALUES_START,
+    load_schema,
+)
+
+
+class FeatureTokenizer:
+    """Tokenizer for a single feature. Handles encode, decode, and OOV tracking."""
+
+    def __init__(self, schema: FeatureSchema) -> None:
+        self._schema = schema
+        self._boundaries: np.ndarray | None = None
+        self._oov_count: int = 0
+        self._fitted: bool = schema.type != "bucketed"
+
+    @property
+    def name(self) -> str:
+        return self._schema.name
+
+    @property
+    def vocab_size(self) -> int:
+        return self._schema.vocab_size
+
+    @property
+    def is_fitted(self) -> bool:
+        return self._fitted
+
+    @property
+    def oov_count(self) -> int:
+        return self._oov_count
+
+    def reset_oov_count(self) -> None:
+        self._oov_count = 0
+
+    def fit(self, values: np.ndarray) -> None:
+        """Compute bucket boundaries from data. Only valid for bucketed features.
+
+        Args:
+            values: 1-D array of raw continuous values to compute boundaries from.
+        """
+        assert self._schema.type == "bucketed", (
+            f"fit() only applies to bucketed features, not {self._schema.type}"
+        )
+        assert self._schema.bucket_range is not None
+        assert self._schema.bucket_method is not None
+
+        num_buckets = self._schema.num_values
+
+        if self._schema.bucket_method == "quantile":
+            quantiles = np.linspace(0.0, 1.0, num_buckets + 1)
+            self._boundaries = np.quantile(values, quantiles).astype(np.float64)
+        elif self._schema.bucket_method == "uniform":
+            lo, hi = self._schema.bucket_range
+            self._boundaries = np.linspace(lo, hi, num_buckets + 1, dtype=np.float64)
+        else:
+            raise ValueError(f"Unknown bucket_method: {self._schema.bucket_method}")
+
+        self._fitted = True
+
+    def fit_uniform_from_range(self) -> None:
+        """Compute uniform boundaries directly from the schema's bucket_range."""
+        assert self._schema.type == "bucketed"
+        assert self._schema.bucket_range is not None
+        lo, hi = self._schema.bucket_range
+        self._boundaries = np.linspace(
+            lo, hi, self._schema.num_values + 1, dtype=np.float64
+        )
+        self._fitted = True
+
+    def encode(self, values: np.ndarray) -> np.ndarray:
+        """Map raw values to token IDs.
+
+        Args:
+            values: array of raw feature values (any shape).
+
+        Returns:
+            int16 array of token IDs, same shape as input.
+        """
+        assert self._fitted, f"Feature '{self.name}' not fitted. Call fit() first."
+        original_shape = values.shape
+        flat = values.ravel()
+
+        if self._schema.type == "bucketed":
+            token_ids, oov_count = self._encode_bucketed(flat)
+        else:
+            token_ids, oov_count = self._encode_categorical(flat)
+
+        self._oov_count += oov_count
+        return token_ids.reshape(original_shape)
+
+    def decode(self, token_ids: np.ndarray) -> np.ndarray:
+        """Map token IDs back to values. Bucketed features return bucket centers.
+
+        Special tokens (MASK, OOV, NULL) decode to NaN.
+
+        Args:
+            token_ids: array of token IDs (any shape).
+
+        Returns:
+            float64 array of decoded values, same shape as input.
+        """
+        original_shape = token_ids.shape
+        flat = token_ids.ravel().astype(np.int64)
+        result = np.full(len(flat), np.nan, dtype=np.float64)
+
+        value_mask = flat >= VALUES_START
+        value_indices = flat[value_mask] - VALUES_START
+
+        if self._schema.type == "bucketed" and self._boundaries is not None:
+            centers = (self._boundaries[:-1] + self._boundaries[1:]) / 2.0
+            valid = value_indices < len(centers)
+            result_positions = np.where(value_mask)[0]
+            result[result_positions[valid]] = centers[value_indices[valid]]
+        else:
+            result[value_mask] = value_indices.astype(np.float64)
+
+        return result.reshape(original_shape)
+
+    def _encode_categorical(self, values: np.ndarray) -> tuple[np.ndarray, int]:
+        """Encode categorical/binary values. Returns (token_ids, oov_count)."""
+        int_values = values.astype(np.int64)
+        oov_mask = (int_values < 0) | (int_values >= self._schema.num_values)
+        token_ids = (int_values + VALUES_START).astype(np.int16)
+        token_ids[oov_mask] = OOV_TOKEN
+        return token_ids, int(oov_mask.sum())
+
+    def _encode_bucketed(self, values: np.ndarray) -> tuple[np.ndarray, int]:
+        """Encode bucketed values using pre-computed boundaries.
+
+        Returns (token_ids, oov_count). Uses np.digitize on internal boundaries
+        so bucket index i covers [boundaries[i], boundaries[i+1]).
+        The last bucket includes its upper bound: [boundaries[-2], boundaries[-1]].
+        """
+        assert self._boundaries is not None
+        internal_bins = self._boundaries[1:-1]
+        bucket_idx = np.digitize(values, internal_bins)
+        # bucket_idx in [0, num_values - 1] for values within boundary range.
+        # Can be num_values for values above the last internal bin, but np.digitize
+        # with N-1 internal bins returns at most N-1 for values < boundaries[-1].
+        # Values exactly at boundaries[-1] get num_values - 1 (last bucket).
+        bucket_idx = np.clip(bucket_idx, 0, self._schema.num_values - 1)
+
+        oov_mask = (values < self._boundaries[0]) | (values > self._boundaries[-1])
+        token_ids = (bucket_idx + VALUES_START).astype(np.int16)
+        token_ids[oov_mask] = OOV_TOKEN
+        return token_ids, int(oov_mask.sum())
+
+    def get_state(self) -> dict[str, Any]:
+        """Serializable state for fingerprinting and persistence."""
+        state: dict[str, Any] = {
+            "name": self._schema.name,
+            "type": self._schema.type,
+            "num_values": self._schema.num_values,
+            "vocab_size": self._schema.vocab_size,
+        }
+        if self._boundaries is not None:
+            state["boundaries"] = self._boundaries.tolist()
+        return state
+
+
+class TransactionTokenizer:
+    """Orchestrates tokenization across all features in the schema.
+
+    Usage:
+        schema = load_schema("data/schema.yaml")
+        tokenizer = TransactionTokenizer(schema)
+
+        # Fit bucketed features from raw data
+        tokenizer.fit_feature("amount", raw_amounts)
+        tokenizer.fit_feature("days_since_last", raw_days)
+
+        # Encode
+        token_ids = tokenizer.encode_feature("amount", raw_values)
+
+        # Save state and compute fingerprint
+        tokenizer.save_state("data/synthetic/tokenizer_state.json")
+        fp = tokenizer.compute_fingerprint("data/schema.yaml")
+    """
+
+    def __init__(self, schema: SchemaConfig) -> None:
+        self._schema = schema
+        self._tokenizers: dict[str, FeatureTokenizer] = {}
+        for feature in schema.features:
+            self._tokenizers[feature.name] = FeatureTokenizer(feature)
+
+    @property
+    def feature_names(self) -> list[str]:
+        return self._schema.feature_names()
+
+    @property
+    def num_features(self) -> int:
+        return self._schema.num_features
+
+    def get_feature_tokenizer(self, name: str) -> FeatureTokenizer:
+        return self._tokenizers[name]
+
+    def fit_feature(self, name: str, values: np.ndarray) -> None:
+        """Fit bucket boundaries for a single bucketed feature."""
+        self._tokenizers[name].fit(values)
+
+    def is_all_fitted(self) -> bool:
+        return all(t.is_fitted for t in self._tokenizers.values())
+
+    def encode_feature(self, name: str, values: np.ndarray) -> np.ndarray:
+        """Encode raw values for one feature. Returns int16 token IDs."""
+        return self._tokenizers[name].encode(values)
+
+    def decode_feature(self, name: str, token_ids: np.ndarray) -> np.ndarray:
+        """Decode token IDs for one feature. Returns float64 values."""
+        return self._tokenizers[name].decode(token_ids)
+
+    def inject_nulls(
+        self, token_ids: np.ndarray, null_mask: np.ndarray
+    ) -> np.ndarray:
+        """Replace positions where null_mask is True with NULL token.
+
+        Args:
+            token_ids: int16 array of token IDs.
+            null_mask: boolean array, same shape as token_ids.
+
+        Returns:
+            Copy of token_ids with NULLs injected.
+        """
+        result = token_ids.copy()
+        result[null_mask] = NULL_TOKEN
+        return result
+
+    @property
+    def oov_counts(self) -> dict[str, int]:
+        return {name: t.oov_count for name, t in self._tokenizers.items()}
+
+    def reset_oov_counts(self) -> None:
+        for t in self._tokenizers.values():
+            t.reset_oov_count()
+
+    def get_state(self) -> dict[str, Any]:
+        """Full tokenizer state for persistence and fingerprinting."""
+        return {
+            "num_features": self._schema.num_features,
+            "num_transactions": self._schema.num_transactions,
+            "features": [
+                self._tokenizers[name].get_state()
+                for name in self._schema.feature_names()
+            ],
+        }
+
+    def save_state(self, path: str | Path) -> None:
+        """Save tokenizer state to JSON for fingerprint computation (D4)."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        state = self.get_state()
+        with open(path, "w") as fh:
+            json.dump(state, fh, indent=2, sort_keys=True)
+
+    @classmethod
+    def from_state(cls, state_path: str | Path, schema: SchemaConfig) -> "TransactionTokenizer":
+        """Reconstruct tokenizer from saved state."""
+        with open(state_path) as fh:
+            state = json.load(fh)
+
+        tokenizer = cls(schema)
+        for feat_state in state["features"]:
+            name = feat_state["name"]
+            ft = tokenizer._tokenizers[name]
+            if "boundaries" in feat_state:
+                ft._boundaries = np.array(feat_state["boundaries"], dtype=np.float64)
+                ft._fitted = True
+        return tokenizer
+
+    def compute_fingerprint(self, *config_paths: str | Path) -> str:
+        """SHA256 fingerprint of tokenizer state plus config files (D4).
+
+        The fingerprint includes the tokenizer's bucket boundaries and vocab
+        sizes, plus the raw bytes of any additional config files (schema,
+        generator config, split indices). Eval refuses to run if fingerprints
+        don't match.
+
+        Args:
+            config_paths: paths to config files to include in the hash.
+
+        Returns:
+            Hex-encoded SHA256 digest.
+        """
+        hasher = hashlib.sha256()
+
+        state_bytes = json.dumps(self.get_state(), sort_keys=True).encode("utf-8")
+        hasher.update(state_bytes)
+
+        for path in sorted(str(p) for p in config_paths):
+            with open(path, "rb") as fh:
+                hasher.update(fh.read())
+
+        return hasher.hexdigest()

src/demo/__init__.py

@@ -0,0 +1 @@
+"""Interactive inference demo for the LFM2 transaction foundation model."""

src/demo/app.py

@@ -0,0 +1,828 @@
+"""Interactive inference demo for the LFM2 Transaction Foundation Model.
+
+Demonstrates multi-head predictions (fraud, next merchant, amount range, MCC)
+on synthetic payment sequences. Includes side-by-side pretrained-vs-random-init
+comparison showing the value of self-supervised pretraining.
+
+Usage:
+    python -m src.demo.app --checkpoint PATH [--data-dir PATH] [--port PORT]
+
+Runs on CPU. No GPU or internet required. Inference < 100ms per customer.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from src.data.schema import SchemaConfig, load_schema
+from src.demo.decode import TransactionDecoder
+from src.demo.merchant_catalog import DemoMerchantCatalog
+from src.demo.profile_inference import format_profile_html, infer_profile
+from src.demo.render import (
+    format_amount_predictions,
+    format_fraud_score,
+    format_mcc_predictions,
+    format_merchant_predictions,
+    format_timeline,
+    render_comparison_header,
+    render_integration_guide,
+    render_production_architecture,
+    render_why_liquid,
+)
+from src.model.lfm2_small import LFM2Small, ModelConfig
+from src.model.task_heads import (
+    AnyHead,
+    DownstreamHead,
+    HeadConfig,
+    MultiHeadModel,
+    TiedEmbeddingHead,
+)
+
+
+# ---------------------------------------------------------------------------
+# Data loading
+# ---------------------------------------------------------------------------
+
+
+class DemoData:
+    """Loads test-set sequences and labels into memory."""
+
+    def __init__(self, data_dir: Path, schema: SchemaConfig) -> None:
+        token_ids = np.load(data_dir / "token_ids.npy", mmap_mode="r")
+        seq_labels = np.load(data_dir / "sequence_labels.npy")
+        splits = np.load(data_dir / "split_indices.npz")
+
+        test_idx = splits["test"]
+        self.token_ids = np.array(token_ids[test_idx])
+        self.labels = seq_labels[test_idx].astype(int)
+        self.num_customers = len(test_idx)
+
+        self.fraud_indices = np.where(self.labels == 1)[0]
+        self.legit_indices = np.where(self.labels == 0)[0]
+
+        self._curated = self._find_curated_examples()
+
+    def _find_curated_examples(self) -> dict[str, int]:
+        """Pick interesting examples for quick navigation.
+
+        Five legitimate profiles to show breadth of normal spending patterns,
+        three fraud profiles representing the main attack archetypes.
+        """
+        examples: dict[str, int] = {}
+
+        # Legitimate customers — varied indices for behavioral variety
+        if len(self.legit_indices) > 0:
+            examples["Typical Customer"] = int(self.legit_indices[0])
+        if len(self.legit_indices) > 100:
+            examples["Frequent Shopper"] = int(self.legit_indices[100])
+        if len(self.legit_indices) > 200:
+            examples["Weekend Spender"] = int(self.legit_indices[200])
+        if len(self.legit_indices) > 300:
+            examples["High-Spend Loyalist"] = int(self.legit_indices[300])
+        if len(self.legit_indices) > 500:
+            examples["Low Activity"] = int(self.legit_indices[500])
+
+        # Fraud archetypes
+        if len(self.fraud_indices) > 0:
+            examples["Fraud: Card Testing"] = int(self.fraud_indices[0])
+        if len(self.fraud_indices) > 50:
+            examples["Fraud: Account Takeover"] = int(self.fraud_indices[50])
+        if len(self.fraud_indices) > 100:
+            examples["Fraud: High Value"] = int(self.fraud_indices[100])
+
+        return examples
+
+    @property
+    def curated_names(self) -> list[str]:
+        return list(self._curated.keys())
+
+    def get_curated_index(self, name: str) -> int:
+        return self._curated[name]
+
+
+# ---------------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------------
+
+
+def build_model(
+    model_yaml: Path,
+    schema: SchemaConfig,
+    finetune_yaml: Path,
+) -> MultiHeadModel:
+    """Construct MultiHeadModel with 4 downstream heads."""
+    import yaml
+
+    backbone = LFM2Small(ModelConfig.from_yaml(model_yaml), schema)
+
+    with open(finetune_yaml) as f:
+        ft_config = yaml.safe_load(f)
+
+    heads: dict[str, AnyHead] = {}
+    for name, hcfg in ft_config["heads"].items():
+        config = HeadConfig(
+            name=name,
+            output_dim=hcfg["output_dim"],
+            loss_type=hcfg["loss"],
+            pool_strategy=hcfg["pool"],
+            target_type=hcfg["target"],
+            weight=hcfg.get("weight", 1.0),
+            mlp_hidden=hcfg.get("mlp_hidden", 128),
+            dropout=hcfg.get("dropout", 0.1),
+        )
+        if hcfg.get("tied", False):
+            heads[name] = TiedEmbeddingHead(
+                config, backbone.config.hidden_size, schema.num_features,
+                backbone.embedding.value_tables,
+            )
+        else:
+            heads[name] = DownstreamHead(
+                config, backbone.config.hidden_size, schema.num_features,
+            )
+
+    return MultiHeadModel(backbone, heads)
+
+
+def load_model_checkpoint(model: MultiHeadModel, checkpoint_path: Path | None) -> str:
+    """Load fine-tuned weights. Returns status message."""
+    if checkpoint_path is None:
+        return "No checkpoint loaded"
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {checkpoint_path}. "
+            f"Run with --checkpoint PATH or --checkpoint none to skip."
+        )
+
+    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(ckpt["model_state_dict"], strict=True)
+    step = ckpt.get("step", "?")
+    return f"step {step}"
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def run_inference(
+    model: MultiHeadModel,
+    token_ids: np.ndarray,
+) -> dict[str, np.ndarray]:
+    """Run all 4 heads on a single customer sequence."""
+    tensor = torch.from_numpy(token_ids).unsqueeze(0).long()
+    predictions = model(tensor)
+
+    results: dict[str, np.ndarray] = {}
+    for name, logits in predictions.items():
+        if name == "fraud":
+            prob = torch.sigmoid(logits).squeeze().numpy()
+            results[name] = np.array([float(prob)])
+        else:
+            probs = F.softmax(logits, dim=-1).squeeze(0).numpy()
+            results[name] = probs
+
+    return results
+
+
+# ---------------------------------------------------------------------------
+# App builder
+# ---------------------------------------------------------------------------
+
+
+def create_app(
+    pretrained_model: MultiHeadModel,
+    random_model: MultiHeadModel,
+    data: DemoData,
+    decoder: TransactionDecoder,
+    merchant_catalog: DemoMerchantCatalog,
+    checkpoint_status: str,
+) -> gr.Blocks:
+    """Build the Gradio app with pretrained vs random-init comparison."""
+
+    pretrained_model.eval()
+    random_model.eval()
+
+    def on_customer_select(
+        curated_name: str | None,
+        customer_idx: int,
+        mode: str,
+    ) -> tuple[str, str, str, str, str, str, str, str, str, str, str]:
+        """Run both models, return all outputs for comparison."""
+
+        if mode == "Curated examples" and curated_name:
+            idx = data.get_curated_index(curated_name)
+        else:
+            idx = int(customer_idx)
+
+        idx = max(0, min(idx, data.num_customers - 1))
+        token_ids = data.token_ids[idx]
+        is_fraud = bool(data.labels[idx])
+
+        summary = decoder.summarize_customer(token_ids, is_fraud)
+
+        t0 = time.perf_counter()
+        pre_results = run_inference(pretrained_model, token_ids)
+        rand_results = run_inference(random_model, token_ids)
+        latency_ms = (time.perf_counter() - t0) * 1000
+
+        timeline_html = format_timeline(decoder, token_ids)
+
+        # Pretrained predictions
+        pre_fraud = format_fraud_score(float(pre_results["fraud"][0]))
+        pre_merchant = format_merchant_predictions(
+            pre_results["next_merchant"], merchant_catalog, k=5,
+        )
+        pre_amount = format_amount_predictions(pre_results["amount_range"], k=5)
+        pre_mcc = format_mcc_predictions(pre_results["mcc"], k=5)
+
+        # Random-init predictions
+        rand_fraud = format_fraud_score(float(rand_results["fraud"][0]))
+        rand_merchant = format_merchant_predictions(
+            rand_results["next_merchant"], merchant_catalog, k=5,
+        )
+        rand_amount = format_amount_predictions(rand_results["amount_range"], k=5)
+        rand_mcc = format_mcc_predictions(rand_results["mcc"], k=5)
+
+        # Behavioral profile
+        profile_match = infer_profile(token_ids)
+        profile_html = format_profile_html(profile_match)
+
+        # Combined comparison HTML for each head
+        fraud_compare = _side_by_side("Fraud Score", pre_fraud, rand_fraud)
+        merchant_compare = _side_by_side("Next Merchant", pre_merchant, rand_merchant)
+        amount_compare = _side_by_side("Amount Range", pre_amount, rand_amount)
+        mcc_compare = _side_by_side("Merchant Category", pre_mcc, rand_mcc)
+
+        latency_str = (
+            f"Inference: {latency_ms:.1f}ms (both models) on CPU | "
+            f"Customer #{idx} | Ground truth: {'FRAUD' if is_fraud else 'Legitimate'}"
+        )
+
+        return (
+            summary, timeline_html, profile_html,
+            fraud_compare, merchant_compare, amount_compare, mcc_compare,
+            latency_str,
+        )
+
+    with gr.Blocks(
+        title="LFM2 Transaction Foundation Model",
+    ) as app:
+        gr.HTML("""
+        <div style="text-align: center; margin-bottom: 16px;">
+            <h1 style="margin: 0; font-size: 24px; font-weight: 700; color: #171717;
+                       letter-spacing: -0.02em;">
+                LFM2 Transaction Foundation Model
+            </h1>
+            <p style="color: #737373; margin: 6px 0 0 0; font-size: 13px;
+                      font-family: JetBrains Mono, ui-monospace, monospace;">
+                Liquid AI &middot; LFM2.5 Architecture &middot; Multi-Head Inference Demo
+            </p>
+        </div>
+        """)
+
+        with gr.Tabs():
+            # ===== Tab 1: Interactive Predictions =====
+            with gr.Tab("Predictions"):
+                gr.HTML("""
+                <div style="padding: 10px 14px; background: #ffffff; border: 1px solid rgba(0,0,0,0.1);
+                            border-radius: 12px; margin: 8px 0; font-size: 12px; color: #525252;">
+                    <b style="color: #171717;">How to read this:</b> Select a customer, see their
+                    transaction history, then compare predictions from the pretrained model (left)
+                    vs random initialization (right). Same architecture, same fine-tuning data.
+                    The only difference is self-supervised pretraining on unlabeled sequences.
+                </div>
+                """)
+                with gr.Accordion("Reference model details", open=False):
+                    gr.HTML(f"""
+                    <div style="padding: 8px 12px; font-family: JetBrains Mono, ui-monospace, monospace;
+                                font-size: 11px; color: #525252; display: flex; gap: 20px; flex-wrap: wrap;">
+                        <span>arch: <b style="color: #171717;">LFM2-small</b> 9.8M params</span>
+                        <span>layers: <b style="color: #10B981;">5 conv</b> + <b style="color: #7c3aed;">3 attn</b></span>
+                        <span>input: 64 tx &times; 15 feat = 960 tokens</span>
+                        <span>checkpoint: <b style="color: #171717;">{checkpoint_status}</b></span>
+                        <span>data: 200K synthetic sequences, 15 features/tx</span>
+                    </div>
+                    """)
+
+                # Customer Selection. Use Radio (not Checkbox) for the
+                # curated-vs-browse toggle: the two modes are mutually exclusive
+                # and Radio's visual state is more reliable across Gradio versions.
+                _CURATED_MODE = "Curated examples"
+                _BROWSE_MODE = "Browse all customers"
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.HTML("<h3 style='margin: 0 0 8px 0;'>Select Customer</h3>")
+                        selection_mode = gr.Radio(
+                            choices=[_CURATED_MODE, _BROWSE_MODE],
+                            value=_CURATED_MODE,
+                            label="Selection mode",
+                            info="Curated: hand-picked legitimate and fraud examples. "
+                                 "Browse: pick any of 20,000 test customers by index.",
+                            elem_classes="liquid-radio",
+                        )
+                        curated_dropdown = gr.Dropdown(
+                            choices=data.curated_names,
+                            value=data.curated_names[0] if data.curated_names else None,
+                            label="Curated Examples",
+                            info="5 legitimate profiles, 3 fraud archetypes",
+                        )
+                        customer_slider = gr.Slider(
+                            minimum=0,
+                            maximum=data.num_customers - 1,
+                            step=1,
+                            value=0,
+                            label=f"Customer Index (0-{data.num_customers - 1})",
+                            info=f"Direct access to any of {data.num_customers:,} test-set customers. "
+                                 f"~3.7% are fraud, rest are legitimate.",
+                            visible=False,
+                        )
+                        run_btn = gr.Button(
+                            "Run Inference", variant="primary", size="lg",
+                            elem_id="run-inference-btn",
+                        )
+
+                    with gr.Column(scale=2):
+                        summary_text = gr.Textbox(
+                            label="Customer Profile", interactive=False, lines=2,
+                        )
+                        profile_output = gr.HTML(label="Behavioral Profile")
+                        latency_text = gr.Textbox(
+                            label="Performance", interactive=False, lines=1,
+                        )
+
+                def toggle_selector(mode: str) -> tuple[Any, Any]:
+                    use_cur = (mode == _CURATED_MODE)
+                    return gr.update(visible=use_cur), gr.update(visible=not use_cur)
+
+                selection_mode.change(
+                    toggle_selector, inputs=[selection_mode],
+                    outputs=[curated_dropdown, customer_slider],
+                )
+
+                # Transaction Timeline
+                gr.HTML("""<h3 style='margin: 16px 0 8px 0; color: #171717;'>Transaction History</h3>
+                <div style="font-size: 11px; color: #737373; margin-bottom: 4px;">
+                    64 most recent transactions. Tx 63 (highlighted) is the most recent.
+                    The model predicts what comes next based on this full sequence.
+                </div>""")
+                timeline_output = gr.HTML()
+
+                # Side-by-side predictions
+                gr.HTML("<h3 style='margin: 16px 0 4px 0; color: #171717; font-size: 18px; font-weight: 600; letter-spacing: -0.01em;'>Model Predictions: Pretrained vs Random Init</h3>")
+                gr.HTML(render_comparison_header())
+
+                fraud_output = gr.HTML()
+                merchant_output = gr.HTML()
+                amount_output = gr.HTML()
+                mcc_output = gr.HTML()
+
+                # Wire callbacks
+                outputs = [
+                    summary_text, timeline_output, profile_output,
+                    fraud_output, merchant_output, amount_output, mcc_output,
+                    latency_text,
+                ]
+
+                run_btn.click(
+                    on_customer_select,
+                    inputs=[curated_dropdown, customer_slider, selection_mode],
+                    outputs=outputs,
+                )
+                curated_dropdown.change(
+                    on_customer_select,
+                    inputs=[curated_dropdown, customer_slider, selection_mode],
+                    outputs=outputs,
+                )
+
+            # ===== Tab 2: Architecture Deep Dive =====
+            with gr.Tab("Architecture"):
+                gr.HTML(render_production_architecture())
+
+            # ===== Tab 3: Why Liquid =====
+            with gr.Tab("Why Liquid AI"):
+                gr.HTML(render_why_liquid())
+
+            # ===== Tab 4: Integration Guide =====
+            with gr.Tab("Integration"):
+                gr.HTML(render_integration_guide())
+
+    return app
+
+
+def _side_by_side(title: str, pretrained_html: str, random_html: str) -> str:
+    """Render pretrained vs random-init predictions side by side."""
+    _mono = "JetBrains Mono, ui-monospace, monospace"
+    return f"""
+    <div style="margin-bottom: 16px;">
+        <div style="font-size: 13px; font-weight: 600; color: #171717; margin-bottom: 8px;
+                    letter-spacing: -0.01em;">
+            {title}
+        </div>
+        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 8px;">
+            <div style="background: #ffffff; border: 1px solid rgba(16,185,129,0.25);
+                        border-radius: 12px; padding: 12px; box-shadow: 0 1px 3px rgba(0,0,0,0.04);">
+                <div style="font-family: {_mono}; font-size: 10px; color: #10B981;
+                            font-weight: 600; margin-bottom: 8px; text-transform: uppercase;
+                            letter-spacing: 0.05em;">
+                    &#10003; Pretrained
+                </div>
+                {pretrained_html}
+            </div>
+            <div style="background: #ffffff; border: 1px solid rgba(0,0,0,0.08);
+                        border-radius: 12px; padding: 12px; box-shadow: 0 1px 3px rgba(0,0,0,0.04);">
+                <div style="font-family: {_mono}; font-size: 10px; color: #a3a3a3;
+                            font-weight: 600; margin-bottom: 8px; text-transform: uppercase;
+                            letter-spacing: 0.05em;">
+                    &#10007; Random Init
+                </div>
+                {random_html}
+            </div>
+        </div>
+    </div>
+    """
+
+
+
+# ---------------------------------------------------------------------------
+# CLI entrypoint
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="LFM2 Transaction Foundation Model — Interactive Inference Demo",
+    )
+    parser.add_argument(
+        "--checkpoint", type=Path,
+        default=Path("experiments/v2_tied/finetune_20260516_190905/checkpoints/step_004999.pt"),
+        help="Path to fine-tuned MultiHeadModel checkpoint (.pt file)",
+    )
+    parser.add_argument(
+        "--data-dir", type=Path, default=Path("data/synthetic"),
+        help="Directory containing token_ids.npy, sequence_labels.npy, split_indices.npz",
+    )
+    parser.add_argument(
+        "--model-config", type=Path, default=Path("configs/model.yaml"),
+        help="Model backbone YAML config",
+    )
+    parser.add_argument(
+        "--schema", type=Path, default=Path("data/schema.yaml"),
+        help="Feature schema YAML",
+    )
+    parser.add_argument(
+        "--finetune-config", type=Path,
+        default=Path("experiments/v2_tied/finetune_20260516_190905/finetune_config.yaml"),
+        help="Fine-tune head config YAML",
+    )
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--share", action="store_true", help="Create public Gradio link")
+
+    args = parser.parse_args()
+
+    print("Loading schema and merchant catalog...")
+    schema = load_schema(args.schema)
+    merchant_catalog = DemoMerchantCatalog(schema)
+
+    print("Loading test data...")
+    demo_data = DemoData(args.data_dir, schema)
+    print(f"  {demo_data.num_customers} test customers "
+          f"({len(demo_data.fraud_indices)} fraud, {len(demo_data.legit_indices)} legitimate)")
+
+    print("Building pretrained model...")
+    pretrained_model = build_model(args.model_config, schema, args.finetune_config)
+    checkpoint_status = load_model_checkpoint(pretrained_model, args.checkpoint)
+    print(f"  Pretrained: {checkpoint_status}")
+
+    print("Building random-init baseline model...")
+    random_model = build_model(args.model_config, schema, args.finetune_config)
+    print("  Random-init: fresh weights (no checkpoint)")
+
+    total_params = sum(p.numel() for p in pretrained_model.parameters())
+    print(f"  Parameters per model: {total_params:,}")
+
+    print("Building decoder...")
+    decoder = TransactionDecoder(schema, merchant_catalog)
+
+    print("Launching demo...")
+    app = create_app(
+        pretrained_model, random_model, demo_data,
+        decoder, merchant_catalog, checkpoint_status,
+    )
+    _liquid_theme = gr.themes.Soft(
+        primary_hue="neutral",
+        secondary_hue="neutral",
+        neutral_hue="neutral",
+        font=gr.themes.GoogleFont("Inter"),
+        font_mono=gr.themes.GoogleFont("JetBrains Mono"),
+    ).set(
+        body_background_fill="#f5f5f5",
+        body_background_fill_dark="#f5f5f5",
+        body_text_color="#171717",
+        body_text_color_dark="#171717",
+        body_text_color_subdued="#737373",
+        body_text_color_subdued_dark="#737373",
+        block_background_fill="#ffffff",
+        block_background_fill_dark="#ffffff",
+        block_border_color="rgba(0,0,0,0.1)",
+        block_border_color_dark="rgba(0,0,0,0.1)",
+        block_label_background_fill="#f5f5f5",
+        block_label_background_fill_dark="#f5f5f5",
+        block_label_text_color="#525252",
+        block_label_text_color_dark="#525252",
+        block_title_text_color="#171717",
+        block_title_text_color_dark="#171717",
+        block_shadow="0 1px 3px rgba(0,0,0,0.04)",
+        block_shadow_dark="0 1px 3px rgba(0,0,0,0.04)",
+        input_background_fill="#ffffff",
+        input_background_fill_dark="#ffffff",
+        input_background_fill_focus="#ffffff",
+        input_background_fill_focus_dark="#ffffff",
+        input_border_color="rgba(0,0,0,0.1)",
+        input_border_color_dark="rgba(0,0,0,0.1)",
+        input_border_color_focus="#171717",
+        input_border_color_focus_dark="#171717",
+        input_placeholder_color="#a3a3a3",
+        input_placeholder_color_dark="#a3a3a3",
+        panel_background_fill="#fafafa",
+        panel_background_fill_dark="#fafafa",
+        panel_border_color="rgba(0,0,0,0.06)",
+        panel_border_color_dark="rgba(0,0,0,0.06)",
+        border_color_primary="rgba(0,0,0,0.1)",
+        border_color_primary_dark="rgba(0,0,0,0.1)",
+        button_primary_background_fill="#171717",
+        button_primary_background_fill_dark="#171717",
+        button_primary_background_fill_hover="#404040",
+        button_primary_background_fill_hover_dark="#404040",
+        button_primary_text_color="#ffffff",
+        button_primary_text_color_dark="#ffffff",
+        button_secondary_background_fill="#ffffff",
+        button_secondary_background_fill_dark="#ffffff",
+        button_secondary_text_color="#525252",
+        button_secondary_text_color_dark="#525252",
+        button_secondary_border_color="rgba(0,0,0,0.1)",
+        button_secondary_border_color_dark="rgba(0,0,0,0.1)",
+        checkbox_background_color="#ffffff",
+        checkbox_background_color_dark="#ffffff",
+        checkbox_border_color="rgba(0,0,0,0.2)",
+        checkbox_border_color_dark="rgba(0,0,0,0.2)",
+        checkbox_background_color_selected="#171717",
+        checkbox_background_color_selected_dark="#171717",
+        checkbox_label_background_fill="#ffffff",
+        checkbox_label_background_fill_dark="#ffffff",
+        checkbox_label_text_color="#171717",
+        checkbox_label_text_color_dark="#171717",
+        slider_color="#171717",
+        slider_color_dark="#171717",
+        table_border_color="rgba(0,0,0,0.06)",
+        table_border_color_dark="rgba(0,0,0,0.06)",
+        table_even_background_fill="#fafafa",
+        table_even_background_fill_dark="#fafafa",
+        table_odd_background_fill="#ffffff",
+        table_odd_background_fill_dark="#ffffff",
+        shadow_spread="0px",
+        shadow_spread_dark="0px",
+        color_accent_soft="rgba(0,0,0,0.04)",
+        color_accent_soft_dark="rgba(0,0,0,0.04)",
+    )
+
+    _liquid_css = """
+    /* Force light mode regardless of system preference */
+    :root, .dark { color-scheme: light !important; }
+
+    .gradio-container {
+        background: #f5f5f5 !important;
+        max-width: 1280px !important;
+        margin: auto !important;
+        padding: 1.5rem !important;
+    }
+
+    /* Tabs: pill-style matching Liquid design system */
+    .tabs { background: transparent !important; }
+    .tab-nav {
+        background: #f5f5f5 !important;
+        border: none !important;
+        border-bottom: 1px solid rgba(0,0,0,0.1) !important;
+        gap: 4px !important;
+        padding: 4px 0 !important;
+    }
+    .tab-nav button {
+        font-weight: 500 !important;
+        font-size: 14px !important;
+        color: #737373 !important;
+        background: transparent !important;
+        border: none !important;
+        border-bottom: 2px solid transparent !important;
+        padding: 8px 16px !important;
+        border-radius: 0 !important;
+        transition: all 0.15s ease !important;
+    }
+    .tab-nav button:hover {
+        color: #171717 !important;
+        background: rgba(0,0,0,0.03) !important;
+    }
+    .tab-nav button.selected {
+        color: #171717 !important;
+        font-weight: 600 !important;
+        border-bottom: 2px solid #171717 !important;
+        background: transparent !important;
+    }
+
+    /* Block/component overrides */
+    .block { border-radius: 12px !important; }
+    .block.padded { background: #ffffff !important; }
+
+    /* Input, textarea, dropdown */
+    input, textarea, select, .wrap {
+        background: #ffffff !important;
+        color: #171717 !important;
+        border-color: rgba(0,0,0,0.1) !important;
+    }
+    .secondary-wrap, .wrap-inner {
+        background: #ffffff !important;
+    }
+
+    /* Labels */
+    label, .label-wrap, span.svelte-1gfkn6j {
+        color: #171717 !important;
+    }
+    .info {
+        color: #737373 !important;
+    }
+
+    /* Action buttons: pill style applied only to explicit primary actions.
+       Scoped by elem_id so it doesn't bleed into Gradio radio/checkbox
+       options (which Gradio also renders as <button class="primary">). */
+    #run-inference-btn button,
+    button#run-inference-btn {
+        background: #171717 !important;
+        color: #ffffff !important;
+        border-radius: 9999px !important;
+        border: none !important;
+        font-weight: 500 !important;
+        letter-spacing: -0.02em !important;
+    }
+    #run-inference-btn button:hover,
+    button#run-inference-btn:hover {
+        background: #404040 !important;
+    }
+
+    /* Radio group: render as a clean vertical list with native circle
+       indicators. Scoped via elem_classes="liquid-radio" so we can target
+       reliably without depending on Gradio's internal Svelte-hashed class
+       names. Defeats Gradio's default of styling the selected option as a
+       depressed dark button -- a form selection should not look like an
+       action button. */
+    .liquid-radio,
+    .liquid-radio > * {
+        background: transparent !important;
+        border: none !important;
+        box-shadow: none !important;
+    }
+    .liquid-radio .wrap,
+    .liquid-radio fieldset {
+        display: flex !important;
+        flex-direction: column !important;
+        gap: 2px !important;
+        padding: 0 !important;
+    }
+    /* Each option label -- override Gradio's button-like rendering */
+    .liquid-radio label {
+        display: flex !important;
+        align-items: center !important;
+        gap: 10px !important;
+        padding: 8px 10px !important;
+        background: transparent !important;
+        background-color: transparent !important;
+        background-image: none !important;
+        border: none !important;
+        border-radius: 8px !important;
+        cursor: pointer !important;
+        font-size: 13px !important;
+        font-weight: 400 !important;
+        color: #171717 !important;
+        box-shadow: none !important;
+        transition: background 0.1s ease !important;
+    }
+    .liquid-radio label:hover {
+        background: rgba(0,0,0,0.04) !important;
+    }
+    /* The radio input itself -- native circle, dark fill when checked */
+    .liquid-radio input[type="radio"] {
+        appearance: auto !important;
+        -webkit-appearance: radio !important;
+        accent-color: #171717 !important;
+        width: 16px !important;
+        height: 16px !important;
+        min-width: 16px !important;
+        margin: 0 !important;
+        cursor: pointer !important;
+        opacity: 1 !important;
+    }
+    /* Selected option: subtle background tint + slightly bolder text.
+       Three selectors because different Gradio versions mark the selected
+       option differently: .selected class, [aria-checked="true"], or
+       :has(input:checked). */
+    .liquid-radio label.selected,
+    .liquid-radio label[aria-checked="true"],
+    .liquid-radio label:has(input:checked) {
+        background: rgba(0,0,0,0.04) !important;
+        color: #171717 !important;
+        font-weight: 500 !important;
+    }
+    /* Kill any inherited button.primary/button.secondary styling that
+       Gradio may apply to radio option wrappers in some versions. */
+    .liquid-radio button,
+    .liquid-radio button.primary,
+    .liquid-radio button.secondary {
+        background: transparent !important;
+        color: #171717 !important;
+        border: none !important;
+        border-radius: 8px !important;
+        box-shadow: none !important;
+        font-weight: 400 !important;
+        text-align: left !important;
+        justify-content: flex-start !important;
+    }
+    .liquid-radio button.primary,
+    .liquid-radio button.selected {
+        background: rgba(0,0,0,0.04) !important;
+        font-weight: 500 !important;
+    }
+
+    /* Checkbox */
+    .checkbox-item { color: #171717 !important; }
+
+    /* Dropdown */
+    .dropdown-arrow { color: #525252 !important; }
+    ul.options { background: #ffffff !important; border-color: rgba(0,0,0,0.1) !important; }
+    ul.options li { color: #171717 !important; }
+    ul.options li:hover, ul.options li.selected {
+        background: #f5f5f5 !important;
+    }
+
+    /* Textbox display (non-editable) */
+    .textbox textarea[disabled], .textbox input[disabled] {
+        background: #fafafa !important;
+        color: #171717 !important;
+        opacity: 1 !important;
+    }
+
+    /* Remove dark shadows/borders */
+    .block { box-shadow: 0 1px 3px rgba(0,0,0,0.04) !important; }
+
+    /* Accordion/group headers */
+    .form { background: #ffffff !important; border-color: rgba(0,0,0,0.06) !important; }
+
+    /* Override any remaining dark backgrounds */
+    [class*="dark:"], .dark * {
+        --tw-bg-opacity: 1 !important;
+    }
+    """
+
+    _force_light_js = """
+    () => {
+        document.documentElement.classList.remove('dark');
+        document.documentElement.style.colorScheme = 'light';
+        const meta = document.createElement('meta');
+        meta.name = 'color-scheme';
+        meta.content = 'light';
+        document.head.appendChild(meta);
+    }
+    """
+
+    # Try the requested port first, then walk up to 10 ports.
+    # Gradio errors with OSError when a port is taken; we retry transparently
+    # so a stale background process doesn't block a fresh launch.
+    for port in range(args.port, args.port + 10):
+        try:
+            app.launch(
+                server_port=port,
+                share=args.share,
+                theme=_liquid_theme,
+                css=_liquid_css,
+                js=_force_light_js,
+            )
+            break
+        except OSError as e:
+            if "Cannot find empty port" in str(e) or "Address already in use" in str(e):
+                print(f"  port {port} in use, trying {port + 1}...")
+                continue
+            raise
+    else:
+        raise RuntimeError(
+            f"No free port in range {args.port}-{args.port + 9}. "
+            f"Kill stale processes: lsof -ti:{args.port} | xargs kill"
+        )
+
+
+if __name__ == "__main__":
+    main()

src/demo/decode.py

@@ -0,0 +1,234 @@
+"""Decode token IDs back to human-readable transaction descriptions.
+
+Handles the reserved-token offset (0=MASK, 1=OOV, 2=NULL, 3+=values)
+and maps each feature's raw value to readable text.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from src.data.generator import AMOUNT_RANGE_LABELS
+from src.data.schema import SchemaConfig, VALUES_START
+from src.demo.merchant_catalog import DemoMerchantCatalog, MCC_NAMES
+
+
+DOW_NAMES: list[str] = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
+
+ENTRY_MODE_DISPLAY: dict[str, str] = {
+    "card_present": "Card Present",
+    "card_not_present": "Online",
+    "contactless": "Contactless",
+    "chip": "Chip",
+    "manual_key": "Manual",
+}
+
+COUNTRY_NAMES: list[str] = [
+    "US", "UK", "CA", "DE", "FR", "JP", "AU", "BR", "IN", "MX",
+    "IT", "ES", "NL", "CH", "SE", "NO", "DK", "FI", "KR", "SG",
+    "HK", "TW", "NZ", "IE", "BE", "AT", "PT", "PL", "CZ", "GR",
+    "IL", "AE", "SA", "TH", "MY", "PH", "ID", "VN", "CL", "CO",
+    "AR", "ZA", "NG", "EG", "KE", "TR", "RU", "UA", "RO", "HU",
+]
+
+AVS_DISPLAY: dict[str, str] = {
+    "full_match": "AVS Match",
+    "zip_match": "ZIP Match",
+    "address_match": "Addr Match",
+    "no_match": "AVS No Match",
+    "not_checked": "AVS N/A",
+}
+
+CVV_DISPLAY: dict[str, str] = {
+    "match": "CVV Match",
+    "no_match": "CVV No Match",
+    "not_provided": "No CVV",
+}
+
+
+@dataclass
+class DecodedTransaction:
+    """A single transaction with all features decoded to display strings."""
+
+    index: int
+    hour: str
+    dow: str
+    days_since_last: str
+    is_recurring: str
+    mcc: str
+    merchant_name: str
+    merchant_category: str
+    customer_merchant_count: str
+    entry_mode: str
+    amount_range: str
+    card_product: str
+    country: str
+    avs: str
+    cvv: str
+    device_hash: str
+    customer_tenure: str
+
+
+class TransactionDecoder:
+    """Converts raw token_ids (T, F) into human-readable transactions."""
+
+    def __init__(self, schema: SchemaConfig, merchant_catalog: DemoMerchantCatalog) -> None:
+        self.schema = schema
+        self.merchants = merchant_catalog
+        self._feature_names = schema.feature_names()
+
+    def _decode_value(self, feature_idx: int, token_id: int) -> str:
+        """Decode a single token_id for a given feature index."""
+        feat = self.schema.features[feature_idx]
+
+        if token_id == 0:
+            return "[MASK]"
+        if token_id == 1:
+            return "[OOV]"
+        if token_id == 2:
+            return "[NULL]"
+
+        value = token_id - VALUES_START
+
+        if feat.name == "hour":
+            if 0 <= value <= 23:
+                h = value % 12 or 12
+                ampm = "AM" if value < 12 else "PM"
+                return f"{h} {ampm}"
+            return f"H{value}"
+
+        if feat.name == "dow":
+            return DOW_NAMES[value] if 0 <= value < 7 else f"D{value}"
+
+        if feat.name == "days_since_last":
+            if value == 0:
+                return "Same day"
+            if value <= 5:
+                return f"{value}d ago"
+            bucket_size = 365 / 30
+            approx = int(value * bucket_size)
+            return f"~{approx}d ago"
+
+        if feat.name == "is_recurring":
+            return "Recurring" if value == 1 else "One-time"
+
+        if feat.name == "mcc":
+            return MCC_NAMES.get(value, f"MCC-{value}")
+
+        if feat.name == "merchant_id":
+            info = self.merchants.get(value)
+            return info.name
+
+        if feat.name == "customer_merchant_count":
+            if value == 0:
+                return "1st visit"
+            if value < 5:
+                return f"{value + 1} visits"
+            bucket_size = 500 / 20
+            approx = int(value * bucket_size)
+            return f"~{approx} visits"
+
+        if feat.name == "entry_mode":
+            if feat.values and value in feat.values:
+                raw = feat.values[value]
+                return ENTRY_MODE_DISPLAY.get(raw, raw)
+            return f"Entry-{value}"
+
+        if feat.name == "amount":
+            range_idx = value // 16
+            return AMOUNT_RANGE_LABELS.get(range_idx, f"${value}")
+
+        if feat.name == "card_product":
+            if feat.values and value in feat.values:
+                raw = feat.values[value]
+                return raw.replace("_", " ").title()
+            return f"Card-{value}"
+
+        if feat.name == "country":
+            return COUNTRY_NAMES[value] if 0 <= value < len(COUNTRY_NAMES) else f"Country-{value}"
+
+        if feat.name == "avs":
+            if feat.values and value in feat.values:
+                raw = feat.values[value]
+                return AVS_DISPLAY.get(raw, raw)
+            return f"AVS-{value}"
+
+        if feat.name == "cvv":
+            if feat.values and value in feat.values:
+                raw = feat.values[value]
+                return CVV_DISPLAY.get(raw, raw)
+            return f"CVV-{value}"
+
+        if feat.name == "device_hash":
+            return f"Device-{value:04d}"
+
+        if feat.name == "customer_tenure":
+            months = value * 12
+            if months < 12:
+                return f"<1 year"
+            return f"~{months // 12}yr"
+
+        return str(value)
+
+    def decode_sequence(self, token_ids: np.ndarray) -> list[DecodedTransaction]:
+        """Decode a full (T, F) sequence into readable transactions.
+
+        Returns transactions in reverse chronological order (most recent first).
+        """
+        num_tx = token_ids.shape[0]
+        txns: list[DecodedTransaction] = []
+
+        for t in range(num_tx - 1, -1, -1):
+            row = token_ids[t]
+
+            merchant_val = int(row[5]) - VALUES_START
+            merchant_info = self.merchants.get(max(0, merchant_val))
+
+            txn = DecodedTransaction(
+                index=t,
+                hour=self._decode_value(0, int(row[0])),
+                dow=self._decode_value(1, int(row[1])),
+                days_since_last=self._decode_value(2, int(row[2])),
+                is_recurring=self._decode_value(3, int(row[3])),
+                mcc=self._decode_value(4, int(row[4])),
+                merchant_name=self._decode_value(5, int(row[5])),
+                merchant_category=merchant_info.category,
+                customer_merchant_count=self._decode_value(6, int(row[6])),
+                entry_mode=self._decode_value(7, int(row[7])),
+                amount_range=self._decode_value(8, int(row[8])),
+                card_product=self._decode_value(9, int(row[9])),
+                country=self._decode_value(10, int(row[10])),
+                avs=self._decode_value(11, int(row[11])),
+                cvv=self._decode_value(12, int(row[12])),
+                device_hash=self._decode_value(13, int(row[13])),
+                customer_tenure=self._decode_value(14, int(row[14])),
+            )
+            txns.append(txn)
+
+        return txns
+
+    def summarize_customer(
+        self, token_ids: np.ndarray, is_fraud: bool,
+    ) -> str:
+        """One-line behavioral summary derived from actual token data."""
+        num_tx = token_ids.shape[0]
+
+        merchant_ids = set()
+        mcc_counts: dict[str, int] = {}
+        for t in range(num_tx):
+            mid = int(token_ids[t, 5]) - VALUES_START
+            merchant_ids.add(mid)
+            mcc_val = int(token_ids[t, 4]) - VALUES_START
+            cat = MCC_NAMES.get(mcc_val, f"Cat-{mcc_val}")
+            mcc_counts[cat] = mcc_counts.get(cat, 0) + 1
+
+        top_cats = sorted(mcc_counts.items(), key=lambda x: -x[1])[:3]
+        cat_str = ", ".join(c[0] for c in top_cats)
+
+        label = "FRAUD" if is_fraud else "Legitimate"
+        return (
+            f"{num_tx} transactions | {len(merchant_ids)} unique merchants | "
+            f"Top categories: {cat_str} | Label: {label}"
+        )

src/demo/merchant_catalog.py

@@ -0,0 +1,177 @@
+"""Deterministic merchant name catalog for the inference demo.
+
+Maps 10K merchant_id values to plausible business names based on their
+MCC assignments from the data generator. Uses a seeded RNG so the same
+merchant always gets the same name across runs.
+
+The catalog is built by:
+  1. Rebuilding the same MerchantCatalog the generator used (same seed=42)
+  2. Assigning human-readable names from per-MCC name pools
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from src.data.generator import MerchantCatalog
+from src.data.schema import SchemaConfig, load_schema
+
+
+MCC_NAMES: dict[int, str] = {
+    0: "Transit", 1: "Coffee Shop", 2: "Gas Station", 3: "Parking",
+    4: "Convenience Store", 5: "Discount Store", 6: "Dollar Store",
+    7: "Vending", 8: "Laundromat", 9: "Dry Cleaner",
+    10: "Airlines", 11: "Hotels", 12: "Car Rental", 13: "Cruise Lines",
+    14: "Travel Agency", 15: "Taxi/Rideshare", 16: "Railway", 17: "Bus Lines",
+    18: "Tolls", 19: "Freight/Shipping",
+    20: "Amazon/Online Retail", 21: "Electronics Online", 22: "Clothing Online",
+    23: "Home Goods Online", 24: "Marketplace",
+    25: "Books/Media", 26: "Beauty/Cosmetics", 27: "Pet Supplies",
+    28: "Sporting Goods", 29: "Toys/Games",
+    30: "Casual Dining", 31: "Fast Casual", 32: "Fine Dining",
+    33: "Bars/Pubs", 34: "Bakery/Cafe",
+    35: "Catering", 36: "Food Delivery", 37: "Meal Kit",
+    38: "Ice Cream/Dessert", 39: "Food Truck",
+    40: "Luxury Goods", 41: "Department Store", 42: "Jewelry",
+    43: "Watches", 44: "Designer Fashion",
+    45: "Clothing Retail", 46: "Shoes", 47: "Accessories",
+    48: "Sportswear", 49: "Thrift/Vintage",
+    50: "Office Supplies", 51: "Wholesale Club", 52: "Industrial Supply",
+    53: "Printing/Copy", 54: "Packaging",
+    55: "Legal Services", 56: "Accounting", 57: "Consulting",
+    58: "Marketing/Advertising", 59: "Temp Staffing",
+    60: "Grocery", 61: "Supermarket", 62: "Pharmacy",
+    63: "Health Food", 64: "Liquor Store",
+    65: "Doctor/Medical", 66: "Dentist", 67: "Gym/Fitness",
+    68: "Spa/Wellness", 69: "Veterinary",
+    70: "Fast Food", 71: "Pizza", 72: "Movie Theater",
+    73: "Amusement Park", 74: "Arcade/Gaming",
+    75: "Streaming/Digital", 76: "Music/Audio", 77: "News/Media",
+    78: "Cloud/SaaS", 79: "App Store",
+    80: "Subscription Box", 81: "Software License", 82: "Web Hosting",
+    83: "Domain/DNS", 84: "VPN/Security",
+    85: "Auto Parts", 86: "Auto Service", 87: "Car Wash",
+    88: "Auto Insurance", 89: "Dealership",
+    90: "Hardware Store", 91: "Furniture", 92: "Appliances",
+    93: "Garden/Nursery", 94: "Flooring/Tile",
+    95: "University/College", 96: "Online Course", 97: "Tutoring",
+    98: "Government", 99: "Utilities",
+}
+
+MERCHANT_NAME_POOLS: dict[int, list[str]] = {
+    0: ["Metro Transit", "City Bus", "Subway Pass", "BART", "TransitGo"],
+    1: ["Starbucks", "Dunkin'", "Peet's Coffee", "Blue Bottle", "Tim Hortons"],
+    2: ["Shell", "Chevron", "BP", "ExxonMobil", "Costco Gas"],
+    3: ["ParkMobile", "SpotHero", "LAZ Parking", "ParkWhiz", "Ace Parking"],
+    10: ["Delta Airlines", "United Airlines", "Southwest", "JetBlue", "American Airlines"],
+    11: ["Marriott", "Hilton", "Hyatt", "IHG Hotels", "Best Western"],
+    12: ["Enterprise", "Hertz", "Avis", "National Car", "Budget Rent"],
+    20: ["Amazon", "eBay", "Etsy", "Walmart.com", "Target.com"],
+    21: ["Best Buy", "Newegg", "B&H Photo", "Apple Store", "Samsung"],
+    22: ["SHEIN", "ASOS", "Nordstrom", "H&M Online", "Zara Online"],
+    30: ["Applebee's", "Chili's", "Olive Garden", "Red Lobster", "TGI Friday's"],
+    31: ["Chipotle", "Panera Bread", "Shake Shack", "Sweetgreen", "Cava"],
+    32: ["Ruth's Chris", "Capital Grille", "Nobu", "Fleming's", "Morton's"],
+    33: ["Local Bar & Grill", "Sports Bar", "Craft Brewery", "Wine Bar", "Dive Bar"],
+    40: ["Louis Vuitton", "Gucci", "Tiffany & Co", "Hermès", "Cartier"],
+    41: ["Macy's", "Nordstrom", "Bloomingdale's", "Saks Fifth", "Neiman Marcus"],
+    45: ["Zara", "H&M", "Uniqlo", "Gap", "Old Navy"],
+    46: ["Nike", "Foot Locker", "DSW", "Adidas", "New Balance"],
+    48: ["Lululemon", "REI", "Dick's Sporting", "Under Armour", "Patagonia"],
+    50: ["Staples", "Office Depot", "Amazon Business", "Uline", "W.B. Mason"],
+    51: ["Costco", "Sam's Club", "BJ's Wholesale", "Restaurant Depot", "Sysco"],
+    55: ["Smith & Associates", "Baker Law", "Davis Legal", "Park LLP", "Cohen Firm"],
+    60: ["Walmart", "Kroger", "Safeway", "Whole Foods", "Trader Joe's"],
+    61: ["Albertsons", "Publix", "H-E-B", "Wegmans", "Aldi"],
+    62: ["CVS", "Walgreens", "Rite Aid", "Pharmacy Plus", "MedShop"],
+    67: ["Planet Fitness", "LA Fitness", "Equinox", "CrossFit", "Gold's Gym"],
+    70: ["McDonald's", "Burger King", "Wendy's", "Taco Bell", "KFC"],
+    71: ["Domino's", "Papa John's", "Pizza Hut", "Little Caesars", "Marco's"],
+    72: ["AMC Theatres", "Regal Cinema", "Cinemark", "Alamo Draft", "IMAX"],
+    75: ["Netflix", "Spotify", "Disney+", "YouTube Premium", "Hulu"],
+    76: ["Apple Music", "Tidal", "SoundCloud", "Audible", "Pandora"],
+    78: ["AWS", "Google Cloud", "Azure", "Salesforce", "Slack"],
+    80: ["HelloFresh", "BarkBox", "FabFitFun", "Birchbox", "Dollar Shave"],
+    81: ["Microsoft 365", "Adobe CC", "JetBrains", "Zoom Pro", "Notion"],
+    85: ["AutoZone", "O'Reilly Auto", "NAPA", "Advance Auto", "Pep Boys"],
+    86: ["Jiffy Lube", "Meineke", "Firestone", "Midas", "Valvoline"],
+    90: ["Home Depot", "Lowe's", "Ace Hardware", "Menards", "True Value"],
+    91: ["IKEA", "Wayfair", "Pottery Barn", "Crate & Barrel", "West Elm"],
+    95: ["State University", "Community College", "Coursera", "edX", "Udemy"],
+    96: ["Masterclass", "Skillshare", "LinkedIn Learning", "Pluralsight", "Codecademy"],
+    98: ["IRS", "DMV", "City Hall", "Court Fees", "Passport Office"],
+    99: ["Electric Co", "Water Utility", "Gas Utility", "Internet/Cable", "Phone Bill"],
+}
+
+
+@dataclass
+class MerchantInfo:
+    """Human-readable merchant identity."""
+
+    merchant_id: int
+    name: str
+    mcc_code: int
+    category: str
+
+
+class DemoMerchantCatalog:
+    """Maps merchant_id -> human-readable name + category.
+
+    Deterministic: same seed always produces the same mapping.
+    """
+
+    def __init__(self, schema: SchemaConfig, seed: int = 42) -> None:
+        rng = np.random.default_rng(seed)
+        n_merchants = schema.get_feature("merchant_id").num_values
+        n_mccs = schema.get_feature("mcc").num_values
+
+        catalog = MerchantCatalog.build(n_merchants, n_mccs, rng)
+
+        self._merchant_map: dict[int, MerchantInfo] = {}
+        name_rng = np.random.default_rng(seed + 1000)
+
+        mcc_counters: dict[int, int] = {}
+        for mid in range(n_merchants):
+            mcc = int(catalog.mcc_assignments[mid])
+            mcc_counters[mcc] = mcc_counters.get(mcc, 0) + 1
+
+            pool = MERCHANT_NAME_POOLS.get(mcc)
+            category = MCC_NAMES.get(mcc, f"Category {mcc}")
+
+            if pool:
+                idx = name_rng.integers(0, len(pool))
+                base_name = pool[idx]
+                count = mcc_counters[mcc]
+                if count <= len(pool):
+                    name = pool[(count - 1) % len(pool)]
+                else:
+                    name = f"{base_name} #{count}"
+            else:
+                name = f"{category} #{mcc_counters[mcc]}"
+
+            self._merchant_map[mid] = MerchantInfo(
+                merchant_id=mid,
+                name=name,
+                mcc_code=mcc,
+                category=category,
+            )
+
+    def get(self, merchant_id: int) -> MerchantInfo:
+        info = self._merchant_map.get(merchant_id)
+        if info is None:
+            return MerchantInfo(merchant_id, f"Merchant #{merchant_id}", -1, "Unknown")
+        return info
+
+    def get_name(self, merchant_id: int) -> str:
+        return self.get(merchant_id).name
+
+    def get_category(self, merchant_id: int) -> str:
+        return self.get(merchant_id).category
+
+    @classmethod
+    def from_schema_path(cls, schema_path: str | Path) -> DemoMerchantCatalog:
+        schema = load_schema(schema_path)
+        return cls(schema)

src/demo/profile_inference.py

@@ -0,0 +1,325 @@
+"""Infer behavioral profile from transaction sequence features.
+
+Analyzes MCC distribution, temporal patterns, entry modes, and amount
+patterns to classify a customer into one of 16 behavioral archetypes.
+This gives meetings a "we know what the model is reasoning about" beat.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from src.data.schema import VALUES_START
+from src.demo.merchant_catalog import MCC_NAMES
+
+
+@dataclass
+class ProfileMatch:
+    """Inferred behavioral profile with confidence signals."""
+
+    name: str
+    description: str
+    confidence: float
+    evidence: list[str]
+
+
+PROFILE_SIGNATURES: list[dict[str, object]] = [
+    {
+        "name": "Commuter",
+        "description": "Daily transit, coffee, and gas with predictable weekday patterns",
+        "mcc_signal": {0, 1, 2, 3},
+        "hour_signal": {7, 8, 9, 17, 18},
+        "dow_signal": "weekday_heavy",
+        "entry_mode_signal": "contactless",
+    },
+    {
+        "name": "Traveler",
+        "description": "Frequent flyer with hotels, restaurants, and multi-country activity",
+        "mcc_signal": {10, 11, 12, 13, 14},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": None,
+        "country_signal": "multi_country",
+    },
+    {
+        "name": "Online Shopper",
+        "description": "E-commerce heavy with evening browsing and consistent device usage",
+        "mcc_signal": {20, 21, 22, 23, 24, 25},
+        "hour_signal": {20, 21, 22, 23},
+        "dow_signal": None,
+        "entry_mode_signal": "card_not_present",
+    },
+    {
+        "name": "Restaurant Regular",
+        "description": "Lunch and dinner out at local restaurants, moderate spend",
+        "mcc_signal": {30, 31, 32, 33, 34},
+        "hour_signal": {12, 13, 18, 19, 20},
+        "dow_signal": None,
+        "entry_mode_signal": None,
+    },
+    {
+        "name": "Luxury Spender",
+        "description": "High-value purchases at department stores and luxury retailers",
+        "mcc_signal": {40, 41, 42, 43, 44},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": None,
+        "amount_signal": "high",
+    },
+    {
+        "name": "Small Business Owner",
+        "description": "Office supplies, wholesale, business services during work hours",
+        "mcc_signal": {50, 51, 52, 53, 54, 55, 56, 57},
+        "hour_signal": {9, 10, 11, 14, 15, 16},
+        "dow_signal": "weekday_heavy",
+        "entry_mode_signal": None,
+    },
+    {
+        "name": "Retiree",
+        "description": "Daytime errands: groceries, pharmacy, low-moderate spend",
+        "mcc_signal": {60, 61, 62, 63},
+        "hour_signal": {9, 10, 11, 14, 15},
+        "dow_signal": None,
+        "entry_mode_signal": None,
+        "amount_signal": "low",
+    },
+    {
+        "name": "Student",
+        "description": "Fast food, entertainment, streaming, small online purchases",
+        "mcc_signal": {70, 71, 72, 73, 74, 75},
+        "hour_signal": {11, 12, 20, 21, 22, 23},
+        "dow_signal": None,
+        "entry_mode_signal": None,
+        "amount_signal": "low",
+    },
+    {
+        "name": "Gig Worker",
+        "description": "Irregular hours, frequent gas and food, high transaction frequency",
+        "mcc_signal": {2, 30, 31, 70, 71},
+        "hour_signal": None,
+        "dow_signal": "all_days",
+        "entry_mode_signal": "contactless",
+        "frequency_signal": "high",
+    },
+    {
+        "name": "Suburban Parent",
+        "description": "Family groceries, kids activities, school supplies, weekend patterns",
+        "mcc_signal": {60, 61, 28, 90, 91},
+        "hour_signal": {16, 17, 18, 19},
+        "dow_signal": "weekend_spike",
+        "entry_mode_signal": None,
+    },
+    {
+        "name": "Digital Nomad",
+        "description": "Remote worker across countries: coworking, cafes, SaaS tools",
+        "mcc_signal": {1, 34, 78, 81, 82},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": "card_not_present",
+        "country_signal": "multi_country",
+    },
+    {
+        "name": "Healthcare Worker",
+        "description": "Shift-based schedule: early morning and late evening clusters",
+        "mcc_signal": {62, 65, 66},
+        "hour_signal": {5, 6, 7, 22, 23},
+        "dow_signal": "all_days",
+        "entry_mode_signal": None,
+    },
+    {
+        "name": "Subscription Maximalist",
+        "description": "Heavy recurring: streaming, SaaS, meal kits, gym memberships",
+        "mcc_signal": {75, 76, 77, 78, 80, 81},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": "card_not_present",
+        "recurring_signal": "high",
+    },
+    {
+        "name": "High Net Worth",
+        "description": "Infrequent large transactions: fine dining, luxury hotels, international",
+        "mcc_signal": {32, 11, 40, 42},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": None,
+        "amount_signal": "very_high",
+        "country_signal": "multi_country",
+    },
+    {
+        "name": "Seasonal Worker",
+        "description": "Concentrated spending bursts, basic cards, domestic only",
+        "mcc_signal": {60, 2, 70},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": None,
+        "amount_signal": "low",
+    },
+    {
+        "name": "Urban Cashless",
+        "description": "City dweller: transit, delivery, convenience, many small contactless taps",
+        "mcc_signal": {0, 1, 36, 70, 15},
+        "hour_signal": None,
+        "dow_signal": None,
+        "entry_mode_signal": "contactless",
+        "amount_signal": "low",
+        "frequency_signal": "high",
+    },
+]
+
+
+def infer_profile(token_ids: np.ndarray) -> ProfileMatch:
+    """Analyze a (T=64, F=15) sequence and infer the most likely behavioral profile."""
+    num_tx = token_ids.shape[0]
+
+    hours = token_ids[:, 0] - VALUES_START
+    dows = token_ids[:, 1] - VALUES_START
+    mccs = token_ids[:, 4] - VALUES_START
+    entry_modes = token_ids[:, 7] - VALUES_START
+    amounts = token_ids[:, 8] - VALUES_START
+    countries = token_ids[:, 10] - VALUES_START
+    recurring = token_ids[:, 3] - VALUES_START
+
+    mcc_counts: dict[int, int] = {}
+    for m in mccs:
+        mcc_counts[int(m)] = mcc_counts.get(int(m), 0) + 1
+
+    hour_counts = np.bincount(np.clip(hours, 0, 23), minlength=24)
+    dow_counts = np.bincount(np.clip(dows, 0, 6), minlength=7)
+    entry_counts = np.bincount(np.clip(entry_modes, 0, 4), minlength=5)
+
+    mean_amount = float(np.mean(amounts))
+    num_countries = len(set(int(c) for c in countries))
+    recurring_rate = float(np.mean(recurring == 1))
+    unique_merchants = len(set(int(token_ids[t, 5] - VALUES_START) for t in range(num_tx)))
+
+    best_score = -1.0
+    best_profile = PROFILE_SIGNATURES[0]
+    best_evidence: list[str] = []
+
+    for sig in PROFILE_SIGNATURES:
+        score = 0.0
+        evidence: list[str] = []
+
+        mcc_signal = sig.get("mcc_signal", set())
+        if mcc_signal:
+            mcc_hits = sum(mcc_counts.get(m, 0) for m in mcc_signal)
+            mcc_frac = mcc_hits / num_tx
+            score += mcc_frac * 4.0
+            if mcc_frac > 0.3:
+                top_mcc = max(mcc_signal, key=lambda m: mcc_counts.get(m, 0))
+                evidence.append(f"{mcc_frac*100:.0f}% transactions at {MCC_NAMES.get(top_mcc, 'key')} merchants")
+
+        hour_signal = sig.get("hour_signal")
+        if hour_signal:
+            hour_hits = sum(int(hour_counts[h]) for h in hour_signal if h < 24)
+            hour_frac = hour_hits / num_tx
+            score += hour_frac * 2.0
+            if hour_frac > 0.4:
+                evidence.append(f"Peak activity at expected hours")
+
+        dow_signal = sig.get("dow_signal")
+        if dow_signal == "weekday_heavy":
+            weekday_frac = float(dow_counts[:5].sum()) / num_tx
+            if weekday_frac > 0.75:
+                score += 1.5
+                evidence.append(f"{weekday_frac*100:.0f}% weekday transactions")
+        elif dow_signal == "weekend_spike":
+            weekend_frac = float(dow_counts[5:].sum()) / num_tx
+            if weekend_frac > 0.35:
+                score += 1.5
+                evidence.append(f"High weekend activity ({weekend_frac*100:.0f}%)")
+        elif dow_signal == "all_days":
+            if dow_counts.min() > num_tx * 0.08:
+                score += 1.0
+
+        entry_signal = sig.get("entry_mode_signal")
+        if entry_signal == "contactless":
+            if entry_counts[2] > num_tx * 0.3:
+                score += 1.5
+                evidence.append(f"Contactless-heavy ({entry_counts[2]*100//num_tx}%)")
+        elif entry_signal == "card_not_present":
+            if entry_counts[1] > num_tx * 0.4:
+                score += 1.5
+                evidence.append(f"Primarily online ({entry_counts[1]*100//num_tx}%)")
+
+        amount_signal = sig.get("amount_signal")
+        if amount_signal == "low" and mean_amount < 60:
+            score += 1.0
+            evidence.append(f"Low avg spend")
+        elif amount_signal == "high" and mean_amount > 120:
+            score += 1.5
+            evidence.append(f"High avg spend")
+        elif amount_signal == "very_high" and mean_amount > 160:
+            score += 2.0
+            evidence.append(f"Very high avg spend")
+
+        country_signal = sig.get("country_signal")
+        if country_signal == "multi_country" and num_countries > 3:
+            score += 2.0
+            evidence.append(f"{num_countries} countries visited")
+
+        recurring_sig = sig.get("recurring_signal")
+        if recurring_sig == "high" and recurring_rate > 0.3:
+            score += 1.5
+            evidence.append(f"{recurring_rate*100:.0f}% recurring")
+
+        freq_signal = sig.get("frequency_signal")
+        if freq_signal == "high" and unique_merchants > 15:
+            score += 1.0
+
+        if score > best_score:
+            best_score = score
+            best_profile = sig
+            best_evidence = evidence
+
+    max_possible = 10.0
+    confidence = min(1.0, best_score / max_possible)
+
+    if not best_evidence:
+        top_mcc = max(mcc_counts, key=mcc_counts.get) if mcc_counts else 0
+        best_evidence.append(f"Primary category: {MCC_NAMES.get(top_mcc, 'Unknown')}")
+
+    return ProfileMatch(
+        name=str(best_profile["name"]),
+        description=str(best_profile["description"]),
+        confidence=confidence,
+        evidence=best_evidence,
+    )
+
+
+def format_profile_html(match: ProfileMatch) -> str:
+    """Render profile match using Liquid design tokens."""
+    conf_pct = match.confidence * 100
+    if conf_pct > 60:
+        conf_color = "#10B981"
+    elif conf_pct > 30:
+        conf_color = "#F59E0B"
+    else:
+        conf_color = "#737373"
+
+    evidence_items = "".join(
+        f'<li style="margin-bottom: 2px;">{e}</li>' for e in match.evidence
+    )
+    _mono = "JetBrains Mono, ui-monospace, SFMono-Regular, monospace"
+
+    return f"""
+    <div style="padding: 12px 14px; background: #ffffff; border-radius: 12px;
+                border: 1px solid rgba(0,0,0,0.1); border-left: 3px solid {conf_color};
+                box-shadow: 0 1px 3px rgba(0,0,0,0.04);">
+        <div style="font-size: 16px; font-weight: 600; color: #171717;
+                    letter-spacing: -0.01em; margin-bottom: 2px;">
+            {match.name}
+        </div>
+        <div style="font-size: 12px; color: #525252; margin-bottom: 6px;">
+            {match.description}
+        </div>
+        <div style="font-family: {_mono}; font-size: 11px; color: #737373; margin-bottom: 4px;">
+            confidence: <span style="color: {conf_color}; font-weight: 600;">{conf_pct:.0f}%</span>
+        </div>
+        <ul style="font-size: 11px; color: #525252; margin: 0; padding-left: 14px; line-height: 1.5;">
+            {evidence_items}
+        </ul>
+    </div>
+    """

src/demo/render.py

@@ -0,0 +1,833 @@
+"""HTML rendering for the inference demo UI.
+
+Styling follows the Liquid AI design system from liquid-lfm-cloud:
+  - Monochrome primary scale (#171717 primary, #f5f5f5 background)
+  - Semantic accents: emerald (#10B981), amber (#F59E0B), red (#EF4444), blue (#3B82F6)
+  - JetBrains Mono for technical/metric elements, system sans-serif for body
+  - 16px rounded cards with subtle shadows
+  - Pill-shaped buttons and badges
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+
+from src.data.generator import AMOUNT_RANGE_LABELS
+from src.data.schema import VALUES_START
+from src.demo.decode import TransactionDecoder
+from src.demo.merchant_catalog import DemoMerchantCatalog, MCC_NAMES
+
+# Liquid design tokens — light mode (liquid-lfm-cloud design system)
+_BG = "#f5f5f5"
+_BG_CARD = "#ffffff"
+_BG_CARD_ALT = "#fafafa"
+_BORDER = "rgba(0,0,0,0.1)"
+_BORDER_SUBTLE = "rgba(0,0,0,0.05)"
+_TEXT = "#171717"
+_TEXT_MUTED = "#525252"
+_TEXT_DIM = "#737373"
+_ACCENT_BLUE = "#3B82F6"
+_ACCENT_GREEN = "#10B981"
+_ACCENT_AMBER = "#F59E0B"
+_ACCENT_RED = "#EF4444"
+_RADIUS_CARD = "16px"
+_RADIUS_SM = "8px"
+_FONT_MONO = "JetBrains Mono, ui-monospace, SFMono-Regular, monospace"
+_FONT_SANS = "-apple-system, BlinkMacSystemFont, Segoe UI, Roboto, sans-serif"
+
+
+def format_fraud_score(prob: float) -> str:
+    """Format fraud probability as colored gauge bar."""
+    pct = prob * 100
+    if pct < 20:
+        color = _ACCENT_GREEN
+        risk = "LOW RISK"
+    elif pct < 60:
+        color = _ACCENT_AMBER
+        risk = "MEDIUM RISK"
+    else:
+        color = _ACCENT_RED
+        risk = "HIGH RISK"
+
+    bar_width = max(2, min(100, int(pct)))
+    return f"""
+    <div style="margin: 8px 0;">
+        <div style="font-family: {_FONT_MONO}; font-size: 24px; font-weight: 600;
+                    color: {color}; margin-bottom: 4px; letter-spacing: -0.02em;">
+            {pct:.1f}%
+            <span style="font-size: 12px; font-weight: 500; opacity: 0.8;
+                         letter-spacing: 0.05em;">{risk}</span>
+        </div>
+        <div style="background: #e5e5e5; border-radius: 9999px; height: 8px;
+                    width: 100%; overflow: hidden;">
+            <div style="background: {color}; height: 100%; width: {bar_width}%;
+                        border-radius: 9999px; transition: width 0.3s ease;"></div>
+        </div>
+    </div>
+    """
+
+
+def format_topk_predictions(
+    probs: np.ndarray,
+    k: int,
+    label_fn: Callable[[int], str],
+) -> str:
+    """Format top-k predictions as styled rows with probability bars."""
+    top_indices = np.argsort(probs)[::-1][:k]
+
+    rows = ""
+    for i, idx in enumerate(top_indices):
+        p = probs[idx] * 100
+        label = label_fn(int(idx))
+        bar_width = max(2, int(p * 2.5))
+        opacity = 1.0 - i * 0.15
+        weight = "600" if i == 0 else "400"
+        # Show "<0.1%" for probabilities that round to 0.0% at 1 decimal place.
+        # This happens with high-cardinality heads (10K merchants) when the model
+        # hasn't learned a meaningful distribution (e.g. random-init baseline).
+        p_str = f"{p:.1f}%" if p >= 0.05 else "<0.1%"
+        rows += f"""
+        <div style="display: flex; align-items: center; gap: 8px; padding: 4px 0;">
+            <div style="flex: 1; font-size: 13px; font-weight: {weight};
+                        color: {_TEXT}; opacity: {opacity};">{label}</div>
+            <div style="width: 50px; text-align: right; font-family: {_FONT_MONO};
+                        font-size: 12px; color: {_TEXT_MUTED};">{p_str}</div>
+            <div style="width: 120px;">
+                <div style="background: {_ACCENT_BLUE}; height: 6px; width: {bar_width}%;
+                            border-radius: 9999px; opacity: {opacity};"></div>
+            </div>
+        </div>
+        """
+
+    return f'<div style="padding: 4px 0;">{rows}</div>'
+
+
+def format_merchant_predictions(
+    probs: np.ndarray,
+    merchant_catalog: DemoMerchantCatalog,
+    k: int = 5,
+) -> str:
+    """Top-k merchant predictions with names and categories."""
+    def label_fn(idx: int) -> str:
+        if idx < VALUES_START:
+            return "[special]"
+        mid = idx - VALUES_START
+        info = merchant_catalog.get(mid)
+        return f"{info.name} ({info.category})"
+
+    return format_topk_predictions(probs, k, label_fn)
+
+
+def format_amount_predictions(probs: np.ndarray, k: int = 5) -> str:
+    """Top-k amount range predictions."""
+    def label_fn(idx: int) -> str:
+        return AMOUNT_RANGE_LABELS.get(idx, f"Range {idx}")
+
+    return format_topk_predictions(probs, k, label_fn)
+
+
+def format_mcc_predictions(probs: np.ndarray, k: int = 5) -> str:
+    """Top-k MCC predictions with category names."""
+    def label_fn(idx: int) -> str:
+        if idx < VALUES_START:
+            return "[special]"
+        mcc_val = idx - VALUES_START
+        return MCC_NAMES.get(mcc_val, f"MCC-{mcc_val}")
+
+    return format_topk_predictions(probs, k, label_fn)
+
+
+def format_timeline(decoder: TransactionDecoder, token_ids: np.ndarray) -> str:
+    """Render transaction sequence as a scrollable table."""
+    txns = decoder.decode_sequence(token_ids)
+
+    header = f"""
+    <div style="max-height: 380px; overflow-y: auto; border: 1px solid {_BORDER};
+                border-radius: {_RADIUS_CARD}; background: {_BG_CARD};">
+    <table style="width: 100%; border-collapse: collapse; font-size: 12px;
+                  font-family: {_FONT_MONO};">
+        <thead>
+            <tr style="border-bottom: 1px solid {_BORDER}; position: sticky; top: 0;
+                       background: {_BG_CARD}; z-index: 1;">
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase; letter-spacing: 0.05em;">Tx</th>
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase;">When</th>
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase;">Merchant</th>
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase;">Category</th>
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase;">Amount</th>
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase;">Method</th>
+                <th style="padding: 8px; text-align: left; color: {_TEXT_DIM};
+                           font-size: 10px; text-transform: uppercase;">Country</th>
+            </tr>
+        </thead>
+        <tbody>
+    """
+
+    rows = ""
+    for txn in txns:
+        when = f"{txn.dow} {txn.hour}"
+        highlight = f"background: rgba(59, 130, 246, 0.08);" if txn.index == 63 else ""
+        rows += f"""
+        <tr style="border-bottom: 1px solid {_BORDER_SUBTLE}; {highlight}">
+            <td style="padding: 6px 8px; color: {_TEXT_DIM};">{txn.index}</td>
+            <td style="padding: 6px 8px; color: {_TEXT_MUTED};">{when}</td>
+            <td style="padding: 6px 8px; color: {_TEXT}; font-weight: 500;">{txn.merchant_name}</td>
+            <td style="padding: 6px 8px; color: {_TEXT_DIM};">{txn.merchant_category}</td>
+            <td style="padding: 6px 8px; color: {_TEXT_MUTED};">{txn.amount_range}</td>
+            <td style="padding: 6px 8px; color: {_TEXT_DIM};">{txn.entry_mode}</td>
+            <td style="padding: 6px 8px; color: {_TEXT_DIM};">{txn.country}</td>
+        </tr>
+        """
+
+    return header + rows + "</tbody></table></div>"
+
+
+def render_production_architecture() -> str:
+    """LFM2.5 production architecture and how it applies to payments."""
+    _purple = "#7c3aed"
+    _purple_bg = "rgba(124,58,237,0.08)"
+    _purple_border = "rgba(124,58,237,0.25)"
+
+    def _layer_cell(label: str, idx: int, is_attn: bool) -> str:
+        bg = _purple_bg if is_attn else "rgba(16,185,129,0.08)"
+        bc = _purple_border if is_attn else "rgba(16,185,129,0.25)"
+        color = _purple if is_attn else _ACCENT_GREEN
+        return f"""<div style="flex: 1; padding: 6px 2px; background: {bg};
+            border: 1px solid {bc}; border-radius: 4px; text-align: center; min-width: 0;">
+            <div style="font-family: {_FONT_MONO}; font-size: 8px; color: {color};
+                        font-weight: 600;">{label}</div>
+            <div style="font-family: {_FONT_MONO}; font-size: 7px; color: {_TEXT_DIM};">L{idx}</div>
+        </div>"""
+
+    layers_1_2b = [
+        ("C", 0, False), ("C", 1, False), ("A", 2, True), ("C", 3, False),
+        ("C", 4, False), ("A", 5, True), ("C", 6, False), ("C", 7, False),
+        ("A", 8, True), ("C", 9, False), ("A", 10, True), ("C", 11, False),
+        ("A", 12, True), ("C", 13, False), ("A", 14, True), ("C", 15, False),
+    ]
+    layer_cells = "".join(_layer_cell(l, i, a) for l, i, a in layers_1_2b)
+
+    return f"""
+    <div style="max-width: 1100px; margin: 0 auto; padding: 16px;">
+
+      <!-- Header -->
+      <h2 style="margin: 0 0 4px 0; color: {_TEXT}; font-size: 22px; font-weight: 700;
+                 letter-spacing: -0.02em;">
+        LFM2.5 for Payment Sequences
+      </h2>
+      <p style="color: {_TEXT_DIM}; font-size: 13px; margin: 0 0 24px 0; line-height: 1.5;">
+        The published LFM2.5 architecture adapted for structured transaction data.
+        Same hybrid conv-attention backbone. New per-feature embedding layer.
+      </p>
+
+      <!-- Config comparison table -->
+      <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-bottom: 24px;">
+        <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      text-transform: uppercase; letter-spacing: 0.05em; margin-bottom: 10px;">
+            Production Target</div>
+          <table style="width: 100%; font-size: 12px; border-collapse: collapse;">
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Parameters</td>
+                <td style="padding: 3px 0; color: {_TEXT}; font-weight: 600; text-align: right;">1.2B</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Hidden dim</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">2048</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Layers</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">
+                  <span style="color: {_ACCENT_GREEN};">10 conv</span> + <span style="color: {_purple};">6 attn</span></td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Attention</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">32Q / 8KV (GQA)</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">MLP</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">SwiGLU 12288</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Sequence</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">128 tx &times; 30 feat = 3,840</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Latency target</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_ACCENT_GREEN}; font-weight: 600; text-align: right;">&lt; 50ms (H100)</td></tr>
+          </table>
+        </div>
+
+        <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      text-transform: uppercase; letter-spacing: 0.05em; margin-bottom: 10px;">
+            This Demo (Reference)</div>
+          <table style="width: 100%; font-size: 12px; border-collapse: collapse;">
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Parameters</td>
+                <td style="padding: 3px 0; color: {_TEXT}; font-weight: 600; text-align: right;">9.8M</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Hidden dim</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">256</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Layers</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">
+                  <span style="color: {_ACCENT_GREEN};">5 conv</span> + <span style="color: {_purple};">3 attn</span></td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Attention</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">4Q / 2KV (GQA)</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">MLP</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">SwiGLU 1024</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Sequence</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">64 tx &times; 15 feat = 960</td></tr>
+            <tr><td style="padding: 3px 0; color: {_TEXT_DIM};">Measured latency</td>
+                <td style="padding: 3px 0; font-family: {_FONT_MONO}; color: {_TEXT}; text-align: right;">&lt; 80ms (CPU)</td></tr>
+          </table>
+        </div>
+      </div>
+
+      <!-- LFM2.5-1.2B layer diagram -->
+      <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 24px;">
+        <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                    text-transform: uppercase; letter-spacing: 0.05em; margin-bottom: 12px;">
+          LFM2.5-1.2B Layer Pattern (16 layers)</div>
+        <div style="display: flex; gap: 3px; margin-bottom: 8px;">
+          {layer_cells}
+        </div>
+        <div style="display: flex; gap: 16px; font-family: {_FONT_MONO}; font-size: 10px;">
+          <span style="color: {_ACCENT_GREEN};">&#9632; Conv (10): local patterns, k=3, O(n)</span>
+          <span style="color: {_purple};">&#9632; Attention (6): global context, GQA, O(n&sup2;)</span>
+        </div>
+        <div style="margin-top: 8px; font-size: 11px; color: {_TEXT_DIM}; line-height: 1.5;">
+          First and last layers are convolutional. Attention is densest mid-stack.
+          The LM head reads from a local-conv output, but 6 attention layers have already
+          encoded global context upstream.
+        </div>
+      </div>
+
+      <!-- Three key architectural adaptations -->
+      <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 12px; margin-bottom: 24px;">
+
+        <!-- Embedding adaptation -->
+        <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="font-size: 14px; font-weight: 600; color: {_ACCENT_BLUE}; margin-bottom: 6px;">
+            Per-Feature Embedding</div>
+          <p style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5; margin: 0 0 8px 0;">
+            Text LFM2 has one embedding table. Payment LFM2 has one table per feature
+            (hour, merchant, MCC, amount, ...) plus a feature-type table. Summed to produce
+            the same (B, T, D) the backbone expects.</p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;">
+            value_tables[f](token) + type_table(f)</div>
+        </div>
+
+        <!-- KV-cache advantage -->
+        <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="font-size: 14px; font-weight: 600; color: {_ACCENT_GREEN}; margin-bottom: 6px;">
+            KV-Cache Advantage</div>
+          <p style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5; margin: 0 0 8px 0;">
+            Only 6 of 16 layers need KV cache (attention layers only). A pure-attention
+            model at the same depth needs 16 layers of cache. At 3,840 token sequences
+            with 1,024 concurrent requests:</p>
+          <div style="font-family: {_FONT_MONO}; font-size: 11px;">
+            <span style="color: {_ACCENT_GREEN}; font-weight: 600;">LFM2: ~25 GB</span>
+            <span style="color: {_TEXT_DIM};"> vs </span>
+            <span style="color: {_ACCENT_RED};">Pure attn: ~64 GB</span>
+          </div>
+        </div>
+
+        <!-- Weight-tied heads -->
+        <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="font-size: 14px; font-weight: 600; color: {_purple}; margin-bottom: 6px;">
+            Weight-Tied Heads</div>
+          <p style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5; margin: 0 0 8px 0;">
+            Downstream heads that predict pretrained features (next merchant, MCC)
+            project through the backbone's own embedding table. Zero extra parameters,
+            +50% accuracy vs fresh MLP heads.</p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;">
+            adapter(h) @ embedding.weight.T</div>
+        </div>
+      </div>
+
+      <!-- How it translates -->
+      <div style="padding: 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 16px;">
+        <div style="font-size: 14px; font-weight: 600; color: {_TEXT}; margin-bottom: 10px;">
+          How This Translates to Your Implementation</div>
+        <div style="display: grid; grid-template-columns: auto 1fr; gap: 6px 16px; font-size: 12px;">
+          <div style="font-family: {_FONT_MONO}; color: {_ACCENT_BLUE}; font-weight: 600;">Schema</div>
+          <div style="color: {_TEXT_MUTED};">Your features, your vocab sizes, your ordering. We review the schema design.</div>
+          <div style="font-family: {_FONT_MONO}; color: {_ACCENT_BLUE}; font-weight: 600;">Pretrain</div>
+          <div style="color: {_TEXT_MUTED};">Self-supervised on your unlabeled transactions. No fraud labels needed.</div>
+          <div style="font-family: {_FONT_MONO}; color: {_ACCENT_BLUE}; font-weight: 600;">Fine-tune</div>
+          <div style="color: {_TEXT_MUTED};">Attach task heads (fraud, disputes, auth optimization). Multi-task with shared backbone.</div>
+          <div style="font-family: {_FONT_MONO}; color: {_ACCENT_BLUE}; font-weight: 600;">Deploy</div>
+          <div style="color: {_TEXT_MUTED};">Your infra, your GPUs, your compliance. Sub-100ms for real-time authorization decisioning.</div>
+          <div style="font-family: {_FONT_MONO}; color: {_ACCENT_BLUE}; font-weight: 600;">Own</div>
+          <div style="color: {_TEXT_MUTED};">No data leaves your infrastructure. No external API dependency. You own the model.</div>
+        </div>
+      </div>
+
+      <!-- Architecture source -->
+      <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM}; text-align: center;">
+        Architecture: <a href="https://arxiv.org/abs/2511.23404" style="color: {_ACCENT_BLUE}; text-decoration: none;">
+        arXiv 2511.23404</a> &middot;
+        Weights: <a href="https://huggingface.co/LiquidAI" style="color: {_ACCENT_BLUE}; text-decoration: none;">
+        huggingface.co/LiquidAI</a>
+      </div>
+    </div>
+    """
+
+
+def render_comparison_header() -> str:
+    """Header explaining the pretrained vs random-init comparison."""
+    return f"""
+    <div style="padding: 10px 14px; background: rgba(245,158,11,0.06);
+                border: 1px solid rgba(245,158,11,0.2); border-radius: {_RADIUS_SM};
+                margin-bottom: 12px; font-size: 12px; color: {_ACCENT_AMBER}; line-height: 1.5;">
+        <b>Pretrained vs Random Init:</b>
+        <span style="color: {_TEXT_MUTED};">Same architecture, same input, same fine-tuning data.
+        Left: pretrained on 200K unlabeled sequences first. Right: trained from scratch.
+        The difference is the value of self-supervised pretraining.</span>
+    </div>
+    """
+
+
+def render_why_liquid() -> str:
+    """Render the 'Why Liquid AI' value proposition tab.
+
+    Content drawn from docs/architecture-walkthrough.md §14.
+    Uses measured reference-implementation numbers, not marketing claims.
+    """
+    def _table_row(cells: list[str], bold_last: bool = False) -> str:
+        tds = ""
+        for i, c in enumerate(cells):
+            weight = "600" if (bold_last and i == len(cells) - 1) else "400"
+            align = "right" if i > 0 else "left"
+            tds += (
+                f'<td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;'
+                f' color: {_TEXT}; font-weight: {weight}; text-align: {align};">{c}</td>'
+            )
+        return f"<tr style='border-bottom: 1px solid {_BORDER_SUBTLE};'>{tds}</tr>"
+
+    def _table_header(cols: list[str]) -> str:
+        ths = ""
+        for i, c in enumerate(cols):
+            align = "right" if i > 0 else "left"
+            ths += (
+                f'<th style="padding: 6px 10px; font-size: 10px; color: {_TEXT_DIM};'
+                f' text-transform: uppercase; letter-spacing: 0.05em; text-align: {align};'
+                f' font-weight: 600;">{c}</th>'
+            )
+        return f"<tr style='border-bottom: 1px solid {_BORDER};'>{ths}</tr>"
+
+    return f"""
+    <div style="max-width: 1100px; margin: 0 auto; padding: 16px;">
+
+      <h2 style="margin: 0 0 4px 0; color: {_TEXT}; font-size: 22px; font-weight: 700;
+                 letter-spacing: -0.02em;">
+        Why LFM2.5 for Transaction Sequences
+      </h2>
+      <p style="color: {_TEXT_DIM}; font-size: 13px; margin: 0 0 24px 0; line-height: 1.5;">
+        Three claims, each backed by a different kind of evidence. Serving cost is arithmetic.
+        Label efficiency is measured. Architectural fit is a first-principles argument.
+      </p>
+
+      <!-- 1. Serving cost -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          1. Serving Cost Scales Better Than Pure Attention
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 12px 0;">
+          10 of 16 layers use O(n) conv instead of O(n&sup2;) attention. The gap is structural
+          and compounds as sequence length grows. Measured at matched parameter count on the same
+          hardware:
+        </p>
+        <table style="width: 100%; border-collapse: collapse; margin-bottom: 8px;">
+          {_table_header(["Sequence", "Tokens", "Hybrid", "Pure Attn", "Speedup"])}
+          {_table_row(["64 tx &times; 15 feat", "960", "5.96 ms", "9.01 ms", "1.5x"])}
+          {_table_row(["64 tx &times; 30 feat", "1,920", "13.05 ms", "23.84 ms", "1.8x"])}
+          {_table_row(["128 tx &times; 30 feat", "3,840", "35.65 ms", "75.95 ms",
+                       "<b>2.1x</b>"], bold_last=False)}
+          {_table_row(["256 tx &times; 30 feat", "7,680", "119 ms", "283 ms",
+                       "<b>2.4x</b>"], bold_last=False)}
+        </table>
+        <p style="font-size: 11px; color: {_TEXT_DIM}; margin: 0; line-height: 1.5;">
+          At 3,840 tokens (production target), the FLOP analysis predicts 7.7% fewer operations.
+          The actual speedup is larger because conv layers achieve better hardware utilization.
+          At 7,680 tokens, savings reach 17% in FLOPs and 2.4x in wall-clock.
+        </p>
+      </div>
+
+      <!-- 2. Label scarcity -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          2. Pretraining Gets to 98% of Full-Data Quality with 10% of Labels
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 12px 0;">
+          Self-supervised pretraining on unlabeled transactions. No fraud labels needed.
+          Then fine-tune with whatever labels you have. The pretrained model at 10% labels
+          beats the random-init baseline at 100% labels.
+        </p>
+        <table style="width: 100%; border-collapse: collapse; margin-bottom: 8px;">
+          {_table_header(["Labels", "Sequences", "Pretrained PR-AUC", "Baseline PR-AUC", "Delta"])}
+          {_table_row(["1%", "1,700", "<b>0.539</b>", "0.046", "<b>+0.493</b>"])}
+          {_table_row(["10%", "17,000", "<b>0.948</b>", "0.690", "<b>+0.258</b>"])}
+          {_table_row(["100%", "170,000", "<b>0.964</b>", "0.922", "+0.041"])}
+        </table>
+        <p style="font-size: 11px; color: {_TEXT_DIM}; margin: 0; line-height: 1.5;">
+          At 1% labels the baseline has learned nothing (0.046 is random guessing).
+          The pretrained model is already useful. This is the defining advantage for
+          institutions with billions of unlabeled transactions and limited fraud labels.
+        </p>
+      </div>
+
+      <!-- 3. Convergence + multi-head -->
+      <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-bottom: 12px;">
+        <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+            5x Convergence Speed
+          </h3>
+          <p style="color: {_TEXT_MUTED}; font-size: 12px; line-height: 1.6; margin: 0 0 8px 0;">
+            The pretrained model hit 0.959 PR-AUC at step 1,000. The random-init baseline
+            never reached that level in 5,000 steps (peaked at 0.922).
+          </p>
+          <p style="color: {_TEXT_MUTED}; font-size: 12px; line-height: 1.6; margin: 0;">
+            For quarterly model refreshes, this cuts GPU-hours per cycle by 80%.
+            For weekly refreshes, it is the difference between feasible and infeasible.
+          </p>
+        </div>
+
+        <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+            One Backbone, Many Tasks
+          </h3>
+          <p style="color: {_TEXT_MUTED}; font-size: 12px; line-height: 1.6; margin: 0 0 8px 0;">
+            A single forward pass serves fraud detection, next-merchant prediction, amount
+            forecasting, and merchant category classification. Add dispute prediction, default
+            risk, or authorization optimization as new heads without retraining the backbone.
+          </p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;">
+            Tied-embedding heads: +50% merchant accuracy at zero parameter cost
+          </div>
+        </div>
+      </div>
+
+      <!-- 4. Architectural fit -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          3. The Architecture Matches Transaction Data Structure
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0 0 8px 0;">
+          Transaction data is not like text. Information density is concentrated locally
+          (within-transaction feature correlations, adjacent-transaction continuity) with sparse
+          global signal (behavioral baselines across the full history). LFM2.5 allocates
+          O(n) conv to the dense local patterns and O(n&sup2;) attention to the sparse
+          global patterns. A pure transformer allocates O(n&sup2;) compute uniformly
+          across all distances.
+        </p>
+        <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 10px;">
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 4px;">Within Transaction</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Merchant determines MCC. Entry mode correlates with amount. Dense, local,
+              often deterministic. A 3-wide conv kernel captures this.
+            </div>
+          </div>
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 4px;">Adjacent Transactions</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Strong temporal continuity. A customer at Starbucks is likely at a similar
+              merchant next. The conditional distribution of t+1 given t is heavily peaked.
+            </div>
+          </div>
+          <div style="padding: 10px; background: {_BG_CARD_ALT}; border-radius: 8px;">
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 4px;">Distant Transactions</div>
+            <div style="font-size: 11px; color: {_TEXT_DIM}; line-height: 1.4;">
+              Weak but non-zero signal. Behavioral profile matters for fraud baseline,
+              but per-position information density is thin. This is where attention earns its cost.
+            </div>
+          </div>
+        </div>
+      </div>
+
+      <!-- 5. Data ownership -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 12px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 4px 0; font-size: 15px; font-weight: 600;">
+          Your Data, Your Model, Your Infrastructure
+        </h3>
+        <p style="color: {_TEXT_MUTED}; font-size: 13px; line-height: 1.6; margin: 0;">
+          Liquid licenses the architecture and training recipe. You train on your proprietary
+          data behind your firewall. No data leaves your infrastructure. No dependency on
+          external model APIs. The result is a foundation model you own, optimized for your
+          specific transaction patterns, deployed on your hardware.
+        </p>
+      </div>
+
+      <!-- What we claim / don't claim -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 16px;">
+        <h3 style="color: {_TEXT}; margin: 0 0 8px 0; font-size: 14px; font-weight: 600;">
+          What We Claim vs What We Don't
+        </h3>
+        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; font-size: 12px;">
+          <div>
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT};
+                        font-weight: 600; margin-bottom: 6px; text-transform: uppercase;
+                        letter-spacing: 0.05em;">We claim</div>
+            <ul style="margin: 0; padding-left: 14px; color: {_TEXT_MUTED}; line-height: 1.6;">
+              <li>Fewer FLOPs per forward pass above ~1,500 tokens (arithmetic)</li>
+              <li>Cost gap widens with sequence length (structural, measured 2.1-2.4x)</li>
+              <li>Pretraining dramatically improves low-label performance (measured)</li>
+              <li>Full pipeline works end-to-end on LFM2.5 (built it)</li>
+            </ul>
+          </div>
+          <div>
+            <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                        font-weight: 600; margin-bottom: 6px; text-transform: uppercase;
+                        letter-spacing: 0.05em;">We don't claim</div>
+            <ul style="margin: 0; padding-left: 14px; color: {_TEXT_MUTED}; line-height: 1.6;">
+              <li>Quality advantage over pure transformers at matched scale</li>
+              <li>Absolute latency numbers extrapolated from 10M to 1.2B</li>
+              <li>Superiority on all tasks (the advantage is specific to structured sequences)</li>
+              <li>Real-data validation (all results are on synthetic data)</li>
+            </ul>
+          </div>
+        </div>
+      </div>
+
+      <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM}; text-align: center;">
+        Source: <a href="https://arxiv.org/abs/2511.23404" style="color: {_TEXT_DIM};
+        text-decoration: underline;">arXiv 2511.23404</a> &middot;
+        Reference implementation measured at 9.85M params, Apple Silicon
+      </div>
+    </div>
+    """
+
+
+def render_integration_guide() -> str:
+    """Render high-abstraction integration architecture flow.
+
+    Distills docs/integration-guide.md into a visual pipeline overview.
+    Monochrome design -- no rainbow pills or per-card color coding.
+    """
+
+    def _phase_card(num: str, title: str, body: str, detail: str) -> str:
+        return f"""
+        <div style="padding: 14px 16px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                    border-radius: {_RADIUS_CARD};">
+          <div style="display: flex; align-items: baseline; gap: 8px; margin-bottom: 6px;">
+            <span style="font-family: {_FONT_MONO}; font-size: 11px; color: {_TEXT_DIM};
+                         font-weight: 600;">{num}</span>
+            <span style="font-size: 14px; font-weight: 600; color: {_TEXT};">{title}</span>
+          </div>
+          <p style="font-size: 12px; color: {_TEXT_MUTED}; line-height: 1.5; margin: 0 0 8px 0;">
+            {body}</p>
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      padding: 6px 8px; background: {_BG_CARD_ALT}; border-radius: 6px;
+                      line-height: 1.5;">
+            {detail}</div>
+        </div>"""
+
+    def _gotcha(num: str, title: str, desc: str) -> str:
+        return f"""
+        <div style="display: flex; gap: 8px; padding: 5px 0;
+                    border-bottom: 1px solid {_BORDER_SUBTLE};">
+          <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM};
+                      font-weight: 600; min-width: 18px;">{num}.</div>
+          <div>
+            <span style="font-size: 12px; font-weight: 600; color: {_TEXT};">{title}</span>
+            <span style="font-size: 12px; color: {_TEXT_MUTED};"> -- {desc}</span>
+          </div>
+        </div>"""
+
+    return f"""
+    <div style="max-width: 1100px; margin: 0 auto; padding: 16px;">
+
+      <h2 style="margin: 0 0 4px 0; color: {_TEXT}; font-size: 22px; font-weight: 700;
+                 letter-spacing: -0.02em;">
+        Integration Architecture
+      </h2>
+      <p style="color: {_TEXT_DIM}; font-size: 13px; margin: 0 0 20px 0; line-height: 1.5;">
+        Seven phases from raw transaction data to production fraud scoring.
+        Each phase has a clear input, output, and set of decisions.
+      </p>
+
+      <!-- Pipeline flow -->
+      <div style="display: flex; align-items: center; justify-content: center; gap: 6px;
+                  margin-bottom: 24px; padding: 10px 0; flex-wrap: wrap;">
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Schema</span>
+        <span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Tokenize</span>
+        <span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Embed</span>
+        <span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Backbone</span>
+        <span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Pretrain</span>
+        <span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Heads</span>
+        <span style="color: {_TEXT_DIM}; font-size: 12px;">&rarr;</span>
+        <span style="padding: 5px 12px; background: {_TEXT}; color: #fff;
+                     border-radius: 9999px; font-family: {_FONT_MONO};
+                     font-size: 10px; font-weight: 600;">Deploy</span>
+      </div>
+
+      <!-- Phase cards -->
+      <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-bottom: 20px;">
+        {_phase_card("1", "Schema Design",
+            "Define your features, vocab sizes, and feature ordering. This is the contract "
+            "between your data team and your ML team. Three feature types: categorical "
+            "(direct index), ordinal (hour, day-of-week), and bucketed-continuous (amount "
+            "into quantile bins). Spend two weeks here.",
+            "3 reserved tokens per feature (MASK, OOV, NULL) &nbsp;|&nbsp; "
+            "Order features by semantic family for the conv window &nbsp;|&nbsp; "
+            "Embed schema fingerprint in checkpoint metadata")}
+
+        {_phase_card("2", "Tokenization",
+            "Convert raw transaction fields into integer token IDs. Categorical features "
+            "map directly to vocab indices. Continuous values (amount, days-since-last) get "
+            "quantile-bucketed into N bins. High-cardinality features (merchant_id) need the "
+            "long tail bucketed or factored into orthogonal features.",
+            "amount &rarr; 16-256 quantile bins &nbsp;|&nbsp; "
+            "merchant_id: top 10K distinct, rest into ~1K frequency buckets &nbsp;|&nbsp; "
+            "Unseen values at inference &rarr; OOV token (ID 1)")}
+
+        {_phase_card("3", "Structured Embedding",
+            "One embedding table per feature (sized to its vocab) plus a feature-type table. "
+            "Summed to produce the (B, T*F, D) tensor the backbone expects. "
+            "No raw continuous features -- everything goes through an embedding table.",
+            "value_tables[f](token) + type_table(f) &nbsp;|&nbsp; "
+            "High-cardinality features dominate param budget -- "
+            "bucket the merchant long tail")}
+
+        {_phase_card("4", "Backbone Config",
+            "Start from a published LFM2 scale point (350M, 700M, 1.2B, 2.6B). "
+            "Keep the conv-to-attention ratio, GQA config, SwiGLU MLP, and RoPE theta. "
+            "Do not deviate without a specific reason.",
+            "10:6 conv:attn at 1.2B &nbsp;|&nbsp; 32Q/8KV GQA &nbsp;|&nbsp; "
+            "QK-RMSNorm before RoPE &nbsp;|&nbsp; theta=1M &nbsp;|&nbsp; "
+            "First and last layers are conv")}
+
+        {_phase_card("5", "Pretraining",
+            "Self-supervised on your unlabeled transactions. No fraud labels needed. "
+            "Causal next-feature prediction or masked-transaction prediction. "
+            "Average per-feature losses (do NOT sum).",
+            "AdamW lr=3e-4, betas=(0.9, 0.95), wd=0.1 &nbsp;|&nbsp; "
+            "Cosine decay to 10% &nbsp;|&nbsp; BF16 (loss in FP32) &nbsp;|&nbsp; "
+            "Chinchilla: ~20 tokens per parameter")}
+
+        {_phase_card("6", "Downstream Heads",
+            "Attach task-specific heads. The critical decision: "
+            "TiedEmbeddingHead for features that were in pretraining vocab "
+            "(next_merchant, MCC). Fresh MLP for everything else (fraud, disputes).",
+            "Tied head: adapter(h) @ embedding.weight.T (+50% accuracy) &nbsp;|&nbsp; "
+            "Pool: last_tx_mean for sequence tasks, pre_last_tx for next-tx &nbsp;|&nbsp; "
+            "Dual LR: backbone 5e-5, heads 1e-3")}
+      </div>
+
+      <!-- Deployment gets its own full-width card -->
+      <div style="margin-bottom: 20px;">
+        {_phase_card("7", "Deployment",
+            "KV cache only on 6 attention layers (not 16). Dynamic batching with "
+            "5-10ms collection window. Conv-dominant models quantize cleanly to INT8. "
+            "Same model runs on GPU or CPU.",
+            "LFM2 @ 3,840 tokens: ~25 GB KV at 1K concurrent &nbsp;|&nbsp; "
+            "Pure attn: ~64 GB &nbsp;|&nbsp; "
+            "Sub-100ms on H100 for real-time auth decisioning")}
+      </div>
+
+      <!-- Gotchas -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 20px;">
+        <div style="font-size: 14px; font-weight: 600; color: {_TEXT}; margin-bottom: 8px;">
+          Top Gotchas
+        </div>
+        {_gotcha("1", "Per-feature loss summing",
+                 "high-cardinality features dominate. Average, don't sum.")}
+        {_gotcha("2", "Fresh MLP for pretrained features",
+                 "next_merchant stays at random. Use TiedEmbeddingHead.")}
+        {_gotcha("3", "Schema-checkpoint drift",
+                 "training on schema v3, deploying on v2. Embed fingerprint in metadata.")}
+        {_gotcha("4", "Frozen backbone",
+                 "PR-AUC drops from 0.96 to 0.16. Plan to fine-tune.")}
+        {_gotcha("5", "Pool strategy leak",
+                 "last_tx_mean for next-tx prediction leaks the target. Use pre_last_tx.")}
+      </div>
+
+      <!-- Engagement model -->
+      <div style="padding: 16px 20px; background: {_BG_CARD}; border: 1px solid {_BORDER};
+                  border-radius: {_RADIUS_CARD}; margin-bottom: 16px;">
+        <div style="font-size: 14px; font-weight: 600; color: {_TEXT}; margin-bottom: 10px;">
+          Typical Engagement
+        </div>
+        <table style="width: 100%; border-collapse: collapse;">
+          <tr style="border-bottom: 1px solid {_BORDER};">
+            <th style="padding: 6px 10px; text-align: left; font-size: 10px; color: {_TEXT_DIM};
+                       text-transform: uppercase; letter-spacing: 0.05em; font-weight: 600;">Phase</th>
+            <th style="padding: 6px 10px; text-align: left; font-size: 10px; color: {_TEXT_DIM};
+                       text-transform: uppercase; font-weight: 600;">Duration</th>
+            <th style="padding: 6px 10px; text-align: left; font-size: 10px; color: {_TEXT_DIM};
+                       text-transform: uppercase; font-weight: 600;">What Happens</th>
+          </tr>
+          <tr style="border-bottom: 1px solid {_BORDER_SUBTLE};">
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">Discovery</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">2-4 weeks</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Schema definition, data sample (~1M sequences), compliance review, architectural fit assessment</td>
+          </tr>
+          <tr style="border-bottom: 1px solid {_BORDER_SUBTLE};">
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">POC</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">2 weeks</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Pretrain + fine-tune on your data sample, measurement report, go/no-go recommendation</td>
+          </tr>
+          <tr style="border-bottom: 1px solid {_BORDER_SUBTLE};">
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">Production</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">3-6 months</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Engineering team builds, Liquid provides architectural support, weekly design review</td>
+          </tr>
+          <tr>
+            <td style="padding: 5px 10px; font-size: 12px; font-weight: 600;
+                       color: {_TEXT};">Scale</td>
+            <td style="padding: 5px 10px; font-family: {_FONT_MONO}; font-size: 11px;
+                       color: {_TEXT_MUTED};">Ongoing</td>
+            <td style="padding: 5px 10px; font-size: 12px; color: {_TEXT_MUTED};">
+              Operations, monitoring, retraining cadence, architecture evolution</td>
+          </tr>
+        </table>
+      </div>
+
+      <div style="font-family: {_FONT_MONO}; font-size: 10px; color: {_TEXT_DIM}; text-align: center;">
+        Full guide: docs/integration-guide.md &middot;
+        Architecture: <a href="https://arxiv.org/abs/2511.23404" style="color: {_TEXT_DIM};
+        text-decoration: underline;">arXiv 2511.23404</a>
+      </div>
+    </div>
+    """

src/model/embedding.py

@@ -0,0 +1,67 @@
+"""Structured-feature embedding layer for transaction sequences.
+
+Replaces the single text-token embedding table of public LFM2.5 with
+per-feature value tables plus a feature-type table, summed. Each feature
+has its own embedding table sized to that feature's vocabulary. A shared
+feature-type table (15 rows) tells the model which feature a token represents.
+
+Input:  (B, T, F) int tensor of token IDs, where T=64 transactions, F=15 features
+Output: (B, T*F, D) float tensor of embeddings, where D=hidden_dim
+"""
+
+import torch
+import torch.nn as nn
+
+from src.data.schema import SchemaConfig
+
+
+class StructuredEmbedding(nn.Module):
+    """Per-feature value embeddings + feature-type embeddings, summed.
+
+    The value_tables are exposed as a ModuleList so the per-feature LM heads
+    can tie weights to them. The type_table is NOT tied to anything.
+    """
+
+    def __init__(self, schema: SchemaConfig, hidden_dim: int) -> None:
+        super().__init__()
+        self.num_features = schema.num_features
+        self.num_transactions = schema.num_transactions
+        self.hidden_dim = hidden_dim
+
+        self.value_tables = nn.ModuleList([
+            nn.Embedding(feature.vocab_size, hidden_dim)
+            for feature in schema.features
+        ])
+
+        self.type_table = nn.Embedding(schema.num_features, hidden_dim)
+
+        self._vocab_sizes = [f.vocab_size for f in schema.features]
+
+    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+        """Embed structured token IDs into a flat sequence.
+
+        Args:
+            token_ids: (B, T, F) int tensor. T=num_transactions, F=num_features.
+                       Each value must be in [0, vocab_size) for its feature.
+
+        Returns:
+            (B, T*F, D) float tensor of summed value + type embeddings.
+        """
+        B, T, F = token_ids.shape
+        assert F == self.num_features, (
+            f"Expected {self.num_features} features, got {F}"
+        )
+
+        # Feature type indices: [0, 1, 2, ..., F-1], broadcast across batch and time
+        type_indices = torch.arange(F, device=token_ids.device)  # (F,)
+        type_emb = self.type_table(type_indices)  # (F, D)
+
+        feature_embeddings = []
+        for f_idx in range(F):
+            feat_tokens = token_ids[:, :, f_idx]  # (B, T)
+            val_emb = self.value_tables[f_idx](feat_tokens)  # (B, T, D)
+            feature_embeddings.append(val_emb + type_emb[f_idx])  # (B, T, D) + (D,) broadcast
+
+        # (B, T, F, D) -> (B, T*F, D)
+        stacked = torch.stack(feature_embeddings, dim=2)  # (B, T, F, D)
+        return stacked.reshape(B, T * F, self.hidden_dim)  # (B, 960, D)

src/model/heads.py

@@ -0,0 +1,65 @@
+"""Prediction heads for the transaction foundation model.
+
+PerFeatureLMHeads: 15 linear projections for next-token prediction during
+pretraining. Each head projects hidden_dim -> feature_vocab_size with weights
+tied to the corresponding value embedding table (parameter identity, not copy).
+No bias, matching LFM2's lm_head convention.
+
+Weight tying semantics:
+  - Forward through embedding: output = weight[token_id]  (index select)
+  - Forward through LM head: logits = hidden @ weight.T   (matmul)
+  - Backward: gradients flow through BOTH paths to the SAME Parameter.
+    The optimizer updates it once, accumulating both contributions.
+  - The feature-type embedding table is NOT tied to any head (it has no
+    corresponding prediction target).
+
+The backbone's final RMSNorm (embedding_norm) is applied to hidden states
+BEFORE they reach these heads, matching LFM2's architecture. That
+normalization lives in the backbone (F8), not here.
+"""
+
+import torch
+import torch.nn as nn
+
+from src.model.embedding import StructuredEmbedding
+
+
+class PerFeatureLMHeads(nn.Module):
+    """Per-feature prediction heads with tied embedding weights.
+
+    Each of the 15 heads is a bias-free linear projection from hidden_dim to
+    that feature's vocab_size. The weight matrix is the SAME nn.Parameter as
+    the corresponding value embedding table.
+    """
+
+    def __init__(self, embedding: StructuredEmbedding) -> None:
+        super().__init__()
+        self.num_features = embedding.num_features
+        self.hidden_dim = embedding.hidden_dim
+
+        self.heads = nn.ModuleList()
+        for f_idx in range(self.num_features):
+            vocab_size = embedding.value_tables[f_idx].num_embeddings
+            head = nn.Linear(embedding.hidden_dim, vocab_size, bias=False)
+            # Tie weights: head.weight IS embedding.value_tables[f].weight.
+            # Both the embedding lookup and the LM head projection operate on
+            # the same underlying tensor. PyTorch deduplicates in parameters().
+            head.weight = embedding.value_tables[f_idx].weight
+            self.heads.append(head)
+
+    def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:
+        """Project hidden states to per-feature logits.
+
+        Args:
+            hidden_states: (B, S, D) from the backbone (after embedding_norm).
+
+        Returns:
+            List of 15 tensors. heads[f] has shape (B, S, vocab_size_f).
+            During training, the loss function selects positions per feature:
+            position p predicts feature (p+1) % num_features.
+        """
+        return [head(hidden_states) for head in self.heads]
+
+    def get_vocab_sizes(self) -> list[int]:
+        """Vocab size for each feature, useful for loss computation."""
+        return [head.out_features for head in self.heads]

src/model/lfm2_small.py

@@ -0,0 +1,387 @@
+"""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))

src/model/task_heads.py

@@ -0,0 +1,250 @@
+"""Downstream task heads for multi-head fine-tuning.
+
+Four heads sharing the pretrained LFM2Small backbone:
+
+  1. fraud          -- P(fraud) from last-transaction pool (BCE)
+  2. next_merchant  -- merchant_id of last tx from prefix (CE, self-supervised)
+  3. amount         -- amount bucket of last tx from prefix (CE, self-supervised)
+  4. mcc            -- MCC of last tx from prefix (CE, self-supervised)
+
+Two head implementations:
+  - DownstreamHead        -- fresh 2-layer MLP. Doesn't inherit pretrained
+                             knowledge stored in the tied LM-head weights.
+                             Default for fraud (which has no LM-head analog).
+  - TiedEmbeddingHead     -- pool -> adapter -> matmul through the backbone's
+                             tied embedding table for the target feature.
+                             Recovers the pretrained next-feature signal that
+                             a fresh MLP cannot. Use when the target feature
+                             has a corresponding weight-tied LM head in
+                             pretraining (anything with target_type "feature:N").
+
+Pool strategies:
+  last_tx_mean -- mean of last transaction's 15 positions. These have seen
+                  the full sequence (causal masking), so the representation
+                  encodes full-sequence context. Used for sequence-level tasks.
+  pre_last_tx  -- hidden state at position S - num_features - 1 (end of tx 62).
+                  Sees tx 0-62 only, appropriate for predicting tx 63's features
+                  without target leakage. Verified safe: both causal attention
+                  and left-padded Conv1d respect this boundary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class HeadConfig:
+    """Single head configuration, loaded from finetune.yaml."""
+
+    name: str
+    output_dim: int
+    loss_type: str          # "bce" | "ce"
+    pool_strategy: str      # "last_tx_mean" | "pre_last_tx"
+    target_type: str        # "sequence_label" | "feature:<idx>"
+    weight: float = 1.0
+    mlp_hidden: int = 128
+    dropout: float = 0.1
+
+
+class DownstreamHead(nn.Module):
+    """Pool hidden states -> 2-layer MLP -> task prediction.
+
+    Param count per head: ~33K (small output) to ~65K (merchant_id, 5003 classes).
+    """
+
+    def __init__(self, config: HeadConfig, hidden_dim: int, num_features: int) -> None:
+        super().__init__()
+        self.config = config
+        self.num_features = num_features
+
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_dim, config.mlp_hidden),
+            nn.ReLU(),
+            nn.Dropout(config.dropout),
+            nn.Linear(config.mlp_hidden, config.output_dim),
+        )
+
+    def pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """(B, S, D) -> (B, D) via head-specific pooling."""
+        nf = self.num_features
+        if self.config.pool_strategy == "last_tx_mean":
+            return hidden_states[:, -nf:, :].mean(dim=1)
+        if self.config.pool_strategy == "pre_last_tx":
+            return hidden_states[:, -(nf + 1), :]
+        raise ValueError(f"Unknown pool strategy: {self.config.pool_strategy}")
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.pool(hidden_states))
+
+    def compute_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        if self.config.loss_type == "bce":
+            return F.binary_cross_entropy_with_logits(
+                logits.squeeze(-1), targets.float(),
+            )
+        return F.cross_entropy(logits, targets)
+
+    def extract_targets(
+        self,
+        token_ids: torch.Tensor,
+        sequence_labels: torch.Tensor | None,
+        aux_targets: dict[str, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Get this head's targets from input data.
+
+        Args:
+            token_ids: (B, T, F).
+            sequence_labels: (B,) binary fraud labels, or None.
+            aux_targets: optional dict of auxiliary per-sequence targets
+                         (e.g. "amount_range" from amount_range_labels.npy).
+        """
+        if self.config.target_type == "sequence_label":
+            assert sequence_labels is not None
+            return sequence_labels
+        if self.config.target_type == "amount_range":
+            assert aux_targets is not None and "amount_range" in aux_targets
+            return aux_targets["amount_range"]
+        feat_idx = int(self.config.target_type.split(":")[1])
+        return token_ids[:, -1, feat_idx]
+
+
+class TiedEmbeddingHead(nn.Module):
+    """Downstream head that projects through the backbone's tied embedding table.
+
+    The pretrained LM head for feature F is the value-embedding table for F
+    (weight-tied). The fresh-MLP DownstreamHead discards that projection by
+    learning a new one from scratch. This head preserves it: pool the backbone
+    hidden states, run a small adapter, then matmul through the same embedding
+    table the backbone reads from at input time.
+
+    Why this matters:
+      - At fine-tune time the adapter learns a small distribution shift, not
+        a 256-d -> vocab_size projection from scratch.
+      - Gradients flow back to the embedding table from both the input-embed
+        side and this head, just as during pretraining.
+      - When the backbone is frozen, the embedding table is frozen too, so
+        this head reduces to "small adapter on top of pretrained features."
+
+    Constraints:
+      - Only valid for target_type "feature:N". The tied table is keyed to
+        a specific feature index.
+      - Output dim is implicitly the vocab_size of that feature.
+    """
+
+    def __init__(
+        self,
+        config: HeadConfig,
+        hidden_dim: int,
+        num_features: int,
+        value_tables: nn.ModuleList,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.num_features = num_features
+
+        if not config.target_type.startswith("feature:"):
+            raise ValueError(
+                f"TiedEmbeddingHead requires target_type 'feature:N', got "
+                f"{config.target_type!r}",
+            )
+        self.feature_idx = int(config.target_type.split(":")[1])
+        # value_tables is the same nn.ModuleList in StructuredEmbedding. By
+        # assigning it as a module attribute we share parameters with the
+        # backbone — no parameter duplication in state_dict.
+        self.value_tables = value_tables
+
+        # The adapter is the only thing this head learns from scratch. Two
+        # layers with a nonlinearity gives enough flexibility for a small
+        # distribution-shift correction without obscuring the tied projection.
+        self.adapter = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(config.dropout),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+
+    def pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        nf = self.num_features
+        if self.config.pool_strategy == "last_tx_mean":
+            return hidden_states[:, -nf:, :].mean(dim=1)
+        if self.config.pool_strategy == "pre_last_tx":
+            return hidden_states[:, -(nf + 1), :]
+        raise ValueError(f"Unknown pool strategy: {self.config.pool_strategy}")
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        pooled = self.pool(hidden_states)              # (B, D)
+        adapted = self.adapter(pooled)                  # (B, D)
+        # Project through the tied embedding table. F.linear computes
+        # adapted @ weight.T, exactly matching how the pretrained LM head
+        # produces logits. No bias, matching PerFeatureLMHeads.
+        weight = self.value_tables[self.feature_idx].weight  # (vocab_f, D)
+        return F.linear(adapted, weight)
+
+    def compute_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        # Tied heads only support CE — they predict over a tied vocabulary.
+        return F.cross_entropy(logits, targets)
+
+    def extract_targets(
+        self,
+        token_ids: torch.Tensor,
+        sequence_labels: torch.Tensor | None,
+        aux_targets: dict[str, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        return token_ids[:, -1, self.feature_idx]
+
+
+# DownstreamHead and TiedEmbeddingHead share a duck-typed interface.
+# Union type for the head dict in MultiHeadModel.
+AnyHead = DownstreamHead | TiedEmbeddingHead
+
+
+class MultiHeadModel(nn.Module):
+    """Pretrained backbone + downstream task heads.
+
+    Calls backbone.backbone_forward() once, then each head pools and
+    predicts independently. Backbone may be frozen or slow-learned.
+    """
+
+    def __init__(
+        self, backbone: nn.Module, heads: dict[str, AnyHead],
+    ) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.heads = nn.ModuleDict(heads)
+
+    def forward(self, token_ids: torch.Tensor) -> dict[str, torch.Tensor]:
+        hidden = self.backbone.backbone_forward(token_ids)
+        return {name: head(hidden) for name, head in self.heads.items()}
+
+    def compute_losses(
+        self,
+        predictions: dict[str, torch.Tensor],
+        token_ids: torch.Tensor,
+        sequence_labels: torch.Tensor | None,
+        aux_targets: dict[str, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, dict[str, float]]:
+        """Weighted sum of per-head losses."""
+        device = next(iter(predictions.values())).device
+        total = torch.tensor(0.0, device=device)
+        per_head: dict[str, float] = {}
+
+        for name, head in self.heads.items():
+            targets = head.extract_targets(
+                token_ids, sequence_labels, aux_targets,
+            ).to(device)
+            loss = head.compute_loss(predictions[name], targets)
+            total = total + head.config.weight * loss
+            per_head[name] = loss.item()
+
+        return total, per_head
+
+    def head_param_count(self) -> dict[str, int]:
+        """Parameter count per head (excludes shared backbone)."""
+        return {
+            name: sum(p.numel() for p in head.parameters())
+            for name, head in self.heads.items()
+        }

src/training/finetune.py

@@ -0,0 +1,504 @@
+"""Multi-head fine-tuning loop.
+
+Loads a pretrained LFM2Small backbone and attaches downstream task heads:
+  - fraud: binary classification (supervised, BCE)
+  - next_merchant, amount, mcc: feature prediction (self-supervised, CE)
+
+Supports:
+  - Joint multi-head training with per-head loss weighting
+  - Separate LR for backbone (5e-5) vs heads (1e-3)
+  - Optional backbone freezing
+  - Random-init baseline (omit --checkpoint)
+  - Label-scarcity subsampling (--label-fraction)
+  - ROC-AUC / PR-AUC for classification, top-k accuracy for categorical
+"""
+
+from __future__ import annotations
+
+import logging
+import shutil
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import yaml
+from sklearn.metrics import average_precision_score, roc_auc_score
+from torch.amp import GradScaler, autocast
+from torch.utils.data import DataLoader, Dataset
+
+from src.data.schema import load_schema
+from src.model.lfm2_small import LFM2Small, ModelConfig
+from src.model.task_heads import (
+    AnyHead,
+    DownstreamHead,
+    HeadConfig,
+    MultiHeadModel,
+    TiedEmbeddingHead,
+)
+from src.training.trainer_utils import (
+    MetricsLogger,
+    create_scheduler,
+    load_checkpoint,
+    nan_guard,
+    save_checkpoint,
+    setup_deterministic,
+)
+
+log = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+class FinetuneDataset(Dataset):
+    """Token IDs + sequence-level fraud labels + optional amount_range targets."""
+
+    def __init__(
+        self,
+        token_ids: np.ndarray,
+        sequence_labels: np.ndarray,
+        indices: np.ndarray,
+        amount_range_targets: np.ndarray | None = None,
+    ) -> None:
+        self.token_ids = token_ids[indices]
+        self.labels = sequence_labels[indices]
+        if amount_range_targets is not None:
+            self.ar_targets = amount_range_targets[indices]
+        else:
+            self.ar_targets = np.full(len(indices), -1, dtype=np.int8)
+
+    def __len__(self) -> int:
+        return len(self.token_ids)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return (
+            torch.from_numpy(self.token_ids[idx].astype(np.int64)),
+            torch.tensor(self.labels[idx], dtype=torch.float32),
+            torch.tensor(int(self.ar_targets[idx]), dtype=torch.int64),
+        )
+
+
+# ---------------------------------------------------------------------------
+# Head construction
+# ---------------------------------------------------------------------------
+
+
+def build_heads(
+    config: dict[str, Any],
+    hidden_dim: int,
+    num_features: int,
+    backbone: LFM2Small | None = None,
+) -> dict[str, AnyHead]:
+    """Instantiate downstream heads from finetune.yaml.
+
+    If a head config sets ``tied: true``, instantiates TiedEmbeddingHead instead
+    of DownstreamHead. Tied heads need the backbone's value_tables to share
+    parameters, so ``backbone`` must be provided whenever any tied head is
+    configured.
+    """
+    heads: dict[str, AnyHead] = {}
+    for name, hcfg in config["heads"].items():
+        hc = HeadConfig(
+            name=name,
+            output_dim=hcfg["output_dim"],
+            loss_type=hcfg["loss"],
+            pool_strategy=hcfg["pool"],
+            target_type=hcfg["target"],
+            weight=hcfg.get("weight", 1.0),
+            mlp_hidden=hcfg.get("mlp_hidden", 128),
+            dropout=hcfg.get("dropout", 0.1),
+        )
+        if hcfg.get("tied", False):
+            if backbone is None:
+                raise ValueError(
+                    f"Head '{name}' uses tied=true but build_heads was called "
+                    f"without a backbone reference.",
+                )
+            heads[name] = TiedEmbeddingHead(
+                hc, hidden_dim, num_features, backbone.embedding.value_tables,
+            )
+        else:
+            heads[name] = DownstreamHead(hc, hidden_dim, num_features)
+    return heads
+
+
+# ---------------------------------------------------------------------------
+# Optimizer
+# ---------------------------------------------------------------------------
+
+
+def create_finetune_optimizer(
+    model: MultiHeadModel,
+    backbone_lr: float,
+    head_lr: float,
+    betas: tuple[float, float] = (0.9, 0.95),
+    weight_decay: float = 0.1,
+) -> torch.optim.AdamW:
+    """Separate LR groups: backbone (slow) vs heads (fast)."""
+    backbone_decay: list[torch.nn.Parameter] = []
+    backbone_no_decay: list[torch.nn.Parameter] = []
+    head_decay: list[torch.nn.Parameter] = []
+    head_no_decay: list[torch.nn.Parameter] = []
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        is_head = name.startswith("heads.")
+        if is_head:
+            (head_decay if param.dim() >= 2 else head_no_decay).append(param)
+        else:
+            (backbone_decay if param.dim() >= 2 else backbone_no_decay).append(param)
+
+    groups = [
+        {"params": backbone_decay, "lr": backbone_lr, "weight_decay": weight_decay},
+        {"params": backbone_no_decay, "lr": backbone_lr, "weight_decay": 0.0},
+        {"params": head_decay, "lr": head_lr, "weight_decay": weight_decay},
+        {"params": head_no_decay, "lr": head_lr, "weight_decay": 0.0},
+    ]
+    groups = [g for g in groups if g["params"]]
+    return torch.optim.AdamW(groups, betas=betas)
+
+
+# ---------------------------------------------------------------------------
+# Validation
+# ---------------------------------------------------------------------------
+
+
+def _topk_accuracy(logits: torch.Tensor, targets: torch.Tensor, k: int) -> float:
+    topk = logits.topk(k, dim=-1).indices
+    return (topk == targets.unsqueeze(-1)).any(-1).float().mean().item()
+
+
+@torch.no_grad()
+def validate(
+    model: MultiHeadModel,
+    val_loader: DataLoader,
+    device: torch.device,
+    amp_dtype: torch.dtype,
+    use_amp: bool,
+) -> dict[str, float]:
+    """Per-head validation: ROC-AUC/PR-AUC for BCE, top-k accuracy for CE."""
+    model.eval()
+    all_preds: dict[str, list[torch.Tensor]] = {n: [] for n in model.heads}
+    all_targets: dict[str, list[torch.Tensor]] = {n: [] for n in model.heads}
+    total_loss = 0.0
+    n_batches = 0
+
+    for token_ids, labels, ar in val_loader:
+        token_ids = token_ids.to(device)
+        labels = labels.to(device)
+        ar = ar.to(device)
+        aux = {"amount_range": ar}
+
+        with autocast(device.type, dtype=amp_dtype, enabled=use_amp):
+            preds = model(token_ids)
+
+        loss, _ = model.compute_losses(preds, token_ids, labels, aux)
+        total_loss += loss.item()
+        n_batches += 1
+
+        for name, head in model.heads.items():
+            targets = head.extract_targets(token_ids, labels, aux)
+            if head.config.loss_type == "bce":
+                all_preds[name].append(
+                    torch.sigmoid(preds[name].squeeze(-1)).cpu(),
+                )
+            else:
+                all_preds[name].append(preds[name].cpu())
+            all_targets[name].append(targets.cpu())
+
+    metrics: dict[str, float] = {
+        "val/loss": total_loss / max(1, n_batches),
+    }
+
+    for name, head in model.heads.items():
+        p = torch.cat(all_preds[name]).float()
+        t = torch.cat(all_targets[name]).float()
+
+        if head.config.loss_type == "bce":
+            p_np, t_np = p.numpy(), t.numpy().astype(int)
+            if len(np.unique(t_np)) > 1:
+                metrics[f"val/{name}/roc_auc"] = roc_auc_score(t_np, p_np)
+                metrics[f"val/{name}/pr_auc"] = average_precision_score(t_np, p_np)
+        else:
+            metrics[f"val/{name}/top1_acc"] = (
+                (p.argmax(-1) == t).float().mean().item()
+            )
+            if head.config.output_dim > 10:
+                metrics[f"val/{name}/top5_acc"] = _topk_accuracy(p, t, k=5)
+
+    model.train()
+    return metrics
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+
+def _cycle(loader: DataLoader):
+    while True:
+        for batch in loader:
+            yield batch
+
+
+def finetune(
+    pretrain_checkpoint: str | Path | None = None,
+    finetune_config: str | Path = "configs/finetune.yaml",
+    model_config: str | Path = "configs/model.yaml",
+    schema_path: str | Path = "data/schema.yaml",
+    data_dir: str | Path = "data/synthetic",
+    output_dir: str | Path = "experiments",
+    max_steps: int | None = None,
+    label_fraction: float | None = None,
+) -> Path:
+    """Multi-head fine-tuning. Returns path to final checkpoint.
+
+    pretrain_checkpoint=None runs random-init baseline.
+    label_fraction overrides config's active_label_fraction.
+    """
+    with open(finetune_config) as f:
+        config = yaml.safe_load(f)
+    model_cfg = ModelConfig.from_yaml(model_config)
+    schema = load_schema(schema_path)
+    data_dir = Path(data_dir)
+
+    total_steps = (
+        max_steps if max_steps is not None
+        else config["training"]["total_steps"]
+    )
+
+    mode = "finetune" if pretrain_checkpoint else "baseline"
+    run_name = f"{mode}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    run_dir = Path(output_dir) / run_name
+    ckpt_dir = run_dir / "checkpoints"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(finetune_config, run_dir / "finetune_config.yaml")
+
+    repro = config["reproducibility"]
+    setup_deterministic(
+        repro["seed"], repro["deterministic_warn_only"],
+        repro["cublas_workspace_config"],
+    )
+
+    # ---- Data ----
+    token_ids = np.load(data_dir / "token_ids.npy")
+    sequence_labels = np.load(data_dir / "sequence_labels.npy")
+    splits = np.load(data_dir / "split_indices.npz")
+
+    ar_path = data_dir / "amount_range_labels.npy"
+    ar_targets: np.ndarray | None = None
+    if ar_path.exists():
+        ar_all = np.load(ar_path)
+        ar_targets = ar_all[:, -1]
+        log.info("Loaded amount_range targets: %d sequences, 16 classes", len(ar_targets))
+
+    train_indices = splits["train"]
+    frac = (
+        label_fraction if label_fraction is not None
+        else config.get("active_label_fraction", 1.0)
+    )
+    if frac < 1.0:
+        rng = np.random.RandomState(repro["seed"])
+        n_keep = max(1, int(len(train_indices) * frac))
+        train_indices = rng.choice(train_indices, n_keep, replace=False)
+        log.info(
+            "Label scarcity: %d/%d training seqs (%.0f%%)",
+            n_keep, len(splits["train"]), frac * 100,
+        )
+
+    train_ds = FinetuneDataset(token_ids, sequence_labels, train_indices, ar_targets)
+    val_ds = FinetuneDataset(token_ids, sequence_labels, splits["val"], ar_targets)
+
+    bs = config["training"]["batch_size"]
+    train_loader = DataLoader(
+        train_ds, batch_size=bs, shuffle=True,
+        num_workers=4, prefetch_factor=2, pin_memory=True, drop_last=True,
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=bs, shuffle=False,
+        num_workers=2, pin_memory=True,
+    )
+
+    # ---- Model ----
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    backbone = LFM2Small(model_cfg, schema)
+
+    if pretrain_checkpoint is not None:
+        ckpt = load_checkpoint(pretrain_checkpoint, backbone)
+        log.info(
+            "Loaded pretrained backbone from %s (step %s)",
+            pretrain_checkpoint, ckpt.get("step", "?"),
+        )
+    else:
+        log.info("Random-init baseline: no pretrained weights")
+
+    heads = build_heads(
+        config, model_cfg.hidden_size, model_cfg.num_features, backbone,
+    )
+    model = MultiHeadModel(backbone, heads).to(device)
+
+    if config["training"].get("freeze_backbone", False):
+        for param in model.backbone.parameters():
+            param.requires_grad = False
+        log.info("Backbone frozen: training heads only")
+
+    hpc = model.head_param_count()
+    log.info(
+        "Heads: %s | total head params: %d",
+        {k: f"{v:,}" for k, v in hpc.items()}, sum(hpc.values()),
+    )
+
+    # ---- Optimizer + scheduler ----
+    opt_cfg = config["optimizer"]
+    head_lr = opt_cfg.get("head_lr", opt_cfg["lr"] * 20)
+    optimizer = create_finetune_optimizer(
+        model, backbone_lr=opt_cfg["lr"], head_lr=head_lr,
+        betas=tuple(opt_cfg["betas"]), weight_decay=opt_cfg["weight_decay"],
+    )
+    sched_cfg = config["scheduler"]
+    scheduler = create_scheduler(
+        optimizer, sched_cfg["warmup_steps"], total_steps,
+        sched_cfg["min_lr_fraction"],
+    )
+
+    # ---- Mixed precision ----
+    use_amp = device.type == "cuda"
+    amp_dtype = (
+        torch.bfloat16
+        if use_amp and torch.cuda.is_bf16_supported()
+        else torch.float32
+    )
+    scaler = GradScaler(enabled=(use_amp and amp_dtype == torch.float16))
+
+    # ---- Logging ----
+    log_cfg = config.get("logging", {})
+    metrics_logger = MetricsLogger(
+        run_dir / "logs", run_name, log_cfg.get("wandb", "auto"),
+    )
+    metrics_logger.log_config(config)
+
+    # ---- Training loop ----
+    accum_steps = config["training"]["effective_batch_size"] // bs
+    grad_clip = config["training"]["grad_clip"]
+    ckpt_interval = config["checkpointing"]["interval_steps"]
+    val_interval = config["validation"]["interval_steps"]
+    log_interval = log_cfg.get("log_interval", 50)
+
+    data_iter = _cycle(train_loader)
+    model.train()
+    t0 = time.monotonic()
+
+    log.info("Fine-tuning [%s]: steps 0 -> %d", mode, total_steps)
+
+    for step in range(total_steps):
+        acc_loss = 0.0
+        acc_ph: dict[str, float] = {}
+
+        for _ in range(accum_steps):
+            tids, lbls, ar = next(data_iter)
+            tids, lbls, ar = tids.to(device), lbls.to(device), ar.to(device)
+            aux = {"amount_range": ar}
+
+            with autocast(device.type, dtype=amp_dtype, enabled=use_amp):
+                preds = model(tids)
+                loss, ph = model.compute_losses(preds, tids, lbls, aux)
+                scaled = loss / accum_steps
+
+            scaler.scale(scaled).backward()
+            acc_loss += loss.item()
+            for k, v in ph.items():
+                acc_ph[k] = acc_ph.get(k, 0.0) + v
+
+        acc_loss /= accum_steps
+        for k in acc_ph:
+            acc_ph[k] /= accum_steps
+
+        scaler.unscale_(optimizer)
+        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        scheduler.step()
+        optimizer.zero_grad()
+
+        if config["nan_guard"]["enabled"]:
+            nan_guard(torch.tensor(acc_loss), step, run_dir / "debug", model)
+
+        if step % log_interval == 0 or step == 0:
+            gn = (
+                grad_norm.item()
+                if isinstance(grad_norm, torch.Tensor)
+                else float(grad_norm)
+            )
+            lr_bb = optimizer.param_groups[0]["lr"]
+            lr_hd = optimizer.param_groups[-1]["lr"]
+            metrics_logger.log_scalar("loss/total", acc_loss, step)
+            metrics_logger.log_scalar("lr/backbone", lr_bb, step)
+            metrics_logger.log_scalar("lr/heads", lr_hd, step)
+            for k, v in acc_ph.items():
+                metrics_logger.log_scalar(f"loss/{k}", v, step)
+            parts = " ".join(f"{k}={v:.3f}" for k, v in acc_ph.items())
+            log.info(
+                "step %d/%d | loss=%.4f [%s] | lr=%.2e/%.2e | grad=%.2f",
+                step, total_steps, acc_loss, parts, lr_bb, lr_hd, gn,
+            )
+
+        if (step + 1) % ckpt_interval == 0 or step == total_steps - 1:
+            save_checkpoint(
+                ckpt_dir / f"step_{step:06d}.pt",
+                model, optimizer, scheduler, step, config,
+            )
+
+        if (step + 1) % val_interval == 0 or step == total_steps - 1:
+            val_m = validate(
+                model, val_loader, device, amp_dtype, use_amp,
+            )
+            for k, v in val_m.items():
+                metrics_logger.log_scalar(k, v, step)
+            log.info(
+                "step %d | %s", step,
+                " | ".join(f"{k}={v:.4f}" for k, v in val_m.items()),
+            )
+
+    metrics_logger.close()
+    final = ckpt_dir / f"step_{total_steps - 1:06d}.pt"
+    log.info("Fine-tuning complete [%s]. Final: %s", mode, final)
+    return final
+
+
+if __name__ == "__main__":
+    import argparse
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(name)s %(levelname)s %(message)s",
+    )
+    parser = argparse.ArgumentParser(description="Multi-head fine-tuning")
+    parser.add_argument(
+        "--checkpoint", type=str, default=None,
+        help="Pretrained checkpoint path. Omit for random-init baseline.",
+    )
+    parser.add_argument("--config", type=str, default="configs/finetune.yaml")
+    parser.add_argument("--model-config", type=str, default="configs/model.yaml")
+    parser.add_argument("--schema", type=str, default="data/schema.yaml")
+    parser.add_argument("--data-dir", type=str, default="data/synthetic")
+    parser.add_argument("--output-dir", type=str, default="experiments")
+    parser.add_argument("--max-steps", type=int, default=None)
+    parser.add_argument("--label-fraction", type=float, default=None)
+    args = parser.parse_args()
+
+    finetune(
+        pretrain_checkpoint=args.checkpoint,
+        finetune_config=args.config,
+        model_config=args.model_config,
+        schema_path=args.schema,
+        data_dir=args.data_dir,
+        output_dir=args.output_dir,
+        max_steps=args.max_steps,
+        label_fraction=args.label_fraction,
+    )

src/training/trainer_utils.py

@@ -0,0 +1,265 @@
+"""Shared training utilities for pretraining and fine-tuning.
+
+All training loops import from here. No duplicated optimizer construction,
+scheduler logic, or checkpoint I/O across pretrain.py and finetune.py.
+
+Functions:
+  create_optimizer     — AdamW, weight decay on 2D+ params only
+  create_scheduler     — Linear warmup + cosine decay to min_lr_fraction
+  save_checkpoint      — Atomic write (.tmp then rename)
+  load_checkpoint      — Restore with optional fingerprint verification
+  setup_deterministic  — Seed all RNGs, deterministic algorithms
+  nan_guard            — Per-step finiteness check with debug dump
+
+Class:
+  MetricsLogger        — Tensorboard + optional wandb
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import os
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LambdaLR
+
+log = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Optimizer / Scheduler
+# ---------------------------------------------------------------------------
+
+
+def create_optimizer(
+    model: nn.Module,
+    lr: float = 3e-4,
+    betas: tuple[float, float] = (0.9, 0.95),
+    weight_decay: float = 0.1,
+) -> AdamW:
+    """AdamW with weight decay on matrices only (not norms, not biases)."""
+    decay, no_decay = [], []
+    for param in model.parameters():
+        if not param.requires_grad:
+            continue
+        (decay if param.dim() >= 2 else no_decay).append(param)
+
+    return AdamW(
+        [{"params": decay, "weight_decay": weight_decay},
+         {"params": no_decay, "weight_decay": 0.0}],
+        lr=lr, betas=betas,
+    )
+
+
+def create_scheduler(
+    optimizer: AdamW,
+    warmup_steps: int,
+    total_steps: int,
+    min_lr_fraction: float = 0.1,
+) -> LambdaLR:
+    """Linear warmup then cosine decay to min_lr_fraction * peak LR."""
+
+    def lr_lambda(step: int) -> float:
+        if step < warmup_steps:
+            return step / max(1, warmup_steps)
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return min_lr_fraction + (1 - min_lr_fraction) * 0.5 * (1 + math.cos(math.pi * progress))
+
+    return LambdaLR(optimizer, lr_lambda)
+
+
+# ---------------------------------------------------------------------------
+# Checkpointing
+# ---------------------------------------------------------------------------
+
+
+def save_checkpoint(
+    path: str | Path,
+    model: nn.Module,
+    optimizer: AdamW,
+    scheduler: LambdaLR,
+    step: int,
+    config: dict[str, Any],
+    fingerprint: str = "",
+) -> None:
+    """Atomic save: writes to .tmp then os.replace for crash safety."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(".tmp")
+
+    torch.save({
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "scheduler_state_dict": scheduler.state_dict(),
+        "step": step,
+        "config": config,
+        "fingerprint": fingerprint,
+    }, tmp)
+    os.replace(tmp, path)
+    log.info("Checkpoint saved: %s (step %d)", path, step)
+
+
+def load_checkpoint(
+    path: str | Path,
+    model: nn.Module,
+    optimizer: AdamW | None = None,
+    scheduler: LambdaLR | None = None,
+    expected_fingerprint: str | None = None,
+    strict: bool = True,
+) -> dict[str, Any]:
+    """Load checkpoint. Pass optimizer/scheduler=None to skip restoring them.
+
+    strict=False allows loading pretrained weights into a model with extra
+    parameters (e.g. a fraud head added for fine-tuning).
+    Raises ValueError on fingerprint mismatch.
+    """
+    ckpt = torch.load(path, map_location="cpu", weights_only=False)
+
+    if expected_fingerprint is not None:
+        saved = ckpt.get("fingerprint", "")
+        if saved != expected_fingerprint:
+            raise ValueError(
+                f"Fingerprint mismatch: checkpoint='{saved}', expected='{expected_fingerprint}'"
+            )
+
+    model.load_state_dict(ckpt["model_state_dict"], strict=strict)
+    if optimizer is not None and "optimizer_state_dict" in ckpt:
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+    if scheduler is not None and "scheduler_state_dict" in ckpt:
+        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+
+    log.info("Checkpoint loaded: %s (step %s)", path, ckpt.get("step", "?"))
+    return ckpt
+
+
+# ---------------------------------------------------------------------------
+# Reproducibility
+# ---------------------------------------------------------------------------
+
+
+def setup_deterministic(
+    seed: int = 42,
+    warn_only: bool = True,
+    cublas_workspace_config: str = ":4096:8",
+) -> None:
+    """Seed all RNGs and enable deterministic CUDA kernels.
+
+    warn_only=True because depthwise Conv1d backward may lack a deterministic
+    CUDA kernel �� logs a warning instead of crashing (D19e).
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+    torch.use_deterministic_algorithms(True, warn_only=warn_only)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = cublas_workspace_config
+
+    log.info("Deterministic mode: seed=%d, warn_only=%s", seed, warn_only)
+
+
+# ---------------------------------------------------------------------------
+# NaN guard
+# ---------------------------------------------------------------------------
+
+
+class NaNError(RuntimeError):
+    """Raised when loss contains NaN or Inf."""
+    pass
+
+
+def nan_guard(
+    loss: torch.Tensor,
+    step: int,
+    output_dir: str | Path,
+    model: nn.Module | None = None,
+) -> None:
+    """Check loss is finite. Dumps debug state and raises NaNError on failure."""
+    if torch.isfinite(loss):
+        return
+
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    debug_path = output_dir / f"nan_debug_step{step}.pt"
+
+    state: dict[str, Any] = {"step": step, "loss": loss.detach().cpu()}
+    if model is not None:
+        state["model_state_dict"] = {
+            k: v.detach().cpu() for k, v in model.state_dict().items()
+        }
+
+    torch.save(state, debug_path)
+    log.error("NaN/Inf loss at step %d. Debug state: %s", step, debug_path)
+    raise NaNError(f"Loss is {loss.item()} at step {step}. Debug: {debug_path}")
+
+
+# ---------------------------------------------------------------------------
+# Metrics logging
+# ---------------------------------------------------------------------------
+
+
+class MetricsLogger:
+    """Tensorboard writer + optional wandb. Wandb activates on WANDB_API_KEY."""
+
+    def __init__(
+        self,
+        log_dir: str | Path,
+        experiment_name: str = "run",
+        use_wandb: str = "auto",
+        wandb_project: str = "lfm2-transactions",
+    ) -> None:
+        from torch.utils.tensorboard import SummaryWriter
+
+        self.tb = SummaryWriter(log_dir=str(log_dir))
+        self.wandb_run = None
+
+        if use_wandb == "auto" and os.environ.get("WANDB_API_KEY"):
+            try:
+                import wandb
+
+                self.wandb_run = wandb.init(
+                    project=wandb_project, name=experiment_name, dir=str(log_dir),
+                )
+            except ImportError:
+                log.info("wandb not installed, tensorboard only")
+
+    def log_scalar(self, tag: str, value: float, step: int) -> None:
+        self.tb.add_scalar(tag, value, step)
+        if self.wandb_run is not None:
+            import wandb
+
+            wandb.log({tag: value}, step=step)
+
+    def log_per_feature_losses(
+        self,
+        total_loss: float,
+        per_feature_losses: dict[str, float],
+        step: int,
+    ) -> None:
+        self.log_scalar("loss/total", total_loss, step)
+        for name, val in per_feature_losses.items():
+            self.log_scalar(f"loss/{name}", val, step)
+
+    def log_config(self, config: dict[str, Any]) -> None:
+        self.tb.add_text("config", str(config), 0)
+        if self.wandb_run is not None:
+            import wandb
+
+            wandb.config.update(config)
+
+    def close(self) -> None:
+        self.tb.close()
+        if self.wandb_run is not None:
+            import wandb
+
+            wandb.finish()