docs/adr/ADR-001-implementation-framework.md

ADR-001: Implementation Framework for domainTokenizer

Status: Accepted Date: April 29, 2026 Decision: PyTorch + HuggingFace Transformers as primary framework, with JAX/Flax NNX as future scaling path Deciders: domainTokenizer core team

---

Table of Contents

1. Context
2. Goal
3. Options Evaluated
4. Decision
5. Trade-offs and Justification
6. Consequences
7. Implementation Roadmap
8. Appendix A: Framework Usage Across Reference Papers
9. Appendix B: Head-to-Head Comparison Matrix
10. Appendix C: Key Code Patterns

---

1. Context

What We're Building

domainTokenizer is a library for building small Transformer models (24M–330M parameters) that process domain-specific tokens — financial transactions, e-commerce events, healthcare records — instead of natural language text. The architecture follows the validated pattern from Nubank's nuFormer (arXiv: 2507.23267):

Domain Events → Custom Tokenizer → GPT-style Transformer → Foundation Model → Downstream Tasks

The Implementation Question

A video from Google for Developers presents the Keras 3 + JAX/Flax NNX integration as a potential framework, citing:

• Explicit state management (Flax NNX) — useful for tracking sequential transaction state
• Custom training loops — Keras structure + JAX/Optax for domain-specific gradient control
• JIT compilation (@nnx.jit) — high-performance processing of millions of transactions
• Paradigm mixing — Keras layers for standard components + NNX for custom sequential encoders

The question: Is Keras + JAX/Flax NNX the right framework for domainTokenizer, or is there a better choice?

Constraints

1. Custom tokenizer required: We need a tokenizer that maps structured fields (amounts, dates, categories) to special tokens — not a standard text tokenizer
2. Small models: 24M–330M parameters, not 70B+ — framework overhead matters less than developer velocity
3. Production deployment: Models must be servable with low latency for real-time applications (fraud detection, recommendations)
4. GPU hardware: Development on A100/A10G GPUs, not TPUs (standard cloud environment)
5. Team context: ML engineers familiar with Python, PyTorch, and the HuggingFace ecosystem
6. Iteration speed: Need to prototype quickly across multiple domains (finance, e-commerce, healthcare)

---

2. Goal

Choose an implementation framework that:

1. Minimizes time from research to working prototype — weeks, not months
2. Supports custom domain tokenizers as first-class citizens
3. Integrates with the HuggingFace Hub for model sharing, versioning, and community
4. Enables production deployment via standard serving infrastructure (ONNX, TGI, vLLM, etc.)
5. Scales to 330M parameters on 4–8 GPUs without heroic engineering
6. Does not preclude future migration to JAX/TPU if we need to scale beyond 1B parameters

---

3. Options Evaluated

Option A: PyTorch + HuggingFace Transformers

The dominant ecosystem for custom NLP/sequential models. Provides PreTrainedModel, PreTrainedTokenizerFast, Trainer, push_to_hub, ONNX export, and integration with TRL, PEFT, Accelerate, DeepSpeed.

Option B: Keras 3 + JAX Backend + Flax NNX

Google's multi-backend framework. Keras provides high-level APIs; JAX provides XLA compilation and functional transforms; Flax NNX provides PyTorch-like stateful modules on top of JAX.

Option C: Pure JAX + Flax NNX + Optax

Skip Keras entirely. Use Flax NNX for model definition, Optax for optimization, Orbax for checkpointing, and Grain/tf.data for data loading. Google's MaxText framework follows this pattern.

Option D: PyTorch + Custom (no HuggingFace)

Use PyTorch directly without the HuggingFace abstraction layer. Full control but no ecosystem integration.

---

4. Decision

Primary: PyTorch + HuggingFace Transformers (Option A)

Future scaling path: JAX/Flax NNX (Option C) — if and when we need TPU training at >1B parameters

---

5. Trade-offs and Justification

5.1 What the Reference Papers Actually Use

We audited the frameworks used by every paper in the domainTokenizer research corpus. The result is overwhelming:

Paper	Framework	Confidence
nuFormer (Nubank)	PyTorch + HF Transformers (inferred)	~90%
TIGER (Google)	JAX + T5X (official); PyTorch (community reimpl)	100%
ActionPiece (Google DeepMind)	PyTorch + HF Transformers (stated verbatim in paper)	100%
RecFormer (UCSD/Amazon)	PyTorch + HF Transformers (Longformer) (stated verbatim)	100%
Banking Transaction Flow	PyTorch (stated verbatim in appendix)	100%
PLR Embeddings (Yandex)	PyTorch + scikit-learn + Optuna	100%

