PaperlayerView exact revision on arXiv

Source-linked AI summary

TabICL: A Tabular Foundation Model for In-Context Learning on Large Data

Jingang Qu, David Holzmüller, Gaël Varoquaux, Marine Le Morvan

arXiv:2502.05564v2cs.LGcs.AI

TL;DR

Existing tabular foundation models use in-context learning but face computational limits on large tables. TabICL introduces a two-stage embedding architecture and larger-scale pretraining to extend this approach. Across 200 classification datasets, it matches TabPFNv2 while being faster and surpasses TabPFNv2 and CatBoost on larger datasets.

  • Problem

    TabPFNv2’s alternating column- and row-wise attention makes in-context learning computationally difficult for large training sets.

  • Method

    TabICL builds fixed-dimensional row embeddings with column-then-row attention, followed by a transformer for in-context classification.

  • Results

    Across 200 classification datasets, TabICL is comparable to TabPFNv2 while being faster; on 53 datasets with over 10K samples, it surpasses TabPFNv2 and CatBoost.

  • Takeaways & Limitations

    TabICL extends tabular in-context learning to datasets with up to 500K samples and provides a scalable alternative for large-data classification.

  • Takeaways & Limitations

    TabICL is limited to classification, violates column-permutation invariance because of RoPE, and inherits limitations of the TALENT benchmark.

Abstract

from arXiv · show

The long-standing dominance of gradient-boosted decision trees on tabular data is currently challenged by tabular foundation models using In-Context Learning (ICL): setting the training data as context for the test data and predicting in a single forward pass without parameter updates. While TabPFNv2 foundation model excels on tables with up to 10K samples, its alternating column- and row-wise attentions make handling large training sets computationally prohibitive. So, can ICL be effectively scaled and deliver a benefit for larger tables? We introduce TabICL, a tabular foundation model for classification, pretrained on synthetic datasets with up to 60K samples and capable of handling 500K samples on affordable resources. This is enabled by a novel two-stage architecture: a column-then-row attention mechanism to build fixed-dimensional embeddings of rows, followed by a transformer for efficient ICL. Across 200 classification datasets from the TALENT benchmark, TabICL is on par with TabPFNv2 while being systematically faster (up to 10 times), and significantly outperforms all other approaches. On 53 datasets with over 10K samples, TabICL surpasses both TabPFNv2 and CatBoost, demonstrating the potential of ICL for large data. Pretraining code, inference code, and pre-trained models are available at https://github.com/soda-inria/tabicl.

1. Introduction

Tabular foundation models use in-context learning to avoid parameter updates, but existing attention mechanisms limit scalability to large datasets. TabICL addresses this with efficient row embeddings and supports substantially larger classification tables.

  • Tabular classification supports applications across industries including healthcare and finance.
  • In-context learning enables models to capture task patterns without parameter updates, motivating tabular foundation models.
  • TabPFNv2 performs strongly on datasets with up to 10K samples and 500 features, but alternating column- and row-wise attention limits scalability.
  • TabICL uses distribution-aware column embeddings followed by attention-based row interactions to construct semantically grounded row representations.
  • TabICL refines pretraining with tree-based synthetic data, curriculum learning from 1K to 60K samples, and hierarchical classification for datasets exceeding 10 classes.
  • TabICL handles up to 500K samples and 500 features, runs in a single forward pass, and outperforms or matches established methods across the reported benchmark settings.It is reported as up to 10 times faster than TabPFNv2 and surpasses TabPFNv2 and CatBoost on 53 datasets with over 10K samples.

2. Related Work

Related work establishes in-context learning and tabular foundation models as alternatives to conventional tabular methods. TabICL builds on this line by improving scalability and efficiency relative to TabPFN variants.

  • 2.1. Foundation Models and In-Context Learning: TabICL’s architecture combines column-wise embedding, row-wise interaction, and dataset-wise ICL through three transformers.
  • 2.1. Foundation Models and In-Context Learning: In-context learning performs tasks from input-output examples in prompts without explicit training or parameter updates.
  • 2.2. Deep Learning for Tabular Data: Tabular deep learning research increasingly emphasizes cross-table transferability through transferable backbones and dataset-specific components.
  • 2.3. Tabular Foundation Models: TabPFN is a transformer pretrained on synthetic tabular datasets for classification through in-context learning.
  • 2.3. Tabular Foundation Models: TabPFN variants commonly reduce in-context samples through prompt engineering rather than structurally improving the underlying model.
  • 2.3. Tabular Foundation Models: TabICL achieves comparable performance to TabPFNv2 while being more scalable and computationally efficient.

