A Tabular Foundation Model That Learns Joint Distributions, Not Just Predictions

Tabular data is the unglamorous backbone of machine learning in practice. Healthcare records, financial transactions, sensor readings, scientific measurements, the vast majority of real-world prediction tasks involve structured tables where rows are samples and columns are features. Yet the dominant approach to tabular prediction remains task-specific: train a gradient-boosted tree or a neural network on each dataset independently, tune hyperparameters separately, and discard everything when the next dataset arrives. Foundation models for tabular data aim to change this, but most follow a prior-data fitted network approach that learns to predict a target column given context. A new model from Stable AI and Tsinghua University takes a fundamentally different approach.

LimiX-2 introduces what the authors call Contextual Mechanism Networks, a paradigm that shifts the modeling objective from predicting a single target to learning the joint distribution of all variables in a table. Instead of optimizing for p(y | x, D_context), it models p(x, y | D_context), the joint structure underlying how the data was generated. The result is a model that outperforms existing tabular foundation models and dataset-specific baselines on three major benchmarks, while also discovering causal relationships between features as a byproduct of its architecture.

Why Target-Centric Prediction Hits a Ceiling

Conventional tabular prior-data fitted networks treat each prediction task as an isolated problem. Given a labeled context set and a query, they predict the target. This works for supervised classification and regression, but it leaves a lot on the table. The joint distribution of all variables in a table encodes far more information than any single conditional. Two variables might share a causal relationship, a feature might be redundant because it's a noisy copy of another, or the target might depend on an interaction between features that only becomes visible when you model the full dependency structure.

By restricting supervision to a single target column per task, conventional approaches miss these inter-variable dependencies during pretraining. They can impute missing values only as a separate capability, not as something integrated into the core representation. And they can't easily generalize to tasks beyond classification and regression, like conditional reasoning or causal discovery, because their representations don't capture the joint structure that makes those tasks possible.

LimiX-2's approach is to make joint distribution modeling the explicit pretraining objective. The model learns from multiple conditional prediction tasks over the same dataset, using context to capture dependencies shared across them. Supervised prediction becomes a special case of a broader framework for inferring unobserved quantities from available evidence.

Cell-Level Representations with Dual-Axis Attention

The architecture is built around a key design choice: cell-level representations rather than row-level. Most tabular neural networks compress each row into a single vector, losing the fine-grained structure of individual feature values. LimiX-2 keeps every cell as a separate representation, which preserves the ability to model local relations among features within a row and to derive dataset-level statistics from the corresponding cells in context.

The input is a tensor of shape N x F x d, where N is the number of rows, F is the number of features, and d is the embedding dimension (set to 256, up from 192 in the previous LimiX-16M). Each raw cell is mapped into this space through a two-layer MLP with RMSNorm and GELU activation. Missing cells share a single learnable embedding, regardless of column identity. Column identity is provided separately through discriminative feature encoding: each column gets a low-rank code of dimension s = d/4 = 64, which is projected into the embedding space and added to the cell representation. This compression into a compact subspace forces the model to recognize columns by their statistical properties rather than their positions, preventing shortcut memorization.

The backbone is a stack of 24 dual-axis transformer blocks. Each block applies two types of attention in sequence. Sample-axis attention propagates information across rows at each feature position and at the target position. Context rows attend to each other freely, while query rows can only attend to context, preventing information leakage between test samples. The query/key/value mappings are shared among features but not between features and targets, so a query prediction depends only on the sample's own features and the context.

Feature-axis attention is asymmetric: feature representations can attend to both other features and the target, but target representations can only attend to features. This directed information flow means the target readout draws on feature relationships without contaminating the feature representations with target-specific information. The key/value mappings for feature and target attention are not shared, giving the model implicit role separation within a single stream.

A gated feed-forward network (SwiGLU) operates independently on feature and target representations, with separate parameter sets. This replaces the shared MLP from the previous LimiX, providing richer task-specific processing. Multi-head attention is applied across all heads, with a length-dependent rescaling factor per head that adapts to variable context sizes:

s_h = (1 + w_h * log(n)) * beta_h

where n is the sequence length and w_h, beta_h are learnable parameters. This stabilizes attention magnitudes across different context sizes, which is critical when the number of labeled examples varies widely across datasets.

