Spaces:

cdotsanghvi
/

lfm2-transaction-encoder

Sleeping

lfm2-transaction-encoder / encoder /src /model /transaction_encoder_nocompress.py

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

083b138 5 months ago

4.91 kB

"""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

	"""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