[giovannic]

Summary

This PR adds comprehensive specifications for the Transformer

---

Context

Flow Matching for Posterior Estimation

TFMPE uses flow matching to learn the posterior distribution p(θ|x) for simulation-based inference. Flow matching learns a time-dependent vector field v(θ, x, t) that transports samples from a simple base distribution (e.g., standard
normal) to the complex posterior distribution.

The transformer architecture models this vector field by:

1. Encoding context (observations x) through self-attention encoder blocks
2. Encoding parameters (θ at time t) through self-attention encoder blocks
3. Decoding parameters with cross-attention to context to produce the vector field output

This encoder-decoder design allows the model to capture:

• Complex dependencies within parameter structure (self-attention)
• Conditional relationships between observations and parameters (cross-attention)
• Time-dependent dynamics of the flow (time embedding)

Legacy vs New Design

The legacy transformer required manual management of flattened arrays, labels, indices, and masks. This design leverages the Tokens preprocessing interface which already encapsulates this complexity, providing:

• Unified data structure with data, labels, self_attention_mask, functional_inputs
• Automatic mask generation from independence structure
• Cleaner integration between preprocessing and model layers

---

Design Decisions

1. TokensInterface

Accept Tokens objects directly as inputs instead of separate arrays for data, labels, indices, and masks. The Tokens class already encapsulates all necessary information for the transformer, eliminating ~5 function parameters per
call and enforcing consistency with the preprocessing pipeline.

2. Dataclass Configuration

Use a TransformerConfig dataclass with sensible defaults (latent_dim=128, n_encoder=4, n_decoder=4, etc.) to reduce boilerplate. The dataclass validates constraints like latent_dim % n_heads == 0 in __post_init__ and
groups architecture parameters separately from data-dependent parameters (value_dim, n_labels, index_dim).

3. Performance Benchmarking

Implement two types of automated benchmarks:

Speed benchmarks (pytest-benchmark):

• Forward and backward pass timing with statistical analysis
• Parameterized by model size (n_tokens, latent_dim, n_encoder, n_decoder)
• Uses @pytest.mark.speed for standard configs, @pytest.mark.scale for large configs
• Enforces time thresholds as test assertions to catch performance regressions

Memory benchmarks (pytest-memray):

• Peak memory profiling for forward and backward passes
• Same parameterization as speed benchmarks
• Enforces memory thresholds critical for GPU-constrained environments
• Separate markers allow CI to skip expensive large-scale benchmarks

Both benchmark types ensure the transformer scales efficiently to large hierarchical models and prevent performance degradation over time.