3. The TabICL Architecture

TabICL first builds distribution-aware feature embeddings and context-aware row representations, then applies a final transformer for in-context classification. This column-then-row design collapses feature dimensions before ICL, improving scalability while RoPE mitigates representation collapse from similarly distributed features.

  • Architecture: TabICL uses a tabular embedding module followed by an ICL module, with labels used exclusively in the ICL stage.The embedding module models column statistics and cross-column dependencies before classification.
  • Distribution-aware Column-wise Embedding: A shared Set Transformer maps each column’s cell values to distribution-aware weights, biases, and d-dimensional cell embeddings.The set-input formulation enables parameter sharing across columns and uses permutation-invariant cell sets.
  • Context-aware Row-wise Interaction: Learned column embeddings separate features by distributional properties such as skewness and kurtosis.The visualization projects embeddings for 40,000 synthetic features onto their first two principal components.
  • Context-aware Row-wise Interaction: A 3-layer row-wise transformer models feature interactions and concatenates four learnable [CLS] outputs into 512-dimensional row embeddings.The resulting fixed-dimensional representations accommodate varying numbers of features before final ICL.
  • Context-aware Row-wise Interaction: Vanilla permutation-invariant row attention can cause representation collapse when features share the same distribution, as in balance scale.RoPE breaks this symmetry, preserves distinct sample representations, and is followed by column-order ensembling to approximately restore permutation invariance.
  • ICL Module: The final ICL transformer lets training embeddings attend to one another while test embeddings attend only to training embeddings before an MLP predicts class probabilities.Its 12-layer transformer operates on compressed row embeddings rather than the original table cells.
  • Computational Complexity: TabICL has overall complexity O(m2n + n2), compared with TabPFNv2’s O(m2n+n2m) because it collapses feature dimensions before final self-attention.This difference makes TabPFNv2 increasingly expensive for large sample counts and moderate feature counts.

4. Pretraining and Inference

TabICL is pretrained on diverse synthetic datasets generated with structural causal and tree-based models, using staged growth in sample size. Inference additionally supports many-class problems through hierarchical classification and adapts batch size to available GPU memory.

  • Improved Pretraining Synthetic Datasets: TabICL generates synthetic pretraining data with structural causal models to represent realistic dependencies among variables.The process samples a DAG and models each feature from its parent variables with independent noise.
  • Improved Pretraining Synthetic Datasets: Tree-based structural causal models comprise 30% of generation and use XGBoost to model complex interactions and hierarchical dependencies.The remaining 70% uses standard SCMs, balancing the two generation approaches.
  • Improved Pretraining Synthetic Datasets: Pretraining diversifies synthetic dependencies by adding 15 activation functions, Gaussian-process activations, layer-specific activations, and random rescaling.These additions extend the original TabPFN activation set with non-monotone and discontinuous functions.
  • Pretraining Schedule: Synthetic dataset sizes increase through three stages from 1,024 samples to uniformly sampled sizes between 40K and 60K.The procedure adjusts micro-batch size, uses activation checkpointing for datasets exceeding 10K samples, and freezes non-ICL components in stage three.
  • Pretraining Schedule: Pretraining uses 512 datasets per step with at most 100 features and 10 classes, taking 20 days on three 40GB A100 GPUs.FlashAttention and automatic mixed precision are applied globally.
  • Hierarchical Classification: Many-class classification recursively partitions classes into groups of at most 10, with class probabilities obtained by multiplying probabilities along each root-to-leaf path.A k-class problem uses a hierarchy of depth r = ⌈log10 k⌉.
  • Hierarchical Classification: Hierarchical subtasks share the same learned row embeddings and TFicl, improving efficiency during hierarchical inference.Labels are introduced only in the final ICL block, so the hierarchy is constructed during dataset-wise ICL.
  • Inference Efficiency: Inference fits peak activation memory with a polynomial model and dynamically adjusts batch size to sequence length and available GPU memory.Batch size represents different quantities depending on whether processing is column-wise or another inference context.

5. Experiments