Prediction heads are attached at different depths. Masked feature reconstruction uses shallow-layer representations, where local data details are preserved. Classification and regression are decoded from the last layer's target embeddings. Each task gets an independent bottleneck post-adapter. For regression, instead of MSE loss, the model predicts a probability distribution over 5,000 ordered bins and derives the regression value as the weighted sum of bin centers.

Context-Conditional Masked Modeling

The pretraining objective, Context-Conditional Masked Modeling, combines target prediction with masked feature reconstruction under varied observation patterns. Each pretraining episode partitions a synthetic table into context and query sets. For each query row, some feature columns are masked, and the model must reconstruct them from the observed features and the context. The query target is simultaneously predicted from the same conditioning information.

This is denser supervision than standard label prediction. Instead of learning one conditional per task, the model learns inter-variable dependencies across multiple masking patterns and conditioning sets. The masked view uses a shared missing-value embedding plus the column's discriminative feature encoding, so masked and naturally missing cells share a missingness encoding while retaining column identity.

Three masking schemes are combined: individual entries, selected columns across query rows, and blocks of entries. This varies the granularity of prediction targets and the available conditioning information, exposing the model to tasks ranging from recovering isolated values to predicting entire columns to reconstructing groups of missing entries. The interleaving prevents specialization to any single reconstruction setting.

Synthetic Data from Structural Causal Models

LimiX-2 is pretrained exclusively on synthetic datasets generated through structural causal models. The generation pipeline has five stages: hyperparameter sampling, directed acyclic graph generation, SCM propagation, feature and target sampling, and task adaptation.

Graphs are built hierarchically from local causal motifs, including chain, confounding, and collider structures, composed through recursive expansion. Topology-constrained transformations, like edge redirection and local path replacement, further diversify connectivity patterns while preserving acyclicity. Root nodes are sampled from random distributions, and remaining nodes are computed by propagating functional mechanisms along the topological order. The mechanisms include MLPs, CNNs, decision trees, linear mappings, kernel functions, piecewise functions, periodic functions, and multiplicative interactions, with aggregation strategies ranging from simple averaging to neural aggregation.

Task adaptation applies stochastic observation transformations to features and targets. For classification, continuous targets are discretized into random intervals with varying class frequencies and numbers of classes. For regression, targets undergo random scale transformations and skewness adjustments. The resulting datasets span diverse graph structures, functional mechanisms, and observation processes, providing broad coverage of the space of tabular prediction tasks.

Results: Top Elo on Three Benchmarks

LimiX-2 is evaluated on TabArena (51 datasets), TALENT (288 datasets), and BCCO (156 datasets), covering classification and regression across diverse sample sizes, feature dimensionalities, and missingness regimes. The evaluation uses Bradley-Terry Elo modeling with 2,000 bootstrap rounds for confidence intervals.

On TabArena, LimiX-2 achieves an Elo of 1935, exceeding TabFM+ (1818) by 117 points, AutoGluon 1.6 noncommercial 4-hour (1789) by 146 points, and all other compared methods. It holds the highest Elo at every metric: overall, classification (1917), and regression (2206). Its average rank is 5.5 versus 9.0 for TabFM+, with 18.9 aggregated wins versus 5.3. Pairwise win rates against every competing method exceed 60%, reaching 95% or higher against tuned tree and neural baselines.

On TALENT, LimiX-2 reaches an Elo of 1506, surpassing TabFM (1471) by 35 points and AutoGluon 1.6 (1438) by 68 points. Its improvability, the potential for further improvement over the best achievable, is 6.75% versus 9.17% for TabFM, a 26.4% relative reduction. It achieves 84.3 aggregated wins versus 50.1 for TabFM. The gap is largest in regression (1584 vs 1529) and smallest in multiclass classification (1520 vs 1517).

On BCCO, which focuses on challenging and incomplete datasets, LimiX-2 reaches 1432 Elo, exceeding AutoGluon (1376) by 56 points and TabFM (1369) by 63 points. Its regression Elo of 1859 is particularly strong, beating TabFM by 74 points and TabPFN-3 by 168 points. It achieves 50.4 aggregated wins versus 16.7 for TabFM, roughly three times as many.

Across all three benchmarks, LimiX-2 outperforms every compared method including TabFM, TabPFN-3, TabICLv2, EXAONE Tabular, Mitra-v2, Xiaomi-TabLDM, TabDPT, AutoGluon, XGBoost, LightGBM, CatBoost, and Random Forest. Notably, it does this while being four times smaller in parameter count than TabFM.

Causal Skeleton Recovery as a Byproduct

