PaperlayerView exact revision on arXiv

Source-linked AI summary

Orion-MSP: Multi-Scale Sparse Attention for Tabular In-Context Learning

Mohamed Bouadi, Pratinav Seth, Aditya Tanna, Vinay Kumar Sankarapu

arXiv:2511.02818v3cs.AIcs.LG

TL;DR

Tabular ICL models face limitations in representing hierarchical feature interactions, scaling attention to wide tables, and refining representations across sequential components. Orion-MSP combines multi-scale processing, sparse attention, and cross-component memory; across diverse benchmarks it reports state-of-the-art zero-shot performance, including scalability to wide tables.

  • Problem

    Existing tabular ICL architectures use single-scale feature processing, dense attention with quadratic feature-count scaling, and sequential components that limit iterative refinement and bidirectional information flow.

    Source: §1 · p. 2

  • Method

    Orion-MSP combines multi-scale feature processing, block-sparse attention, and Perceiver-style memory for bidirectional information flow while preserving ICL safety.

    Source: §1 · p. 2

  • Results

    Orion-MSP consistently achieves state-of-the-art zero-shot performance across diverse benchmarks, and scales to tables with more than 100 features where existing methods fail from memory constraints.

    Source: §1 · p. 2, §5 · p. 19

  • Takeaways & Limitations

    The analyses report that the combined architecture supports robust generalization, efficient computation, and representation of complex interdependencies while maintaining practical scalability.

    Source: §5 · p. 19

    Orion-MSP’s added architectural complexity offers limited benefit on simpler, low-dimensional datasets.

    Source: Deep Analysis and Interpretation · p. 19

Abstract

from arXiv · show

Tabular data remain the predominant format for real-world applications. Yet, developing effective neural models for tabular data remains challenging due to heterogeneous feature types and complex interactions occurring at multiple scales. Recent advances in tabular in-context learning (ICL), such as TabPFN and TabICL, have achieved state-of-the-art performance comparable to gradient-boosted trees (GBTs) without task-specific fine-tuning. However, current architectures exhibit key limitations: (1) single-scale feature processing that overlooks hierarchical dependencies, (2) dense attention with quadratic scaling in table width, and (3) strictly sequential component processing that prevents iterative representation refinement and cross-component communication. To address these challenges, we introduce Orion-MSP, a tabular ICL architecture featuring three key innovations: (1) multi-scale processing to capture hierarchical feature interactions; (2) block-sparse attention combining windowed, global, and random patterns for scalable efficiency and long-range connectivity; and (3) a Perceiver-style memory enabling safe bidirectional information flow across components. Across diverse benchmarks, Orion-MSP matches or surpasses state-of-the-art performance while scaling effectively to high-dimensional tables, establishing a new standard for efficient tabular in-context learning. The model is publicly available at https://github.com/Lexsi-Labs/Orion-MSP .

1 Introduction

Orion-MSP extends tabular in-context learning with multi-scale feature processing, sparse attention, and cross-component memory. The introduction frames these changes against limitations in existing models and reports competitive accuracy with improved scalability to wide tables.

1 Introduction

  • Tabular ICL draws on foundation-model in-context learning, while gradient-boosted trees remain state of the art for tabular prediction.The introduction links table prevalence and LLM in-context learning to interest in tabular foundation models.Source: §1 · p. 1
  • Existing tabular ICL models process features at one scale, missing hierarchical interactions among individual features, semantic groups, and higher-level data blocks.The introduction gives age and income, demographic attributes, and personal versus behavioral attributes as examples.Source: §1 · p. 2
  • Dense attention scales as O(m^2) with feature count, making memory and computation prohibitive for high-dimensional tables.For m = 100 features, the passage reports 10,000 attention operations per layer.Source: §1 · p. 2
  • TabICL’s sequential column-embedding, row-interaction, and ICL-prediction pipeline prevents downstream dataset insights from refining upstream representations.The pipeline is described as limiting bidirectional information flow and iterative refinement.Source: §1 · p. 2
  • Orion-MSP combines multi-scale processing, block-sparse attention, and Perceiver-style memory to address these limitations while maintaining ICL safety.Its scales cover individual features and groups of 4 and 16; sparse patterns combine windowed, global, and random connections, and memory enables bidirectional flow without test data influencing training representations.Source: §1 · p. 2
  • Orion-MSP achieves competitive accuracy with state-of-the-art tabular ICL methods and scales to tables with more than 100 features where existing methods fail from memory constraints.The model combines hierarchical multi-scale processing, structured sparsity, and cross-component memory.Source: §1 · p. 2

