PaperlayerView exact revision on arXiv

Source-linked AI summary

SOMTab: Set-Order Mamba for Efficient Tabular In-Context Learning

Hao Wang, Siyu Zhang, Wei Ma

arXiv:2608.27882v1cs.LGcs.AI

TL;DR

Attention-heavy tabular in-context learners are effective but costly, raising whether attention is needed at every modeling stage. SOMTab separates representation construction from query-conditioned retrieval using Set-to-Order Mamba and retains attention for final prediction, while DCH-TailMix diversifies pretraining dependencies. Across benchmarks, it approaches strong Transformer-based tabular foundation models with faster inference and lower GPU memory usage.

  • Problem

    Leading PFN-style tabular foundation models rely heavily on attention, whose quadratic sample- or feature-dimension cost makes long-context prediction increasingly expensive.

  • Method

    SOMTab maps unordered table tokens into stable ordered latent slots for Mamba-based representation construction, retains attention for query-context matching, and adds the DCH-TailMix synthetic prior.

  • Results

    SOMTab approaches the predictive performance of strong Transformer-based tabular foundation models while substantially reducing inference cost and exhibiting favorable memory scaling.

  • Takeaways & Limitations

    Tabular ICL can selectively use attention for query-context matching and efficient state-space mixing for representation construction after unordered structure is organized.

  • Takeaways & Limitations

    TabArena runtime comparisons provide contextual positioning rather than a hardware-matched speed ranking because baselines and SOMTab were measured on different GPUs.

Abstract

from arXiv · show

Tabular foundation models based on in-context learning have recently emerged as strong alternatives to task-specific model fitting. However, the current performance frontier remains dominated by attention-heavy architectures, where attention is used throughout the modeling pipeline. This raises a natural question: is attention necessary at every stage of tabular in-context learning? We introduce SOMTab, a Set-Order Mamba architecture for efficient tabular in-context learning. SOMTab separates representation construction from query-conditioned retrieval. For row and column representations, it maps unordered table tokens into stable latent slots and applies Mamba-based state-space mixing to construct compact representations. For final prediction, it retains attention-based in-context learning to preserve query-conditioned retrieval from labeled context examples. We further introduce DCH-TailMix, a synthetic prior that combines degree-corrected graph heterogeneity with mixed heavy-tailed regimes to diversify synthetic dependency structures. Across tabular benchmarks, SOMTab approaches the performance of strong Transformer-based tabular foundation models while achieving faster inference and lower GPU memory usage, yielding a favorable efficiency--accuracy trade-off.

1 Introduction

SOMTab addresses the cost of attention-heavy tabular in-context learning by separating representation construction from query-context matching. It combines Set-to-Order Mamba modules and the DCH-TailMix prior with attention-based final prediction, achieving competitive accuracy with lower computational cost.

  • Motivation: Attention-heavy tabular foundation models support expressive in-context interactions but incur quadratic costs along sample or feature dimensions as contexts grow.This motivates selective rather than universal use of attention in the pipeline.
  • SOMTab architecture: SOMTab organizes unordered table tokens into ordered latent slots and applies Mamba mixing to construct efficient row and column representations.Attention is retained for query-conditioned prediction from labeled context examples.
  • Results: SOMTab achieves competitive performance among strong Transformer-based tabular foundation models while reducing fit-plus-predict runtime and improving peak-memory scaling as context grows.Additional TabArena results show a competitive efficiency–performance trade-off, while ablations support contributions from both the Mamba encoder and DCH-TailMix.
  • Design principle: The resulting design principle is to use attention for query–context matching and efficient state-space operators for representation construction when unordered tables can be organized into stable latent order.This allocates attention to the stage where query-dependent routing is most directly useful.
  • Synthetic prior: DCH-TailMix combines degree-corrected graph heterogeneity with mixed heavy-tailed regimes to expose pretraining to more diverse dependency structures.The prior targets broader feature-target and feature-feature relationships than a single graph-sampling regime.

2 Related Work

