One-sentence Summary

Researchers from Stable AI and Tsinghua University introduce $\text{LimiX-2}$LimiX-2, a $\text{Contextual Mechanism Network}$Contextual Mechanism Network pretrained with $\text{Context-Conditional Masked Modeling}$Context-Conditional Masked Modeling that shifts its organizing principle from $p(y | x, D_{\text{context}})$p(y∣x,Dcontext​) to $p(x, y | D_{\text{context}})$p(x,y∣Dcontext​), leveraging $\text{SCM-generated}$SCM-generated synthetic data for joint modeling and outperforming dataset-specific models and tabular foundation models on $\text{TabArena}$TabArena, $\text{TALENT}$TALENT, and $\text{BCCO}$BCCO while enabling causal skeleton recovery.

Key Contributions

• Introduces Contextual Mechanism Networks (CMNs), a paradigm that shifts tabular learning from target-centric prediction to modeling the joint dependency structure p(x, y | D_context), and instantiates it in LimiX-2, a Transformer-based tabular foundation model pretrained with Context-Conditional Masked Modeling (CCMM) to unify supervised prediction, missing-value imputation, and causal discovery in a single model.
• Pretrains LimiX-2 exclusively on synthetic data from an expanded structural causal model (SCM) generation engine, covering broader graph structures, functional mechanisms, and observation processes than its predecessor, enabling the model to handle diverse tabular regimes without task-specific parameter updates.
• Evaluations on TabArena, TALENT, and BCCO show that LimiX-2 outperforms existing dataset-specific models and tabular foundation models, including surpassing TabFM while using a 4 times smaller parameter size; causal skeleton recovery experiments also demonstrate that its feature attention encodes direct causal relationships, yielding better performance than dedicated causal discovery methods and tree-based feature importance baselines.

Introduction

Structured data underpins prediction and decision making in fields such as healthcare, finance, and scientific discovery. While gradient-boosted trees, deep neural networks, and automated ensembles achieve strong task-specific performance, these methods require separate training and model selection for each dataset, limiting knowledge reuse across tasks. Existing tabular foundation models, such as those based on Prior-Data Fitted Networks (PFNs), enable in-context learning but are confined to predicting a single designated target column, which restricts their ability to reason over the joint distribution of all variables and handle variant data reasoning tasks.

The authors introduce Contextual Mechanism Networks (CMNs), a new paradigm that shifts the modeling focus from a single target to the system of predictive dependencies among variables. By learning from multiple conditional prediction tasks over the same dataset, CMNs explicitly model the joint dependency structure and treat supervised prediction as a special case of broader inference. The authors instantiate this paradigm in LimiX-2, a transformer-based tabular foundation model, pretrained using Context-Conditional Masked Modeling (CCMM) on synthetic data from an expanded structural causal model (SCM) engine. Without task-specific parameter updates, a single LimiX-2 model supports classification, regression, missing-value imputation, and causal discovery. Evaluations on TabArena, TALENT, and BCCO show that LimiX-2 outperforms current tabular foundation models and dataset-specific models, and it surpasses TabFM while being four times smaller in parameter size.

Dataset

The authors construct large-scale pretraining data using a structural causal model (SCM) framework. The pipeline has five stages: hyperparameter sampling, directed acyclic graph (DAG) generation, SCM propagation, data sampling, and task adaptation. The following list details the dataset composition, sources, processing, and usage.

• Dataset composition and sources: The data is fully synthetic. It is generated by varying components at each stage of the SCM pipeline to create diverse datasets with different variable dependencies, feature distributions, and task properties. The pipeline expands the space of graph structures, functional mechanisms, and variable observation processes compared to the previous version, LimiX.

• Hyperparameter sampling: For each generated dataset, the authors sample global properties including sample size, feature dimension (split into continuous and categorical feature counts), and task type (classification or regression). They also randomly draw an evaluation position that splits each dataset into a context part and a query part. The sampling distribution for each hyperparameter is randomly chosen from a family such as normal, uniform, or beta distributions.

• DAG generation: The structural dependencies among variables are represented as DAGs composed of multiple local causal structures, called causal motifs. These motifs encode directed dependencies such as chains, confounding, and collider structures. The DAGs are built through recursive expansion of these motifs at multiple granularities, capturing macro- and micro-level dependencies. Topology-constrained graph transformations, including edge redirection, local path replacement, and node-level structural changes, are randomly applied while preserving acyclicity, leading to diverse connectivity patterns and information-propagation pathways.

• Functional mechanisms of SCM: For each DAG, root node values are sampled from distributions with random types and parameters. Remaining node values are computed by propagating functional computations along the topological order. The value of a node is a function of mapped parent values, edge functions, an aggregation function, and stochastic noise. The authors retain edge functions from the previous version, including MLPs, CNNs, and decision trees, and add linear mappings, kernel functions, piecewise functions, periodic functions, and multiplicative interactions. These basic functions can be composed to model complex relationships. For nodes with multiple parents, aggregation strategies include simple averaging, weighted aggregation, and neural aggregation.