Across the TALENT benchmark, TabICL achieves state-of-the-art accuracy while substantially reducing runtime, and its performance remains strong as dataset size and class count increase. Ablations show benefits from tree-based synthetic data and curriculum learning, with some small-dataset trade-offs.

  • Results: TabICL achieves the best median relative accuracy across datasets, while TabICL and TabPFNv2 outperform competitors by a wide margin.The difference between TabICL and TabPFNv2 is not statistically significant.
  • Results: 1.1 seconds per 1K samples is TabICL’s geometric mean training+inference time, versus around 3 minutes for tuned CatBoost and around 7 minutes for RealMLP and ModernNCA.The reported TabICL time includes training and inference.
  • Speedup: 3–10× faster on large datasets, TabICL outpaces TabPFNv2; on 10,000 samples and 100 features, runtimes are roughly 20 seconds versus 1 minute 40 seconds.On 1,000 samples and 10 features, the corresponding times are 1 second and 2 seconds.
  • Scalability: TabICL’s performance remains strong for larger datasets, whereas TabPFNv2 can fail above 30K samples because of memory usage.TabPFNv2 achieves excellent performance up to 10K samples but was pretrained with up to 2048 training samples.
  • Probability quality: TabICL and TabPFNv2 significantly outperform accuracy-tuned competitors on log loss, with no significant difference between the two models.This indicates comparable probability-estimation performance for the two in-context models.
  • Many classes: TabICL ranks second in mean normalized accuracy on datasets with more than 10 classes, while TabPFNv2 cannot natively handle more than 10 classes.TabICL uses hierarchical classification for these datasets.
  • Ablation Studies: Tree-based synthetic data-generating models improve TabICL’s performance across the 200-dataset evaluation.The comparison evaluates pretraining with versus without tree-based SCMs.
  • Ablation Studies: Curriculum learning improves TabICL’s average rank from 11.4 to 7.46 and finally 6.95 across three stages, but slightly reduces performance on some small datasets.The final average rank is first place.

6. Conclusion

TabICL extends tabular foundation-model ICL to substantially larger datasets while remaining competitive with TabPFNv2 and faster than other tabular methods. Its results support the view that pretraining can provide a competitive advantage even with large data.

  • TabICL is evaluated on datasets with up to 100K training samples and delivers excellent performance without hyperparameter tuning.
  • TabICL achieves comparable performance to TabPFNv2 while being more scalable and faster.
  • Pretraining induces implicit priors that give in-context transformers a competitive advantage even with large data.

A.1. Architecture

TabICL reduces the cost of tabular ICL by collapsing table representations before a transformer stage, while expanding pretraining and inference support beyond TabPFNv2's scale.

  • Core mechanism: TabICL uses column-then-row attention to build fixed-dimensional embeddings, whereas TabPFNv2 alternates attention across the original table dimensions.
  • Representation collapse: TabICL uses RoPE to break symmetry between features with similar distributions, while TabPFNv2 uses random feature identifier vectors.
  • Label fusion: TabICL performs late label fusion, using labels only in the final ICL stage, whereas TabPFNv2 fuses labels early with input features.
  • Prior generation: TabICL extends synthetic pretraining with broader nonlinear SCMs and tree-based SCMs designed to incorporate tree-model inductive biases.
  • Dataset size: TabICL scales pretraining sample counts from 1,024 to 60K through curriculum learning, while TabPFNv2 uses datasets up to 2,048 samples.
  • Dataset size handling: Within 5GB of GPU memory, TabICL handles 100K samples and 500 features, while TabPFNv2 targets up to 10K samples and 500 features.
  • Computational complexity: For n rows and m columns, TabICL has complexity O(m^2n + n^2), compared with TabPFNv2's O(m^2n + n^2m).

B. Further Experiments

Further experiments examine runtime scaling and performance across dataset characteristics. TabICL is faster at scale, behaves well with many features, and shows dataset-dependent behavior across classes and categorical variables.

  • Runtime: Runtime analysis models forward-pass time using the attention complexities of TabPFNv2 and TabICL as functions of rows and columns.
  • Runtime: TabICL's estimated speedup over TabPFNv2 approaches 5 for large datasets and is 1.4 for small datasets.
  • Runtime: The runtime comparison includes only datasets with fewer than 30,000 samples because TabPFNv2 uses subsampling on larger datasets to avoid RAM overflow.
  • Number of classes: TabICL's performance deteriorates on three-class datasets, although the analysis cannot determine whether class count causes this pattern.
  • Number of features: TabICL behaves well with large numbers of features despite tokenizing entire rows before seeing labels.
  • Categorical variables: Performance generally deteriorates with many categorical variables, but TabICL performs slightly better than TabPFNv2 on such datasets.

C. Synthetic Datasets for Pretraining