Prior work establishes tabular foundation models through PFN-style synthetic-task pretraining, while attention-based architectures remain effective but costly. Efficient state-space and set-based alternatives address computation or permutation structure, yet direct Mamba use remains nontrivial because tabular rows and columns lack universal order.

  • Tabular foundation models: PFN-style tabular foundation models pretrain predictors on task distributions and condition inference on labeled training examples for unseen tables.TabPFN and later variants extend this paradigm across accuracy, task coverage, and heterogeneous tabular settings.
  • Efficient architectures: Attention-based architectures enable flexible interactions within and across tabular contexts but become computationally costly when used throughout the pipeline.This motivates hybrid architectures that balance modeling quality and inference efficiency.
  • Efficient architectures: Mamba is difficult to apply directly because rows are exchangeable examples and columns lack a universal ordering across datasets.Set-based architectures handle unordered inputs through permutation-aware attention, but still retain attention as their main mechanism.
  • Synthetic priors: Synthetic priors strongly shape PFN model inductive bias and performance, evolving from causal and Bayesian mechanisms toward richer feature, missingness, graph, and function families.These developments aim to generate more diverse supervised tabular tasks during pretraining.

3 Method

SOMTab separates tabular in-context learning into representation construction and query-conditioned prediction, using Set-to-Order Mamba for efficient encoding while retaining attention for final retrieval. Its architecture handles unordered table structure through latent slots, and DCH-TailMix broadens synthetic dependency patterns with heterogeneous graphs and mixed tail regimes.

  • Attention Allocation in Tabular ICL: SOMTab decomposes tabular ICL into representation construction and query-conditioned prediction, assigning different operators to each role.Representation construction captures column distributions and feature interactions; final prediction retrieves information from labeled context examples.
  • Set-to-Order Mamba Block: SOMTab maps unordered table tokens into fixed-index latent sequences before Mamba mixing, avoiding dependence on arbitrary row or column permutations.Seed slots define the latent order, while the set-to-order mapping is permutation-invariant and optional readback is permutation-equivariant.
  • SOMTab Architecture: The Set-to-Order Mamba block operates across samples within each feature and across features within each sample to produce row representations for in-context prediction.Feature-level summaries are read back to context and query tokens, then updated feature tokens are combined with shared learned row tokens and mixed into sample-level latent slots.
  • SOMTab Architecture: SOMTab injects label information into context rows during representation construction and later adds labeled context pairs to the attention-based reader.The early embedding conditions table representations, while the later embedding supplies labeled context explicitly for final prediction.
  • Complexity: For fixed latent-slot counts, representation construction costs O(Nm(Kv + Ks)), while the final row-level attention reader remains bounded by O(N^2).The representation modules are linear in observed cells; the final ICL reader retains row-level attention.
  • DCH-TailMix Synthetic Prior: DCH-TailMix generates classification tasks from DAGs using degree-corrected graph heterogeneity and mixed heavy-tailed regimes.Its graph variation can produce sparse, modular, dense, correlated, and hub-dominated structures, while TailMix varies task-level and node-level tail behavior.

4 Experiments

SOMTab is evaluated for predictive quality, runtime, memory scaling, and ablation effects across TALENT and TabArena settings. Results show competitive performance with improved efficiency, especially as context size grows, while both the Set-Order Mamba encoder and DCH-TailMix contribute to performance.

  • Predictive Quality–Efficiency Trade-off: SOMTab requires substantially less runtime than both default and single-estimator TabICLv2 configurations while maintaining competitive predictive performance.The matched H200 evaluation isolates runtime differences across ensemble settings.
  • Predictive Quality–Efficiency Trade-off: SOMTab achieves competitive classification performance while maintaining one of the lowest runtimes among high-performing methods.Table 1 aggregates normalized log-loss, accuracy, Macro F1, and median fit-plus-predict runtime across TALENT classification datasets.
  • Predictive Quality–Efficiency Trade-off: TabArena provides contextual efficiency–performance positioning rather than a hardware-matched speed ranking.Public baselines retain H100 configurations, whereas SOMTab is measured on A100 and H200 GPUs.
  • Runtime and Memory Scaling: SOMTab exhibits the lowest runtime across evaluated context lengths, with an increasing advantage as the labeled context grows.Figure 3 varies context rows from 512 to 28,672 with 1,024 queries, 32 features, and 10 classes.
  • Runtime and Memory Scaling: Attention-based architectures show quadratic memory growth that becomes dominant at larger context sizes, while SOMTab improves long-context efficiency.SOMTab can be slightly higher in memory at small contexts because of latent representation modules.
  • Ablation Studies: DCH-TailMix achieves the highest mean accuracy among evaluated prior variants under the same pretraining budget.The variants share the SOMTab architecture and differ only in their synthetic priors.
  • Ablation Studies: Replacing the final attention-based ICL reader with a Mamba-based reader lowers accuracy under the same training configuration.This ablation retains DCH-TailMix and the Set-Order Mamba representation encoders.