5 of 6 papers use PyTorch. The sole JAX user (TIGER) was a Google-internal project using T5X, and even its most popular community reimplementation (781⭐) is in PyTorch.

Even Google DeepMind's own ActionPiece — the paper most relevant to our domain tokenization approach — uses PyTorch + HuggingFace. This is the strongest signal possible.

5.2 Custom Tokenizer Story

This is the decisive factor. domainTokenizer's core innovation is the tokenizer itself. The framework must provide first-class support for custom token vocabularies.

PyTorch + HuggingFace:

• Train custom BPE tokenizer via tokenizers library (Rust-backed, fast)
• Wrap in PreTrainedTokenizerFast → full Trainer compatibility
• Add domain special tokens via add_special_tokens() → auto-resize embeddings
• Push tokenizer to Hub: tokenizer.push_to_hub("org/my-tokenizer")
• Load anywhere: AutoTokenizer.from_pretrained("org/my-tokenizer")
• KL3M (arXiv: 2503.17247) — the gold standard for financial domain tokenizers — is built entirely on this stack

Keras + JAX/Flax NNX:

• No equivalent to PreTrainedTokenizerFast
• No Hub-integrated tokenizer format
• Must build custom tokenizer from scratch with no ecosystem support
• No standard serialization/deserialization for domain vocabularies

Verdict: PyTorch/HF has a complete, production-tested custom tokenizer pipeline. JAX/Keras has nothing — you'd build everything from scratch.

5.3 Production Deployment

Path	PyTorch	JAX/Keras
ONNX export	torch.onnx.export() — one line	Requires TF backend intermediate or experimental jax.export
TensorRT	ONNX → TRT (standard)	Multi-hop, fragile
TGI (HuggingFace inference)	First-class	Not supported
vLLM	First-class	Not supported
Triton Inference Server	Direct ONNX/TorchScript	Via ONNX (workaround)
BentoML	Supported	Supported
Model Hub sharing	push_to_hub() → from_pretrained()	Works but fragmented (.msgpack weights, no Trainer compat)

Verdict: PyTorch has direct, tested paths to every major serving framework. JAX requires multiple intermediate conversions, each introducing failure points.

5.4 Training Speed

At our scale (24M–330M parameters on 4–8 A100s):

Scenario	PyTorch	JAX
Steady-state training throughput	Comparable (torch.compile)	Comparable (XLA JIT)
Variable-length sequences	Native — dynamic shapes	Problematic — recompiles on new shapes; must pad to buckets
Multi-GPU (FSDP)	accelerate + FSDP2 — mature	pmap/shard_map — works but harder to configure
First-run compilation	Instant (eager mode)	5–20s JIT compilation overhead
Debugging	Standard Python debugger	print debugging; cryptic XLA errors

Verdict: At 330M parameters, training speed is a wash. JAX's advantages (XLA kernel fusion, TPU native) only matter at 10B+ parameters on 256+ accelerators. At our scale, developer velocity dominates throughput.

5.5 The JAX Advantage: When It Would Win

JAX/Flax NNX would be the right choice if:

1. Training exclusively on Google TPUs — JAX is the native TPU compiler; PyTorch/XLA is a port with overhead
2. Models >1B parameters — XLA's whole-program optimization shines at scale
3. Fixed-shape workloads — images, fixed-length token sequences (no variable-length padding issues)
4. Need functional transforms — vmap (per-sample gradients), pmap (data parallelism), grad (higher-order derivatives)
5. Google Cloud infrastructure — Vertex AI, TPU VMs, GCS integration

For domainTokenizer's current scope (24M–330M, GPU, variable-length sequences, fast iteration), none of these conditions apply.

5.6 The Keras + JAX Mixing Argument

The Google for Developers video argues for mixing Keras layers (high-level) with NNX modules (custom, high-performance). In theory, this lets you:

• Use Keras for standard Transformer layers
• Use NNX for custom sequential transaction encoders
• Get JIT compilation on the NNX parts

In practice, this creates problems:

1. Two mental models: Keras (layer-oriented, fit/compile) vs. NNX (functional, explicit state) — context switching slows development
2. Limited interop documentation: Keras ↔ NNX examples are thin; edge cases are poorly documented
3. No HF ecosystem integration: You lose Trainer, push_to_hub, PEFT, TRL, Accelerate — the entire ecosystem Nubank and ActionPiece rely on
4. Debugging complexity: Errors in the Keras↔NNX boundary are hard to diagnose

Better approach with PyTorch: Use torch.compile() on performance-critical modules to get JIT compilation benefits without leaving the PyTorch ecosystem. Write custom nn.Module subclasses for domain-specific components. This gives you the same "standard parts + custom parts" architecture without framework mixing.