The pretraining pipeline broadens the synthetic prior, uses randomized feature and activation transformations, and progressively increases dataset size through curriculum learning. These choices are accompanied by diagnostics of generated datasets, positional encoding, and stage-wise ranking.

  • SCM prior: The SCM prior adds standardization, random rescaling, and randomized activation functions alongside the original activation choices.
  • SCM prior: The added activation functions include SELU, SiLU, and a random Fourier-feature-inspired function.
  • SCM prior: The random function is sampled ten times more often than the other activation functions because it can represent many different functions.
  • SCM prior: Figures C.1 and C.2 visualize the activation functions and randomly generated 2D datasets from the resulting prior.
  • Tree-based SCM prior: The tree-based SCM prior replaces linear and activation layers with XGBoost models fitted on random data.
  • Tree-based SCM prior: The tree-based prior samples estimator counts and maximum depths independently, each capped at 4 after exponential sampling.
  • Tree-based SCM prior: Figure C.3 shows datasets generated from the tree-based SCM prior.
  • Rotary Positional Embedding: RoPE applies position-dependent rotations to query and key vectors, encoding relative positional information in attention.

E.2. Memory-efficient inference

TabICL estimates transformer memory to choose batch sizes, processes large datasets in stages with CPU or disk offloading, and reduces resource demands through mixed precision. The approach handles datasets up to 500K samples and 500 features with bounded GPU memory, although CPU usage can remain substantial.

  • Memory estimation: A polynomial regression estimates peak memory for TFcol, TFrow, and TFicl as functions of batch size and sequence length.The estimated memory is measured in megabytes.
  • Staged processing: Input data is reshaped and passed sequentially through column-wise embedding, row-wise interaction, and dataset-wise ICL prediction.Different transformers require appropriate batch sizes for efficient processing.
  • Memory estimation: Memory tracking on an A100 GPU with 40GB was used to fit the regression parameters for adaptive batch-size selection.The tracked variables were batch size and sequence length.
  • Large-dataset inference: 100K samples and 500 features require only 5GB of GPU memory and 25GB of CPU memory with intermediate activation offloading.The setup uses an 80% training split and 20% test split.
  • Large-dataset inference: 500K samples and 500 features require less than 14GB of GPU memory, while approximately 120GB of CPU memory can be reduced through disk offloading.Memory mapping is the stated mechanism for optional disk offloading.
  • Resource reduction: Automatic mixed precision substantially reduces both memory consumption and computation time.GPU fluctuations reflect dynamic batching, while TFcol outputs are progressively offloaded to CPU.

F. Excluded Development Datasets of TabPFNv2

The evaluation excludes TabPFNv2 development datasets from the main comparison because they were used for tuning and model selection. Performance on those datasets is reported separately in a supplementary table.

  • Dataset usage: TabPFNv2 used development datasets for hyperparameter tuning, early stopping, and selecting among multiple candidate models.Examples include learning rate and prior-related parameters.
  • Evaluation protocol: The main paper excludes these development datasets to provide a fair comparison.The reported comparison instead uses high-performing methods on development datasets included in the TALENT benchmark.
  • Evaluation protocol: Table F.1 reports performance on 15 TabPFNv2 development datasets.

G. Random Forest Extension of TabPFNv2 for Large Datasets

The random forest extension reduces the samples processed by tabular foundation models and uses adaptive pruning to select parent or leaf predictions. Applied to TabPFNv2 and TabICL on large datasets, it improves both models, while TabPFNv2 gains more from ensembling.

  • Extension design: The random forest extension recursively partitions feature space using standard decision-tree splitting criteria.At inference, TabPFNv2 is fitted to the training subset reaching each leaf instead of using simple leaf majority voting.
  • Extension design: Predictions route each test sample through all trees and aggregate TabPFNv2 outputs by averaging probabilities or logits.
  • Extension design: Adaptive pruning chooses parent-node or leaf-node predictions using validation performance, preventing overfitting and optimizing tree depth by feature-space region.
  • Large-dataset evaluation: The extension significantly improves both TabPFNv2 and TabICL on datasets with more than 10K samples.Both models use adaptive pruning and an ensemble size of 4 in this evaluation.
  • Ensembling: TabPFNv2 benefits more from ensembling than TabICL across the analyzed datasets.The paper attributes this difference to TabPFNv2’s more sophisticated preprocessing and stronger sensitivity to column permutations.
  • Evaluation protocol: Critical-difference diagrams compare average method ranks using accuracy, AUC, and Log Loss, with lower rank indicating better performance.Wilcoxon-Holm tests use a significance level of 0.05; some methods appear only in accuracy-based rankings.
  • Large-dataset evaluation: Figure I.5 reports accuracy-based critical differences for 53 large classification datasets with more than 10K samples.
Loading 2502.05564v2…