2 Related Work

Prior tabular ICL methods offer different routes to prediction and scalability, but retain gaps in feature scaling, attention cost, or cross-component feedback. The related work motivates adapting sparse attention, hierarchical processing, and memory to tabular ICL.

Tabular In-Context Learning

  • TabPFN [16] performs zero-shot prediction from synthetic-data pretraining, while TabICL [32] combines column embeddings, row interactions, and split-attention ICL prediction.The passages note TabPFN’s difficulty scaling to larger tables and TabICL’s dense attention and sequential pipeline limitations.Source: §2 · p. 3
  • TabDPT uses retrieval and diffusion-based representation learning, while TabPFN-v2 extends cell-based ICL to datasets above 10,000 samples; both retain scalability costs.TabDPT incurs diffusion overhead and dense attention, and TabPFN-v2 retains quadratic attention costs for high-dimensional tables.Source: §2 · p. 3
  • ContextTab [35] adds contextualized feature embeddings for heterogeneous data but retains single-scale processing and dense attention.The related-work synthesis identifies dense quadratic attention, uniform single-scale processing, and missing cross-component feedback as shared limitations.Source: §2 · p. 3

Sparse Attention Mechanisms

  • BigBird [43], Longformer [1], and Sparse Transformers [5] [20] use structured sparse attention for efficiency, but these methods had not been systematically adapted to tabular feature dimensions.Source: §2 · p. 3

Hierarchical and Multi-Scale Architectures

  • Funnel Transformers [6] and Swin Transformers use multi-scale processing, while Set Transformers [34] [10] use attention pooling; TabICL [32] lacks hierarchical feature grouping.Source: §2 · p. 3

Cross-Component Communication

  • Perceiver [19] and Perceiver IO [18] use latent bottlenecks to share information across modalities, but the cited approaches do not address ICL’s restriction on test examples influencing training representations.Source: §2 · p. 3

3 Proposed Approach: Orion-MSP

Orion-MSP extends TabICL with distribution-aware column embeddings, multi-scale sparse row interaction, and Perceiver memory to support zero-shot tabular prediction and bidirectional component communication. Its sparse attention and training-only memory with split-masked in-context learning target hierarchical feature dependencies and reduced attention complexity while refining representations and predicting test labels without leakage.

3.2 High-level Structure: From Data to ICL

  • Orion-MSP [32] combines column embedding, multi-scale sparse row interaction, bidirectional Perceiver memory, and a dataset-wise predictor for zero-shot prediction without gradient updates.The memory enables iterative refinement without information leakage.Source: §3.2 · p. 4

3.3 Column-wise Embedding

  • A column-wise Set Transformer embeds each cell using its value and column distribution; training-only ISAB inducing points generate distribution-conditioned affine weights and biases without test leakage.The inducing points summarize training-column distributions, and all samples can query them for contextualized embeddings.Source: §3.3 · p. 4, §3.3.2 · p. 6, §3.3.3 · p. 6

3.3.3 Weight and Bias Generation:

  • Inducing-point representations reflect column moments, multimodality, and outlier status, while the Set Transformer is permutation-invariant to sample order.The reported observations associate similar statistical moments with similar inducing-point representations and distinct attention patterns for outliers.Source: §3.3.3 · p. 7