---

6. Consequences

What We Gain

1. Immediate access to the entire HuggingFace ecosystem: Trainer, Accelerate, PEFT (LoRA), TRL, Evaluate, push_to_hub, from_pretrained, ONNX export, TGI serving
2. Copy-paste from reference implementations: ActionPiece, RecFormer, Banking TF, and PLR embeddings are all PyTorch — we can directly reuse their code
3. KL3M tokenizer as starting point: The best financial domain tokenizer already exists in PyTorch/HF format at alea-institute/kl3m-004-128k-cased
4. Standard production deployment: ONNX → TensorRT → Triton, or direct TGI/vLLM serving
5. Community and hiring: PyTorch is the dominant ML framework; finding contributors and documentation is easy
6. torch.compile() for performance: When we need JIT compilation on hot paths, torch.compile() provides 10–30% speedups without leaving the ecosystem

What We Accept

1. No native TPU support: If we later need to train on Google TPUs, we'll need PyTorch/XLA (slower than native JAX) or migrate the model code
2. No functional transforms: vmap (per-sample gradients) isn't available without functorch (experimental). If we need advanced gradient manipulation for meta-learning or Nested Learning (HOPE-style), JAX would be better
3. Potential future migration cost: If we scale beyond 1B parameters and move to TPUs, we'll need to rewrite model code in Flax NNX. This is mitigated by keeping model definitions clean and modular

Migration Strategy (If Needed Later)

If domainTokenizer grows to >1B parameters and we need TPU training:

1. Tokenizer layer stays in Python/HF: Tokenizer is framework-agnostic — it produces integer sequences regardless of whether the model is PyTorch or JAX
2. Model architecture translates 1:1: PyTorch nn.Module → Flax NNX nnx.Module mapping is straightforward for standard Transformer components
3. Training loop changes: PyTorch Trainer → custom Flax NNX training loop with Optax
4. Reference: Google's MaxText (github.com/google/maxtext) provides production-grade JAX Transformer patterns we can follow

Estimated migration effort: 2–4 weeks for a clean, well-separated codebase.

---

7. Implementation Roadmap

Phase 2A: Core Tokenizer Library (Weeks 1–3)

Step 1: Domain Schema Definition

Create a declarative schema format that describes the fields in a domain's event data:

# src/tokenizers/schema.py

from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional

class FieldType(Enum):
    NUMERICAL_CONTINUOUS = "numerical_continuous"    # prices, amounts → magnitude bins
    NUMERICAL_DISCRETE = "numerical_discrete"        # quantities → small fixed vocab
    CATEGORICAL_FIXED = "categorical_fixed"          # categories, days of week → direct mapping
    CATEGORICAL_ENTITY = "categorical_entity"        # products, merchants → Semantic IDs (RQ-VAE)
    TEMPORAL = "temporal"                            # timestamps → calendar decomposition
    TEXT = "text"                                     # descriptions → BPE subwords
    SIGN = "sign"                                    # credit/debit → 2 tokens

@dataclass
class FieldSpec:
    name: str
    field_type: FieldType
    vocab_size: Optional[int] = None          # for fixed categorical
    n_bins: int = 21                          # for numerical (Nubank uses 21)
    calendar_fields: List[str] = field(       # for temporal
        default_factory=lambda: ["month", "dow", "dom", "hour"]
    )

@dataclass
class DomainSchema:
    name: str                                 # e.g., "ecommerce", "finance"
    fields: List[FieldSpec]                   # ordered list of fields per event
    
    @property
    def special_token_count(self) -> int:
        """Total domain-specific special tokens needed."""
        count = 0
        for f in self.fields:
            if f.field_type == FieldType.SIGN:
                count += 2
            elif f.field_type == FieldType.NUMERICAL_CONTINUOUS:
                count += f.n_bins
            elif f.field_type == FieldType.CATEGORICAL_FIXED:
                count += f.vocab_size
            elif f.field_type == FieldType.TEMPORAL:
                count += sum({
                    "month": 12, "dow": 7, "dom": 31, "hour": 24,
                    "quarter": 4, "year": 10
                }.get(cf, 0) for cf in f.calendar_fields)
        return count

# Example: Nubank-style financial schema
FINANCE_SCHEMA = DomainSchema(
    name="finance",
    fields=[
        FieldSpec("amount_sign", FieldType.SIGN),
        FieldSpec("amount", FieldType.NUMERICAL_CONTINUOUS, n_bins=21),
        FieldSpec("timestamp", FieldType.TEMPORAL,
                  calendar_fields=["month", "dow", "dom", "hour"]),
        FieldSpec("description", FieldType.TEXT),
    ]
)