5 Conclusion

SOMTab separates representation construction from query-conditioned retrieval, using Set-Order Mamba for efficient representations while retaining attention for final prediction. DCH-TailMix diversifies synthetic dependency structures, and experiments show competitive performance with lower computational cost.

  • SOMTab separates table representation construction from query-conditioned retrieval in tabular in-context learning.
  • Unordered table tokens are mapped into stable latent sequences and mixed with Mamba, while attention remains for final ICL prediction.
  • DCH-TailMix improves the diversity of dependency structures used during synthetic pretraining.
  • SOMTab approaches strong Transformer-based tabular foundation models in predictive performance while substantially reducing inference cost.
  • The results suggest assigning different sequence-modeling mechanisms to different stages of tabular in-context learning.

A Architecture and Implementation Details

The architecture converts unordered table tokens into fixed-index latent sequences before Mamba mixing, then uses attention-based context retrieval for prediction. Its synthetic prior generates heterogeneous graph structures and heavy-tailed regimes for pretraining.

  • Set-to-Order Mamba Block: SOMTab maps unordered table tokens into fixed-index latent sequences before applying Mamba-based mixing.This avoids representations depending on arbitrary row or feature permutations.
  • Set-to-Order Mamba Block: Cross-attention aggregates input sets into latent slots, Mamba-2 mixes the ordered slots, and readback restores token-aligned outputs.
  • Sample-Level Encoder: The sample-level encoder adds four shared row-summary tokens and processes four ordered latent slots through three Mamba-2 residual blocks.
  • Context-Only ICL Reader: The reader lets all rows query labeled context representations, then maps query representations to class logits after 12 reader blocks.
  • DCH-TailMix Task Generation: DCH-TailMix samples graph families, tail regimes, latent variables, observed features, and targets to produce filtered synthetic classification tasks.
  • DCH Graph Prior: Its graph prior combines modular, low-rank, degree-heterogeneous, and noisy dependency components to generate sparse, dense, correlated, and hub-dominated structures.

B.5 Observation and Classification Adapters

The observation and classification adapters convert latent variables into numerical or categorical features and discretize generated targets into classes. Synthetic tasks are filtered for valid, nondegenerate feature–target relationships before training.

  • Observation Model: Selected latent features are independently converted to categorical form with task-level probability qcat, while numerical features use one-dimensional latent projections.
  • Observation Model: Categorical converters vary cardinality, distance metrics, projections, biases, and separation parameters to diversify observed feature representations.
  • Classification Adapter: One to three target nodes generate a standardized continuous target that is discretized into class labels using observed-value or Gaussian thresholds.
  • Graph-Family Mixture: DCH-TailMix includes four graph families: controlled heavy-tail, modular heavy-tail, low-rank confounder, and dense weak-dependency graphs.
  • Task Validity: Tasks are rejected for nonfinite values, insufficient feature variance, weak feature–target marginal correlation, or incompatible context–query class coverage.
  • Pretraining: The model is trained entirely on DCH-TailMix tasks using a two-stage curriculum spanning task sizes and context fractions.

C.2 Optimization and Compute