3.4 Multi-Scale Sparse Row-Wise Interaction

  • To address variable feature counts, heterogeneous dependencies, and dense-attention cost O(m^2), Orion-MSP groups features across resolutions and uses PMA to learn dataset-adaptive groupings.Fine scales capture feature-level dependencies; coarser scales shorten sequences for broader reasoning, and scale representations are aggregated.Source: §3.4.1 · p. 7, §3.4.2 · p. 8

3.4.1 Motivation and Design Principles

  • Block-sparse attention combines local sliding windows, fully connected special tokens, and optional random links, reducing complexity from O(m^2) to O(m · (w + g + r)).The window preserves local structure, global tokens support long-range information flow, and random links increase reachability.Source: §3.4.1 · p. 8

3.4.5 Transformer Encoder per Scale:

  • Dedicated sparse Transformer encoders process each scale, after which CLS representations are averaged across scales, flattened, and normalized into row embeddings.Equal averaging balances fine- and coarse-grained information.Source: §3.4.5 · p. 9

3.4.6 Computational Complexity

  • For m∈[10, 100], w=8, Nglobal=4, and r=2, attention scales from O(m^2) to near-linear O(m · 14) while preserving local and global feature dependencies.The stated simplification assumes w, Nglobal, and r are much smaller than m.Source: §3.4.6 · p. 10

3.5 Cross-Component Memory with Perceiver Architecture

  • The Perceiver-style memory [19] enables leak-free communication by letting training samples write shared context that both training and test samples can read.This addresses the lack of backward adaptation, contextual refinement, and cross-dataset generalization in isolated sequential components.Source: §3.5 · pp. 10–11

3.6 Dataset-wise In-Context Learning

  • Dataset-wise ICL injects labels only into training representations and uses split-masked attention so test samples condition on training examples without leaking information back to training.It predicts test labels in a single forward pass without gradient-based parameter updates; test samples may also attend to one another for contextual reasoning.Source: §3.6 · pp. 13–14

4 Experimental Evaluation

Orion-MSP achieves the best overall zero-shot rank across the evaluated benchmark suites, with strong accuracy and F1 results. Its performance is examined across dataset sizes, widths, class balance, and application domains, using a common set of datasets available to all evaluated models.

4.1.1 Benchmark Suites and Datasets.

  • The evaluation uses TALENT, OpenML-CC18, and TabZilla, with results restricted to datasets shared by all evaluated models: 154, 63, and 27 datasets, respectively.Some datasets were excluded due to out-of-memory or CUDA issues, primarily affecting TabPFN-based architectures even on H200 GPUs.Source: §4.1.1 · p. 15

4.1.2 Models and Baselines.

  • The comparison includes six tabular foundation models and traditional baselines, with experiments run on NVIDIA L40S GPUs and H200 GPUs for memory-intensive cases.Traditional baselines include XGBoost, LightGBM, CatBoost, and Random Forest.Source: §4.1.2 · p. 15, §4.1.3 · p. 16

4.1.4 Evaluation Metrics.

  • Performance is measured using ACC, AUC-ROC, and weighted F1, while mean rank averages per-dataset accuracy ranks, with lower values indicating stronger consistency.Scalability is assessed against sample size, feature dimensionality, and class imbalance.Source: §4.1.4 · p. 16

4.2 Results

  • ORION-MSP attains the best overall zero-shot rank (3.58), with ACC/F1 of 0.8461/0.8360 on TALENT, 0.8722/0.8676 on OpenML-CC18, and 0.8821/0.8786 on TabZilla.Table 1 reports mean accuracy-based ranks and ACC and weighted F1 for each benchmark suite and in aggregate.Source: Table 1 · p. 16, §4.2 · p. 16
  • Classical methods retain mean accuracies of 0.833–0.861 with ranks around 6.0, while pretrained tabular foundation models generalize better; TabPFN follows Orion-MSP at overall rank 4.61.TabPFN leads TALENT with ACC 0.8514, while Orion-MSP peaks on OpenML-CC18 and TabZilla.Source: §4.2 · pp. 16–17
  • Orion-MSP performs well across small, medium, and large datasets, outperforming most baselines on small and medium data, while XGBoost leads ACC/F1 above 10K samples.It has the highest accuracy on narrow datasets and strong performance on medium and wide datasets, where dense TabICL has diminished scalability.Source: §4.2 · p. 17, Table 2 · p. 17, Table 3 · p. 17
  • Orion-MSP’s strongest gains occur on imbalanced datasets, where it ranks second with ACC = 0.8840 and F1 = 0.8731; in medical data it has the highest ACC = 0.8045 and F1 = 0.7916, while ranking second overall.On finance datasets it ranks first in mean rank (4.60), with ACC = 0.8158 and F1 = 0.8047.Source: §4.2 · pp. 17–18, Table 4 · p. 18, Table 5 · p. 18