# Example: E-commerce schema  
ECOMMERCE_SCHEMA = DomainSchema(
    name="ecommerce",
    fields=[
        FieldSpec("event_type", FieldType.CATEGORICAL_FIXED, vocab_size=5),
        FieldSpec("price", FieldType.NUMERICAL_CONTINUOUS, n_bins=21),
        FieldSpec("quantity", FieldType.NUMERICAL_DISCRETE, vocab_size=11),
        FieldSpec("category_l1", FieldType.CATEGORICAL_FIXED, vocab_size=30),
        FieldSpec("category_l2", FieldType.CATEGORICAL_FIXED, vocab_size=200),
        FieldSpec("timestamp", FieldType.TEMPORAL,
                  calendar_fields=["month", "dow", "dom", "hour"]),
        FieldSpec("product_title", FieldType.TEXT),
    ]
)

Step 2: Per-Field Tokenizers

Implement each field type tokenizer as a standalone module:

# src/tokenizers/field_tokenizers.py

import numpy as np
from typing import List

class SignTokenizer:
    """Tokenizes sign of a numerical value (credit/debit, inflow/outflow)."""
    
    def __init__(self, prefix: str = "SIGN"):
        self.tokens = [f"[{prefix}_POS]", f"[{prefix}_NEG]"]
    
    def __call__(self, value: float) -> str:
        return self.tokens[0] if value >= 0 else self.tokens[1]
    
    @property
    def vocab(self) -> List[str]:
        return self.tokens


class MagnitudeBucketTokenizer:
    """Quantizes continuous values into bins (Nubank-style).
    
    Uses quantile-based binning on the training distribution.
    Follows the Relative Magnitude Tokenization principle from TP-BERTa.
    """
    
    def __init__(self, n_bins: int = 21, prefix: str = "AMT"):
        self.n_bins = n_bins
        self.prefix = prefix
        self.bin_edges = None  # fitted from data
    
    def fit(self, values: np.ndarray):
        """Compute bin edges from training data using quantiles."""
        # Use absolute values for magnitude binning
        abs_vals = np.abs(values[~np.isnan(values)])
        quantiles = np.linspace(0, 100, self.n_bins + 1)
        self.bin_edges = np.percentile(abs_vals, quantiles)
        return self
    
    def __call__(self, value: float) -> str:
        if self.bin_edges is None:
            raise ValueError("Tokenizer not fitted. Call .fit() first.")
        bin_idx = np.searchsorted(self.bin_edges[1:-1], abs(value))
        return f"[{self.prefix}_{bin_idx:02d}]"
    
    @property
    def vocab(self) -> List[str]:
        return [f"[{self.prefix}_{i:02d}]" for i in range(self.n_bins)]


class CalendarTokenizer:
    """Decomposes timestamps into calendar components (Nubank-style)."""
    
    FIELD_VOCABS = {
        "month": ([f"[MON_{i:02d}]" for i in range(1, 13)], lambda dt: dt.month - 1),
        "dow":   ([f"[DOW_{i}]" for i in range(7)], lambda dt: dt.weekday()),
        "dom":   ([f"[DOM_{i:02d}]" for i in range(1, 32)], lambda dt: dt.day - 1),
        "hour":  ([f"[HOUR_{i:02d}]" for i in range(24)], lambda dt: dt.hour),
    }
    
    def __init__(self, fields: List[str] = None):
        self.fields = fields or ["month", "dow", "dom", "hour"]
    
    def __call__(self, timestamp) -> List[str]:
        tokens = []
        for field_name in self.fields:
            vocab, extractor = self.FIELD_VOCABS[field_name]
            idx = extractor(timestamp)
            tokens.append(vocab[idx])
        return tokens
    
    @property
    def vocab(self) -> List[str]:
        all_tokens = []
        for field_name in self.fields:
            all_tokens.extend(self.FIELD_VOCABS[field_name][0])
        return all_tokens


class CategoricalTokenizer:
    """Maps categorical values to fixed vocabulary tokens."""
    
    def __init__(self, categories: List[str], prefix: str = "CAT"):
        self.prefix = prefix
        self.token_map = {cat: f"[{prefix}_{i:03d}]" for i, cat in enumerate(categories)}
        self.unk_token = f"[{prefix}_UNK]"
    
    def __call__(self, value: str) -> str:
        return self.token_map.get(value, self.unk_token)
    
    @property
    def vocab(self) -> List[str]:
        return list(self.token_map.values()) + [self.unk_token]

Step 3: Composite Domain Tokenizer