• Feature and target sampling: A complete SCM defines joint states of all variables, but only a subset is observable in practice. The authors formulate variable sampling as a multi-attribute selection problem. Each dataset is constructed by retrieving a subset of variables as features and prediction targets, based on a assigned design regarding subgraph structure and feature redundancy. Candidate tasks are filtered through a multi-objective selection mechanism to ensure they differ in graph structure and cover prediction problems with diverse statistical properties, broadening the coverage of pretraining tasks.

• Task adaptation: The authors apply stochastic observation transformations to features and target variables. These include linear scaling, monotonic nonlinear transformations, logarithmic transformations, exponential transformations, and random compositions of multiple operators. For classification tasks, the initial continuous target is converted into a categorical target via random discretization, which randomly partitions the target value space into intervals and varies class frequencies and discretization parameters. This yields classification tasks with different numbers of classes and degrees of class imbalance. For regression tasks, targets undergo random scale transformations and adjustments to skewness and tail behavior, covering continuous prediction tasks with diverse functional relationships.

• How the data is used: The generated datasets serve as pretraining data for the model. Each dataset is split into a context part and a query part based on the sampled evaluation position. The model is trained on these diverse tasks, which are designed to cover a wide range of structural complexities and statistical characteristics.

Method

The authors leverage a cell-level design for LimiX-2, encoding each cell into a separate representation to preserve fine-grained tabular structure and support conditional reasoning across variables. For a table with $N$N rows and $F$F columns, raw cells $x_{i,j}^R$xi,jR​ are mapped into a feature representation space $\mathbf{x}_{i,j} \in \mathbb{R}^d$xi,j​∈Rd with an expanded embedding dimension $d=256$d=256. Missing cells share a single learnable embedding $E_{\text{miss}}$Emiss​, while observed cells are processed through a two-layer MLP with RMSNorm and GELU. To distinguish columns that might share similar marginal distributions, the model incorporates Discriminative Feature Encoding. Each column is assigned an $s$s-dimensional code mapped into the embedding space, providing explicit column identity without encoding sequential proximity.

Target variables are encoded into $K=4$K=4 task-embedding slots, each of dimension $d$d. Numerical regression targets use an encoder, while categorical targets use an orthogonally initialized embedding table. A task-type embedding is added to each slot.

The model backbone consists of a stack of $M=24$M=24 dual-axis transformer blocks. Unlike previous versions, LimiX-2 separates the computation paths for feature and task representations.

As shown in the figure below, the architecture processes feature and target embeddings through distinct pathways. Within each block, the computation follows a specific order:

1. Sample-axis attention: Representations are propagated across samples. For target positions, the $K$K task embeddings are concatenated into a unified vector before attention. Context rows attend to each other, while query rows only attend to context rows. The query, key, and value mappings are shared among features but distinct between features and targets.
2. Independent SwiGLU: The shared MLP is replaced by gated Feed-Forward Networks instantiated separately for feature and target representations. The feature FFN operates in $\mathbb{R}^p$Rp, while the target FFN operates on the concatenated slot space $\mathbb{R}^{Kd}$RKd.
3. Asymmetric feature-axis attention: Feature representations attend to both other features and target representations, whereas target representations only attend to feature representations. This asymmetry directs information from features to the task readout.

Multi-head attention utilizes all key and value heads, with queries and keys normalized before computing scores. Queries are rescaled per head by a length-dependent factor $s_h = (1 + w_h \log n) \beta_h$sh​=(1+wh​logn)βh​ to maintain stability across variable context lengths.

Prediction heads are attached at different depths. Masked-feature reconstruction utilizes shallow depth representations to capture local data details. Classification and regression tasks decode from the final-layer representations. For classification, the head emits logits trained with cross-entropy. For regression, the target range is partitioned into $B=5000$B=5000 ordered bins, predicting the probability of each bin to derive the regression value $\hat{y} = \sum_{i=1}^B p_i c_i$y^​=∑i=1B​pi​ci​.

The authors employ Context-Conditional Masked Modeling to capture joint dependency structures. Each pretraining episode partitions a table into disjoint context and query sets. The model estimates conditional probabilities for masked features and query targets based on observed features and context data. To broaden the coverage of observation patterns, the training combines three masking schemes: individual entries, selected columns across query rows, and blocks of entries. Masked cells are replaced by the shared missing-value embedding combined with the column identity code.

To train on diverse variable dependencies and task properties, the authors construct large-scale pretraining data using a Structural Causal Model framework.

As shown in the figure below, the synthetic data generation pipeline comprises five stages:

1. Hyperparameter Sampling: Global properties such as sample size, feature dimension, and task type are sampled from various distributions.
2. DAG Generation: Directed Acyclic Graphs are generated hierarchically using causal motifs to capture macro and micro-level dependencies, enriched by topology-constrained transformations.
3. SCM Propagation: Root node values are sampled from diverse distributions, and remaining nodes are computed by propagating functional mechanisms such as MLPs, CNNs, decision trees, and kernel functions along the graph topology.
4. Feature and Target Sampling: A multi-objective selection mechanism filters candidate tasks based on graph structure and feature redundancy to ensure diverse statistical properties.
5. Task Adaptation: Stochastic observation transformations, such as scaling and nonlinear mappings, are applied. Classification targets are created via random discretization with varying class frequencies, while regression targets undergo scale and skewness adjustments.