Deep Analysis and Interpretation

  • The analysis attributes Orion-MSP’s results to fine- and coarse-grained feature processing, sparse attention, and Perceiver memory, but notes limited benefit from added complexity on simpler, low-dimensional datasets.The authors identify adaptive scale selection and dynamic sparsity control as future directions.Source: Deep Analysis and Interpretation · pp. 18–19

5 Conclusion

Orion-MSP combines multi-scale sparse attention and cross-component memory, and experiments report state-of-the-art zero-shot performance across diverse benchmarks. Its benefits are less pronounced on simple, low-dimensional datasets, motivating future adaptive scaling and sparsity.

5 Conclusion

  • Orion-MSP achieves state-of-the-art zero-shot performance across diverse benchmark suites and domain-specific healthcare and finance datasets, particularly on imbalanced, high-dimensional, and context-rich data.The authors report that it outperforms existing tabular foundation models in challenging real-world scenarios while maintaining practical scalability.Source: §5 · p. 19
  • The hierarchical design, sparse attention, and cross-component memory collectively contribute to robust generalization, efficient computation, and improved representation of complex interdependencies.Source: §5 · p. 19
  • The benefits of multi-scale sparse attention are less pronounced on simple, low-dimensional datasets, where its added architectural complexity may not be fully leveraged.Source: §5 · p. 19
  • Adaptive scale selection and data-aware sparsity scheduling are proposed as future work to adjust model complexity to dataset characteristics.The authors suggest these extensions could improve efficiency and generality.Source: §5 · p. 19

A Pretraining and Implementation Details

Pretraining uses synthetic tabular datasets generated from causal and tree-based priors, followed by a three-stage curriculum. The appendix also specifies implementation hardware and selected component hyperparameters.

A.1 Pretraining Data Generation

  • Synthetic datasets are generated with structural causal model priors to control diversity and cover varied statistical patterns across heterogeneous tabular schemas, distributions, and tasks.Source: §A.1 · p. 22

A.1.1 Structural Causal Model (SCM) Prior

  • Features in the SCM prior depend on parent features through sampled nonlinear activation functions with Gaussian noise, with a diverse activation set used to vary transformations.Source: §A.1.1 · pp. 22–23

A.1.2 Tree-Based SCM Prior

  • A tree-based SCM prior complements the smooth MLP-based prior with piecewise-constant relationships for categorical interactions and hierarchical decision boundaries.Decision-tree depths are sampled from {1, 2, 3, 4}.Source: §A.1.2 · p. 23

A.2 Pretraining Details

  • The three-stage curriculum progresses from fixed 1,024-sample datasets to variable contexts up to 40,000 samples, then trains the ICL component on 40,000–60,000 samples.Stage 3 freezes all components except TFicl; across stages the authors use Adam, gradient clipping at 1.0, and the schedule shown in Figure 3.Source: §A.2 · p. 23

A.3 Implementation Details

  • The implementation is built in PyTorch and trained on NVIDIA H200 GPUs.Source: §A.3 · p. 23

Hyperparameters

  • Column embedding uses dimension d = 128, 128 inducing points, four attention heads, three ISAB blocks, and a feedforward dimension of 256.Source: Hyperparameters · p. 24

Initialization and Training Considerations

The appendix specifies initialization and training safeguards alongside multi-scale row-interaction settings.