Beyond prediction, the CMN paradigm gives LimiX-2 an unexpected capability: causal discovery. Because the feature-axis attention encodes direct causal relationships between variables, the attention patterns can be interpreted as a causal skeleton, the undirected graph of direct dependencies.

Evaluated on six causal discovery benchmarks (Sachs, UF, CausalChamber, PATHFINDER, DIABETES, PIGS), LimiX-2 achieves a mean F1 score of 0.7972 across all datasets, ranking first on every one. It also achieves the lowest structural Hamming distance on five of six datasets. The next-best foundation model, EXAONE Tabular, reaches a mean F1 of 0.6591. Other foundation models like TabFM, TabICLv2, TabPFN-3, and Xiaomi-TabLDM perform substantially worse because they compute feature attention only at the group level, not the cell level, which prevents them from distinguishing causal relationship strength among features at fine granularity.

Dedicated causal discovery methods like PC, GES, and NOTEARS-MLP often produce timeouts or worse performance on the discrete datasets. LimiX-2's skeletons match or surpass those of methods designed exclusively for causal discovery, suggesting that the cell-level, mechanism-oriented representation captures genuine structural information about the data-generating process.

Scaling Laws: No Saturation at 406M Parameters

The authors fit log-linear scaling models of the form Elo = alpha + beta * log2(N / 100), where N is the parameter count in millions. The fit is clean across all five evaluation series, with R-squared values between 0.9617 and 0.9808 and residual RMSE between 3.84 and 9.03 Elo points.

The steepest scaling occurs on TabArena, at 34.68 Elo per parameter doubling. BCCO regression follows at 30.06 Elo per doubling, TALENT classification at 22.16, TALENT regression at 18.26, and BCCO classification at 11.24. Across the range from 12.5M to 406.2M parameters, the model gains 169 Elo on TabArena, 109 on TALENT classification, 89 on TALENT regression, 56 on BCCO classification, and 147 on BCCO regression.

No saturation is observed at the largest scale. The authors describe this as strong empirical motivation for extending into the billion-parameter regime. The previous LimiX-16M model achieves 1345 Elo on TabArena; LimiX-2 at 406.2M reaches 1935, a gain of 590 Elo from a 25x increase in parameters.

Limitations and Open Questions

LimiX-2 is pretrained entirely on synthetic data. The generation engine is designed to cover diverse structures and mechanisms, but it cannot perfectly replicate the distribution of real-world tabular datasets. Whether further expansion of the synthetic space would continue to improve performance, or whether a phase of real-data pretraining is needed, remains an open question.

The model requires no task-specific fine-tuning, but its performance on any individual dataset may still be below a well-tuned task-specific model. The improvability metric quantifies this gap: on TabArena, LimiX-2's improvability is 3.3%, meaning there is still room for a hypothetical oracle to improve on the best achievable per-dataset performance. This is lower than most baselines but not zero.

The causal skeleton recovery results are promising but preliminary. The evaluation uses attention weights as a proxy for causal relationships, which is an indirect approach. Whether this mechanism generalizes to larger, more complex causal structures, or whether it can be extended to directed causal graphs rather than just skeletons, is not established.

The scaling law analysis covers up to 406M parameters. Whether the log-linear trend continues into the billion-parameter regime, and whether the same architectural choices remain optimal at that scale, is an empirical question the authors flag for future work.

What This Means for Structured-Data Practice

For practitioners, the immediate takeaway is that a single pretrained model can now compete with or exceed well-tuned task-specific methods across hundreds of tabular datasets. The fact that LimiX-2 achieves this while being four times smaller than the previous state-of-the-art foundation model suggests that the mechanism-oriented modeling approach is more parameter-efficient than the target-centric approach it replaces.

The causal discovery capability is particularly interesting. In domains like healthcare and finance, understanding why a prediction is made matters as much as the prediction itself. A model that simultaneously predicts and recovers causal structure from attention weights offers a path toward interpretable predictions grounded in the data-generating process, not just correlation.

The scaling results indicate that larger models will likely perform better, with no saturation in sight. For teams considering building on this work, the combination of clean scaling laws, competitive performance at 406M parameters, and an open-source release makes LimiX-2 a practical starting point for further development. The synthetic pretraining pipeline, which generates diverse causal structures and functional mechanisms, may also be useful as a standalone tool for data augmentation or for studying how model architecture interacts with data structure.

Read the paper on arXiv