Experiment

LimiX-2 was evaluated on three public tabular benchmarks (TabArena, TALENT, and BCCO) covering classification, regression, and robustness scenarios, where it achieved the highest Elo ratings, lowest average ranks, and broad pairwise win rates above 50% against all baselines, including tree-based models, AutoML, neural networks, and other tabular foundation models. In a causal skeleton recovery evaluation across six datasets, feature attention from LimiX-2 yielded top F1 scores and lowest structural Hamming distances, outperforming dedicated causal discovery methods and group-level attention models, highlighting the benefit of cell-level representations. A scaling study with six model sizes (12.5M to 406.2M parameters) showed log-linear Elo improvements with a fit of R-squared between 0.96 and 0.98 across all evaluation series, with no saturation observed, supporting extrapolation toward billion-parameter models.

The table outlines three tabular benchmarks, TabArena, TALENT, and BCCO, with their dataset counts, task types, and primary metrics, used to evaluate LimiX-2. TabArena focuses on practical predictive performance, TALENT on diverse task generalization, and BCCO on robustness with incomplete data. Together, they provide a comprehensive assessment of predictive performance, generalization, scalability, and robustness. TabArena includes 51 datasets spanning binary and multiclass classification as well as regression, evaluated using Elo, rank, and win rate. TALENT contains 288 datasets after excluding 12 with more than 10 target classes, covering binary and multiclass classification plus regression. BCCO provides 156 datasets focused on robustness, with 106 classification and 50 regression tasks, addressing missing and incomplete features.

LimiX-2 achieves the top rank across all metrics on the TabArena benchmark, outperforming the runner-up TabFM+ by a substantial margin in Elo score. It also shows lower improvability, better average rank, and a much higher aggregated win count, indicating consistent superiority over all compared methods. LimiX-2 has the highest Elo score, exceeding the runner-up by more than 100 points. LimiX-2's improvability is roughly half that of TabFM+, suggesting more stable performance. LimiX-2's aggregated win count is about 3.6 times larger than TabFM+'s. TabM, iLTM, and other baselines fall far behind in Elo and win count.

LimiX-2 outperforms all baselines on TabArena classification tasks, leading in Elo, improvability, average rank, and win count. The model's improvability is notably lower than TabFM+'s, indicating more stable performance, and it achieves substantially better average rank and aggregated wins. LimiX-2 achieves the highest Elo, surpassing TabFM+ by over 100 points. LimiX-2 has a much lower improvability score (4.3%) compared to TabFM+ (7.5%), suggesting more consistent performance. LimiX-2 ranks better on average (6.0) and wins more aggregated matches (10.6) than all other models, including TabFM+ (9.9 and 4.7).

LimiX-2 achieves the top rank across all predictive metrics on the TabArena benchmark, including the highest Elo score and lowest improvability among compared methods. It also shows strong aggregate performance on classification tasks, with a high average pairwise win rate and the best average rank. These results highlight its consistent superiority over both traditional boosting models and other tabular foundation models. LimiX-2 leads in Elo with a score of 1917 on classification tasks, outperforming the runner-up by a clear margin. The model attains the lowest improvability (4.3%) and the best average rank (6.0) among all methods on classification datasets. LimiX-2 records an average pairwise win rate of 94.5%, indicating it wins or ties nearly every comparison against other models. Traditional baselines like CatBoost and LightGBM trail significantly, with Elo scores around 1370-1390 and higher improvability values.

LimiX-2 outperforms all baselines on TabArena regression tasks, achieving the highest Elo, lowest improvability, best average rank, and most wins. Its improvements are particularly large over TabFM+ and AutoGluon 1.6 (noncommercial, 4h), while also exceeding other variants and models by notable margins. LimiX-2 ranks first on all four reported metrics, with an Elo of 2206, substantially ahead of the next best methods. LimiX-2 exhibits much lower improvability (0.6%) and better average rank (3.8) compared with TabFM+ and AutoGluon 1.6 (noncommercial, 4h). LimiX-2 records 8.3 aggregated wins, far exceeding the win counts of all baselines, which are at most 0.5. Among baselines, TabFM+ and AutoGluon 1.6 (noncommercial, 4h) are the closest competitors, but both trail LimiX-2 by over 140 Elo points.

The evaluation covers three tabular benchmarks: TabArena for predictive tasks, TALENT for generalization, and BCCO for robustness with incomplete data. Across all settings, LimiX-2 consistently outperforms baselines, leading in Elo, win rate, average rank, and stability (lower improvability), with particularly strong margins on regression tasks and a high pairwise win rate on classification. The results confirm LimiX-2's superiority over both traditional boosting models and other tabular foundation models.