Assemble per-field tokenizers into a complete domain tokenizer, wrapped as PreTrainedTokenizerFast:

# src/tokenizers/domain_tokenizer.py

from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from transformers import PreTrainedTokenizerFast

class DomainTokenizerBuilder:
    """Builds a HuggingFace-compatible tokenizer from a DomainSchema."""
    
    def __init__(self, schema: DomainSchema):
        self.schema = schema
        self.field_tokenizers = {}  # name → field tokenizer instance
        self._build_field_tokenizers()
    
    def _build_field_tokenizers(self):
        for field_spec in self.schema.fields:
            if field_spec.field_type == FieldType.SIGN:
                self.field_tokenizers[field_spec.name] = SignTokenizer(field_spec.name.upper())
            elif field_spec.field_type == FieldType.NUMERICAL_CONTINUOUS:
                self.field_tokenizers[field_spec.name] = MagnitudeBucketTokenizer(
                    n_bins=field_spec.n_bins, prefix=field_spec.name.upper()
                )
            elif field_spec.field_type == FieldType.TEMPORAL:
                self.field_tokenizers[field_spec.name] = CalendarTokenizer(field_spec.calendar_fields)
            # ... other types
    
    def fit(self, data):
        """Fit data-dependent tokenizers (magnitude bins, etc.)."""
        for field_spec in self.schema.fields:
            if field_spec.field_type == FieldType.NUMERICAL_CONTINUOUS:
                values = [getattr(event, field_spec.name) for event in data]
                self.field_tokenizers[field_spec.name].fit(np.array(values))
        return self
    
    def build_hf_tokenizer(self, text_corpus=None, bpe_vocab_size=8000) -> PreTrainedTokenizerFast:
        """Build a complete HuggingFace tokenizer.
        
        1. Collect all domain special tokens
        2. Train BPE on text fields (if any)
        3. Merge into a single PreTrainedTokenizerFast
        """
        # Collect all special tokens from field tokenizers
        all_special_tokens = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[BOS]", "[EOS]"]
        for name, tok in self.field_tokenizers.items():
            if hasattr(tok, 'vocab'):
                all_special_tokens.extend(tok.vocab)
        
        # Train BPE on text fields
        bpe_tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))
        bpe_tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
        trainer = trainers.BpeTrainer(
            vocab_size=bpe_vocab_size,
            special_tokens=all_special_tokens,
            min_frequency=2,
        )
        if text_corpus:
            bpe_tokenizer.train_from_iterator(text_corpus, trainer=trainer)
        
        # Wrap as HuggingFace tokenizer
        hf_tokenizer = PreTrainedTokenizerFast(
            tokenizer_object=bpe_tokenizer,
            bos_token="[BOS]",
            eos_token="[EOS]",
            pad_token="[PAD]",
            unk_token="[UNK]",
            mask_token="[MASK]",
        )
        
        return hf_tokenizer
    
    def tokenize_event(self, event) -> List[str]:
        """Convert a single domain event into a list of token strings."""
        tokens = []
        for field_spec in self.schema.fields:
            value = getattr(event, field_spec.name, None)
            if value is None:
                tokens.append("[UNK]")
                continue
            tok = self.field_tokenizers[field_spec.name]
            result = tok(value)
            if isinstance(result, list):
                tokens.extend(result)
            else:
                tokens.append(result)
        return tokens

Phase 2B: Model Architecture (Weeks 3–5)

Step 4: GPT-style Causal Transformer (NoPE)

Implement as a HuggingFace-compatible PreTrainedModel:

# src/models/configuration_domain_transformer.py

from transformers import PretrainedConfig

class DomainTransformerConfig(PretrainedConfig):
    model_type = "domain_transformer"
    
    def __init__(
        self,
        vocab_size: int = 32000,
        hidden_size: int = 256,        # 256 = 24M params, 1024 = 330M (Nubank sizes)
        num_hidden_layers: int = 24,    # Nubank uses 24 for both sizes
        num_attention_heads: int = 16,  # Nubank uses 16 for both sizes
        intermediate_size: int = None,  # defaults to 4 * hidden_size
        max_position_embeddings: int = 2048,
        dropout: float = 0.1,
        use_positional_encoding: bool = False,  # NoPE by default!
        **kwargs
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size or 4 * hidden_size
        self.max_position_embeddings = max_position_embeddings
        self.dropout = dropout
        self.use_positional_encoding = use_positional_encoding
        super().__init__(**kwargs)

# src/models/modeling_domain_transformer.py

import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

class DomainTransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.hidden_size)
        self.attn = nn.MultiheadAttention(
            config.hidden_size, config.num_attention_heads,
            dropout=config.dropout, batch_first=True
        )
        self.ln2 = nn.LayerNorm(config.hidden_size)
        self.mlp = nn.Sequential(
            nn.Linear(config.hidden_size, config.intermediate_size),
            nn.GELU(),
            nn.Linear(config.intermediate_size, config.hidden_size),
            nn.Dropout(config.dropout),
        )
    
    def forward(self, x, attn_mask=None):
        # Pre-norm architecture
        h = self.ln1(x)
        h, _ = self.attn(h, h, h, attn_mask=attn_mask, is_causal=True)
        x = x + h
        x = x + self.mlp(self.ln2(x))
        return x