Initialization and Training Considerations

  • Inducing points are initialized from N(0, 0.022), while linear and attention weights use Xavier/Glorot uniform initialization; training uses column dropout pcol = 0.1 and gradient clipping at max-norm 1.0.Source: Initialization and Training… · p. 24

Hyperparameters

  • Multi-scale row interaction uses six transformer blocks across scales S = {1, 4, 16}, with window size 8, four global tokens, and two random links.Source: Hyperparameters · p. 24

Training Details

The appendix records the sparse-attention implementation, memory configuration, ICL settings, and initialization choices.

Training Details

  • Sparse attention uses PyTorch SDPA with a cosine learning-rate schedule and 5% warmup, using FP16 forward computation and FP32 softmax.Source: Training Details · p. 24

Hyperparameters

  • The Perceiver memory uses 32 slots, two write layers, two read layers, and four attention heads.Source: Hyperparameters · p. 24

Training Considerations

Initialization

  • The ICL transformer has 12 blocks, embedding dimension dR = 512, feedforward dimension 1024, and a maximum of 10 classes; label-encoder, transformer, and decoder weights use Xavier/Glorot uniform initialization.Source: Hyperparameters · p. 25, Initialization · p. 25

B Further Experiments

Orion-MSP has the best average rank across all three benchmarks, while per-dataset accuracy comparisons show strong performance against XGBoost and a favorable mean–variance trade-off against other methods.

B Further Experiments

  • Orion-MSP improves on XGBoost across benchmarks, with positive median relative accuracy on TabZilla and OpenML-cc18 and parity on TALENT.On TALENT, Orion-MSP has lower variance than most neural baselines; the comparison reports per-dataset relative accuracy over XGBoost.Source: B Further Experiments · p. 25
  • Orion-MSP matches or surpasses other tabular foundation models’ central performance and has the best mean–variance trade-off, while TabPFN and TabICL show broader variance and heavier left tails.Mitra and ContextTab show pronounced negative shifts and high variance; boosted-tree methods cluster near parity with XGBoost.Source: B Further Experiments · p. 25
  • The cross-dataset comparison averages per-dataset ranks across 11 methods and uses a Friedman test, followed when significant by Nemenyi post-hoc tests and α = 0.05 critical-difference diagrams.Methods connected by a critical-difference bar are not significantly different.Source: B Further Experiments · p. 25
  • 3.09 on TALENT, 3.32 on TabZilla, and 3.35 on OpenML-cc18 are Orion-MSP’s best average ranks across the three benchmarks.TabICL, TabPFN, and TabDPT typically trail by about 0.5–1.5 rank points; tree ensembles average 4.6–6.2, while ContextTab and Mitra rank around 9–10.Source: B Further Experiments · p. 25
  • The critical-difference diagrams support Orion-MSP’s top-performer status across benchmarks, with performance above tree ensembles on average and matching or exceeding pretrained tabular foundation models.The authors also report a favorable significance profile.Source: B Further Experiments · p. 25

C Datasets

The section points readers to benchmark dataset details and visualizes the evaluated datasets’ row and column distributions.

C Datasets

  • Dataset and benchmark details are summarized alongside Figure 6, which visualizes the evaluated datasets’ row and column distributions.The passage does not specify the distributions’ values.Source: C Datasets · p. 25

OpenML-CC18 Benchmark Datasets

The section identifies dataset tables for OpenML-CC18 and TALENT, while Figure 6 visualizes row and column distributions across evaluated datasets.

OpenML-CC18 Benchmark Datasets

TALENT Benchmark Datasets

TabZilla Benchmark Datasets

The supplied material identifies dataset tables for OpenML-CC18, TALENT, and TabZilla; TabZilla’s datasets are selected for evaluating neural-network performance on tabular data.

TabZilla Benchmark Datasets

OpenML-CC18 Benchmark Datasets

TALENT Benchmark Datasets

  • Table 7 lists TALENT’s automatically discovered datasets from multiple domains.Source: Table 7 · pp. 29–32

Related papers

Loading 2511.02818v3…