Training uses distributed hybrid Muon–AdamW optimization, while evaluation measures predictive metrics, runtime, and peak GPU memory under specified dataset and context-scaling protocols. Timing excludes data loading and other setup costs.

  • Optimization: Training combines Muon for matrix parameters with AdamW for embedding and lower-dimensional parameters, using gradient clipping and zero weight decay.
  • Compute: Training runs on four 40 GB NVIDIA A100 GPUs, with Stage 1 taking approximately 17 days and Stage 2 approximately six days.
  • Evaluation: TALENT evaluation reports Accuracy, Macro F1, and log-loss on test rows, assigning equal weight to each dataset in aggregate results.
  • Timing: Runtime includes fitting when applicable and prediction but excludes dataset loading, external preprocessing, model construction, and checkpoint loading.
  • Timing and Memory: Runtime is normalized per 1000 labeled context examples, with aggregate runtime reported as the median across datasets; peak memory is maximum allocated CUDA memory.
  • Context Scaling: The context-scaling experiment fixes nte = 1024, m = 32, and C = 10 while varying ntr from 512 to 28,672.

D.3 Controlled Early-Training Ablation

The controlled early-training ablation compares synthetic priors and final ICL readers under matched training conditions, using a fixed subset of 15 TALENT datasets. DCH-TailMix improves mean accuracy, while replacing the attention reader with Mamba reduces accuracy.

  • Evaluation design: All variants are randomly initialized, trained for 20K steps, and evaluated every 5K steps on the same 15-dataset subset.The subset contains TALENT datasets where SOMTab is relatively less favorable than TabICLv2 in the main benchmark.
  • Synthetic priors: DCH-TailMix achieves higher mean accuracy than the Cauchy and SCM–Tree priors at every evaluated checkpoint.The comparison uses the same architecture and training budget.
  • Final ICL readers: The Mamba Reader remains less accurate than the attention-based final reader throughout training.This supports using Mamba for representation construction and attention for final in-context retrieval.

E.1 Per-Dataset Accuracy Results

Table 7 reports per-dataset accuracy for 20 TALENT classification datasets, with the final 15 datasets forming the fixed subset for the controlled early-training ablation.

  • Per-dataset results: Per-dataset accuracy is reported for 20 TALENT classification datasets.The final 15 datasets are reserved as the fixed evaluation subset for the controlled early-training ablation.
  • Ablation subset: The final 15 datasets constitute the fixed evaluation subset used in the controlled early-training ablation.

F Additional Evaluation on TabArena

Additional TabArena evaluation compares models using average performance gaps to the dataset-best method and reports runtime as a reference comparison across GPU environments.

  • Performance metric: TabArena averages each model’s per-dataset performance gap to the best-performing method, with smaller gaps indicating closer performance.
  • Efficiency comparison: Runtime comparisons are not strictly hardware-controlled because public results use H100 GPUs, while SOMTab is additionally evaluated on H200 and A100 GPUs.The H200 and A100 measurements provide a broader indication of practical efficiency across GPU generations.

F.1 Evaluation Protocol

The TabArena evaluation follows the public protocol, comparing performance–efficiency trade-offs and pairwise model wins across datasets. Pairwise win rates measure how often the row model outperforms the column model.

  • Evaluation protocol: The evaluation follows the public TabArena protocol and includes the models provided by the benchmark.
  • Figures: Figure 9 compares performance–efficiency trade-offs, while Figure 10 reports pairwise comparisons between models.
  • Pairwise comparison: Pairwise win rates report the fraction of datasets where the row model achieves higher predictive performance than the column model.This compares individual-dataset outcomes rather than only averaged benchmark scores.
  • Related benchmark context: Table 7 reports accuracy on 20 TALENT classification datasets, with the final 15 rows forming the controlled-ablation subset.
  • Performance–efficiency plot: Figure 9 places inference time per 1000 samples on a logarithmic x-axis and average performance gap to the dataset-best method on the y-axis.Lower average performance gaps indicate performance closer to the best method.
Loading 2608.27882v1…