class DomainTransformerForCausalLM(PreTrainedModel):
    config_class = DomainTransformerConfig
    
    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        
        # NoPE: no positional encoding by default (Kazemnejad et al. 2023)
        if config.use_positional_encoding:
            self.embed_positions = nn.Embedding(
                config.max_position_embeddings, config.hidden_size
            )
        else:
            self.embed_positions = None
        
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            DomainTransformerBlock(config)
            for _ in range(config.num_hidden_layers)
        ])
        self.ln_f = nn.LayerNorm(config.hidden_size)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        # Weight tying (standard for small models)
        self.lm_head.weight = self.embed_tokens.weight
        
        self.post_init()
    
    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        x = self.embed_tokens(input_ids)
        
        if self.embed_positions is not None:
            positions = torch.arange(input_ids.size(1), device=input_ids.device)
            x = x + self.embed_positions(positions)
        
        x = self.drop(x)
        
        for block in self.blocks:
            x = block(x, attn_mask=attention_mask)
        
        x = self.ln_f(x)
        logits = self.lm_head(x)
        
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = nn.functional.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        
        return CausalLMOutputWithPast(loss=loss, logits=logits)

# Register with AutoClass for Hub compatibility
DomainTransformerConfig.register_for_auto_class()
DomainTransformerForCausalLM.register_for_auto_class("AutoModelForCausalLM")

Step 5: PLR Numerical Embeddings (for Joint Fusion)

Port from Yandex's implementation:

# src/models/plr_embeddings.py

import torch
import torch.nn as nn
import math

class PeriodicLinearReLU(nn.Module):
    """PLR numerical embeddings (Gorishniy et al. 2022).
    
    Maps scalar x → [sin(2π·w·x + b), cos(2π·w·x + b)] → Linear → ReLU
    Frequencies w and phases b are LEARNED parameters.
    """
    
    def __init__(self, n_features: int, n_frequencies: int = 64, embedding_dim: int = 64):
        super().__init__()
        self.n_features = n_features
        self.n_frequencies = n_frequencies
        
        # Learnable frequencies and phases (per feature)
        self.frequencies = nn.Parameter(
            torch.randn(n_features, n_frequencies) * 0.01
        )
        self.phases = nn.Parameter(
            torch.zeros(n_features, n_frequencies)
        )
        
        # Linear projection: 2*n_frequencies → embedding_dim
        self.linear = nn.Linear(2 * n_frequencies, embedding_dim)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch, n_features) — raw scalar feature values
        Returns:
            (batch, n_features, embedding_dim)
        """
        # x: (B, F) → (B, F, 1)
        x = x.unsqueeze(-1)
        
        # Periodic encoding: (B, F, n_freq)
        angles = 2 * math.pi * self.frequencies.unsqueeze(0) * x + self.phases.unsqueeze(0)
        periodic = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)  # (B, F, 2*n_freq)
        
        # Linear + ReLU: (B, F, embedding_dim)
        return torch.relu(self.linear(periodic))

Phase 2C: Pre-training (Weeks 5–7)

Step 6: Data Pipeline

# src/training/data_pipeline.py

from torch.utils.data import Dataset
from typing import List

class DomainSequenceDataset(Dataset):
    """Converts user event sequences into token sequences for CLM training."""
    
    def __init__(self, user_sequences, tokenizer_builder, hf_tokenizer, max_length=2048):
        self.user_sequences = user_sequences
        self.tokenizer_builder = tokenizer_builder
        self.hf_tokenizer = hf_tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.user_sequences)
    
    def __getitem__(self, idx):
        events = self.user_sequences[idx]
        
        # Tokenize each event into token strings
        token_strings = []
        for event in events:
            event_tokens = self.tokenizer_builder.tokenize_event(event)
            token_strings.extend(event_tokens)
        
        # Convert token strings to IDs via HF tokenizer
        encoding = self.hf_tokenizer(
            " ".join(token_strings),
            max_length=self.max_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        
        input_ids = encoding["input_ids"].squeeze(0)
        
        return {
            "input_ids": input_ids,
            "labels": input_ids.clone(),  # CLM: labels = input shifted by 1
            "attention_mask": encoding["attention_mask"].squeeze(0),
        }

Step 7: Pre-training with HuggingFace Trainer

# src/training/pretrain.py

from transformers import Trainer, TrainingArguments

def pretrain_domain_model(
    model,
    train_dataset,
    eval_dataset=None,
    output_dir="./checkpoints",
    hub_model_id="org/domain-model-24m",
    num_epochs=3,
    batch_size=64,
    learning_rate=3e-4,
    context_length=2048,
):
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=4,
        learning_rate=learning_rate,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        weight_decay=0.01,
        logging_strategy="steps",
        logging_steps=100,
        logging_first_step=True,
        disable_tqdm=True,           # plain text logging for cloud
        eval_strategy="steps" if eval_dataset else "no",
        eval_steps=500,
        save_strategy="steps",
        save_steps=1000,
        save_total_limit=3,
        push_to_hub=True,
        hub_model_id=hub_model_id,
        bf16=True,
        dataloader_num_workers=4,
        report_to="trackio",
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
    )
    
    trainer.train()
    trainer.push_to_hub()

Phase 2D: Joint Fusion Fine-tuning (Weeks 7–9)

Step 8: nuFormer-style Joint Fusion

# src/models/joint_fusion.py

import torch
import torch.nn as nn

class DCNv2CrossLayer(nn.Module):
    """Single cross layer from DCN V2 (Wang et al. 2021)."""
    
    def __init__(self, dim):
        super().__init__()
        self.weight = nn.Linear(dim, dim, bias=True)
    
    def forward(self, x0, x):
        return x0 * self.weight(x) + x  # element-wise product with anchor


class JointFusionModel(nn.Module):
    """nuFormer-style: Transaction Transformer + DCNv2(PLR) → Joint Prediction.
    
    Architecture:
        Transaction Sequence → Pre-trained DomainTransformer → user_embedding
        Tabular Features → PLR → DCNv2 → tab_embedding
        Concatenate → MLP Head → prediction
    """
    
    def __init__(self, transformer_model, n_tabular_features, n_classes=1,
                 plr_frequencies=64, dcn_layers=3, hidden_dim=256):
        super().__init__()
        
        self.transformer = transformer_model  # pre-trained, unfrozen for fine-tuning
        transformer_dim = transformer_model.config.hidden_size
        
        # Tabular branch: PLR → DCNv2
        self.plr = PeriodicLinearReLU(n_tabular_features, plr_frequencies, hidden_dim)
        tab_input_dim = n_tabular_features * hidden_dim
        
        self.dcn_layers = nn.ModuleList([
            DCNv2CrossLayer(tab_input_dim) for _ in range(dcn_layers)
        ])
        self.dcn_deep = nn.Sequential(
            nn.Linear(tab_input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        
        # Joint head
        self.head = nn.Sequential(
            nn.Linear(transformer_dim + hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, n_classes),
        )
    
    def forward(self, input_ids, attention_mask, tabular_features, labels=None):
        # Transaction branch: get last-token embedding
        transformer_output = self.transformer(input_ids, attention_mask=attention_mask)
        user_embedding = transformer_output.logits[:, -1, :]  # last token representation
        
        # Tabular branch: PLR → flatten → DCNv2
        tab_embedded = self.plr(tabular_features)  # (B, F, D)
        tab_flat = tab_embedded.view(tab_embedded.size(0), -1)  # (B, F*D)
        
        x0 = tab_flat
        x = tab_flat
        for cross_layer in self.dcn_layers:
            x = cross_layer(x0, x)
        tab_output = self.dcn_deep(x)  # (B, hidden_dim)
        
        # Joint fusion
        combined = torch.cat([user_embedding, tab_output], dim=-1)
        logits = self.head(combined)
        
        loss = None
        if labels is not None:
            loss = nn.functional.binary_cross_entropy_with_logits(logits.squeeze(-1), labels.float())
        
        return {"loss": loss, "logits": logits}

Phase 3: Domain Demos (Weeks 9–12)

Week	Deliverable	Hardware
9–10	Finance demo: Transaction tokenizer + 24M model pre-trained on synthetic/public financial data + fraud detection fine-tuning	a10g-large
10–11	E-commerce demo: Event tokenizer + 24M model pre-trained on Amazon review sequences + next-purchase prediction	a10g-large
11–12	Evaluation & benchmarking: Compare domain tokenizer vs. text serialization vs. LightGBM baselines on each domain	a10g-large

Phase 4: Scale & Optimize (Weeks 12+)

Task	Details
Scale to 330M params	Increase hidden_size to 1024, train on a100-large
torch.compile()	Apply to attention and MLP blocks for 10–30% speedup
ONNX export	torch.onnx.export() for production serving
Context window experiments	Ablate 512/1024/2048/4096 context lengths
Data source ablation	Test impact of different event types (Nubank found adding low-signal sources hurts)
ActionPiece vocabulary	Implement BPE-like cross-field merging on top of per-field tokens

---

Appendix A: Framework Usage Across Reference Papers

Paper	ArXiv	Framework	Verbatim Evidence
nuFormer (Nubank)	2507.23267	PyTorch + HF (inferred)	All dependencies are PyTorch-based
TIGER (Google)	2305.05065	JAX + T5X	"We use the open-sourced T5X framework"
ActionPiece (DeepMind)	2502.13581	PyTorch + HF	"HuggingFace Transformers and PyTorch" (Appendix H)
RecFormer	2305.13731	PyTorch + HF Longformer	"Longformer implemented by Huggingface" (§3.1.4)
Banking TF	2410.08243	PyTorch	"Pytorch backend is used" (Appendix B)
PLR Embeddings (Yandex)	2203.05556	PyTorch	Repository: pure PyTorch + scikit-learn
KL3M Tokenizers	2503.17247	HF tokenizers + PyTorch	"tokenizers" BPE for HF compatibility

---

Appendix B: Head-to-Head Comparison Matrix

Criterion	PyTorch + HF	JAX/Flax NNX	Keras 3 + JAX
Custom domain tokenizer	✅ PreTrainedTokenizerFast	❌ Build from scratch	❌ Build from scratch
HF Trainer integration	✅ Native	❌ Not compatible	❌ Not compatible
Hub push/pull	✅ push_to_hub()	⚠️ Works, fragmented	⚠️ Limited
PEFT/LoRA	✅ Drop-in	❌ Manual	❌ Manual
ONNX export	✅ One-line	❌ Multi-hop	⚠️ TF backend required
TGI/vLLM serving	✅ First-class	❌ Not supported	❌ Not supported
TPU training	⚠️ PyTorch/XLA (overhead)	✅ Native	✅ Native
JIT compilation	✅ torch.compile()	✅ @nnx.jit	✅ XLA via JAX
Dynamic shapes (NLP)	✅ Native	❌ Recompiles	❌ Recompiles
Debugging	✅ Eager mode, std debugger	⚠️ Challenging	⚠️ Challenging
Reference implementations	5/6 papers	1/6 papers	0/6 papers
Community/hiring pool	🟢 Large	🟡 Small	🟡 Small

---

Appendix C: Key Code Patterns

Adding Domain Special Tokens to an Existing Tokenizer

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")  # or any base tokenizer

# Add all domain special tokens
special_tokens = {
    "additional_special_tokens": [
        # Amount tokens (Nubank-style)
        "[AMT_POS]", "[AMT_NEG]",
        *[f"[AMT_{i:02d}]" for i in range(21)],
        # Calendar tokens  
        *[f"[MON_{i:02d}]" for i in range(1, 13)],
        *[f"[DOW_{i}]" for i in range(7)],
        *[f"[DOM_{i:02d}]" for i in range(1, 32)],
        *[f"[HOUR_{i:02d}]" for i in range(24)],
    ]
}
num_added = tokenizer.add_special_tokens(special_tokens)
print(f"Added {num_added} domain tokens. Vocab size: {len(tokenizer)}")

# CRITICAL: resize model embeddings
model.resize_token_embeddings(len(tokenizer))

Registering a Custom Model for Hub Deployment

# In your package's __init__.py or a registration script:
from transformers import AutoConfig, AutoModelForCausalLM

from .configuration_domain_transformer import DomainTransformerConfig
from .modeling_domain_transformer import DomainTransformerForCausalLM

# Register so AutoClass can find your model
AutoConfig.register("domain_transformer", DomainTransformerConfig)
AutoModelForCausalLM.register(DomainTransformerConfig, DomainTransformerForCausalLM)

# Enable push_to_hub with custom code
DomainTransformerConfig.register_for_auto_class()
DomainTransformerForCausalLM.register_for_auto_class("AutoModelForCausalLM")

# Push: uploads configuration.py, modeling.py, config.json, model.safetensors
model.push_to_hub("org/domain-transformer-24m")

# Load anywhere:
model = AutoModelForCausalLM.from_pretrained("org/domain-transformer-24m", trust_remote_code=True)

---

This ADR is a living document and will be updated as implementation progresses.