Orion-Bix: Bi-Axial Attention for Tabular In-Context Learning

Research

Orion-Bix: Bi-Axial Attention for Tabular In-Context Learning

November 28, 2025

Paris Lab

Mumbai Lab

Share this Post:

current-url

1. Introduction

Tabular data remain the most common form of real-world data, spanning domains such as healthcare, finance, and scientific research. Despite the remarkable progress of deep learning in natural language processing [liu2023pre, yu2024natural] and computer vision [goldblum2023battle], gradient boosted trees (GBTs) remain the predominant state-of-the-art (SOTA) for tabular prediction tasks. In other data modalities, foundation models—particularly Large Language Models (LLMs) [DBLP:conf/nips/00020D24, DBLP:conf/nips/LiZLML0HY24]—have significantly advanced the ability to tackle new tasks and few-shot learning. This is largely due to their remarkable in-context learning (ICL) capabilities [DBLP:conf/iclr/Zhang0YOKZY025, DBLP:conf/uss/CarliniTWJHLRBS21], which enable them to capture patterns directly from prompts without updating their parameters. This success combined with the pervasiveness of tables have spurred interest in tabular foundation models [DBLP:conf/icml/BreugelS24].

Although LLMs are primarily designed to process natural language, recent efforts have explored fine-tuning them for tabular data tasks [tabllm, tmlr]. These approaches typically rely on table serialization, which is the process of converting table rows into text or sentences suitable for tokenization. For instance, [DBLP:conf/nips/0001PS24] fine-tuned a Llama 3-8B model on a large corpus of serialized tables and demonstrated that this strategy can outperform traditional tree-based models in few-shot scenarios. However, such language model–based approaches face inherent challenges. Their limited context windows restrict the number of serialized examples that can be processed simultaneously (e.g., up to 32 or 64 shots in [DBLP:conf/nips/0001PS24]), and it remains uncertain whether LLMs can reliably interpret and reason over numerical values [DBLP:conf/naacl/ThawaniPIS21].

Recently, tabular in-context learning has emerged, adapting the ICL paradigm—central to large language models—to tabular data, enabling pretraining across diverse tables and rapid task adaptation without gradient updates. TabPFN [tabpfn] pioneered this approach by meta-training a transformer on synthetic datasets generated via structural causal models. Its encoder–decoder design lets test samples attend to training examples for zero-shot prediction, but alternating column- and row-wise attentions make large training sets computationally costly. TabDPT [tabdpt] achieves comparable performance using similarity-based retrieval, though its diffusion process adds overhead. TabPFN-v2 [tabpfn2] extended row-based encoding to datasets exceeding 10,000 samples. TabICL [tabicl] further refined this with a table-native transformer comprising column embeddings, row interactions, and an ICL head. Its SetTransformer-based column encoder and label-aware in-context learner achieve state-of-the-art results, but its row encoder (tf_row) applies a single homogeneous attention over all features, ignoring local groupings, multi-scale interactions, and structured aggregation. This limitation is pronounced in high-dimensional, heterogeneous tables where features naturally form semantic groups (e.g., demographics, vitals, labs) and dependencies span multiple scales. Additionally, TabICL treats synthetic tables as individual supervised tasks rather than support/query episodes, producing an implicit few-shot signal that may misalign with test-time objectives.

To address these limitations, we introduce Orion-Bix¹¹ ,which enhances TabICL with two key improvements: (1) a biaxial row encoder replacing tf_row with complementary attention modes—standard cross-feature, grouped, hierarchical, and relational—aggregated via multiple CLS tokens, explicitly modeling local groups, coarse-scale interactions, and global patterns; and (2) a meta-learning training regime that constructs support/query episodes from synthetic tables, optionally using kNN-based support selection, aligning training with few-shot test-time objectives. Orion-Bix preserves TabICL’s strengths in column-wise SetTransformer embeddings and label-aware ICL while addressing its row-level and training limitations.

2. Orion-Bix: Proposed Approach

Orion-Bix retains two core TabICL [tabicl] components: the SetTransformer column encoder and the label-aware in-context learner, while introducing two major changes: (1) a biaxial row encoder that replaces the single-stage tf_row and provides structured multi-scale feature reasoning, and (2) a meta-learning training regime that constructs explicit support/query episodes from synthetic tables, optionally using kNN-based support selection. An overview of the complete architecture is shown in Figure 1.

‍

Figure 1: An overview of ORION-BiX architecture. A column-wise SetTransformer-based embedder maps the input table $X \in \mathbb{R}^{B \times n \times m}$ to column embeddings $E \in \mathbb{R}^{B \times n \times m^\prime \times d}$, where $m^\prime = m + N_{\mathrm{CLS}}$ includes reserved CLS slots. The biaxial row encoder reshapes $E$ into per-row sequences $X^\prime \in \mathbb{R}^{(B \cdot n) \times m^\prime \times d}$ and applies a stack of BiAxialAttentionBlocks combining full cross-feature attention $(X^\prime \rightarrow X_1)$, local grouped attention $(X_1 \rightarrow X_2)$, hierarchical attention across coarse feature partitions $(X_2 \rightarrow X_2^\prime)$, and structured relational attention $(X_2^\prime \rightarrow X_3)$. A multi-CLS attention layer $\mathrm{CLSAttn}(\mathrm{CLS}, X_3)$ aggregates each row into a multi-aspect representation $R \in \mathbb{R}^{B \times n \times (N_{\mathrm{CLS}} \cdot d)}$. The label-aware ICL adds projected support labels to support row embeddings and uses a masked cross-attention Transformer to predict labels for query rows.

‍

2.1. Problem Setting and Data Representation

We consider a tabular dataset as a sequence of rows with heterogeneous columns. For each task (dataset), let:

$ X \in \mathbb{R}^{n \times m} $: a table with $n$ rows and $m$ features (after preprocessing).
$ y \in \{0, \ldots, C - 1\}^n $: labels over $C$ classes.

We split rows into a support set and a query set. The first $n_{\text{train}}$ rows form the support set, the remaining $n_{\text{test}} = n - n_{\text{train}}$ rows form the query set. During meta-training and inference, Orion-BiX receives the full table $X$ and only the support labels

$ y_{\text{train}} = y_{1:n_{\text{train}}}. $

The goal is to predict labels for the query rows conditioned solely on the support rows and their labels, i.e., in a tabular in-context learning regime.

In practice, Orion-BiX operates on a batch of tasks. Inputs have shape

$ X \in \mathbb{R}^{B \times n \times m}, $

where $B$ is the number of tasks (tables) in the batch. A per-task feature-count tensor

$ d \in \mathbb{N}^B $

tracks how many of the $H$ columns are active in each table (the remaining columns are padding).

‍

2.2. Column-wise Embedding

We adopt the original TabICL [tabicl] column-wise embedder to map each feature to a d-dimensional representation per cell. The embedder uses a Set Transformer, TFcol, that treats each column as a permutation-invariant set of values across all rows.

2.2.1. Input and CLS Reservation.

The column embedder receives $ X \in \mathbb{R}^{B \times n \times m} $. We reserve a fixed number $N_{\text{CLS}}$ of feature slots (typically 4) at the beginning of the feature dimension to host class tokens at a later stage.

These slots are padded with a special skip value, so they do not influence the subsequent projections.

2.2.2. Skippable Linear Projection.

Each scalar value $x_{b,n,m}$ is passed through a skippable linear layer:

\[ \tilde{x}_{b,n,m} = \begin{cases} \text{skip\_value}, & \text{if } x_{b,n,m} = \text{skip\_value}, \\ w_h x_{b,n,m} + b_h, & \text{otherwise}, \end{cases} \]

where $w_h$ and $b_h$ are learnable parameters. This yields an intermediate representation $\text{src} \in \mathbb{R}^{B \times m' \times n \times D}$, with $m' = m + N_{\text{CLS}}$ and embedding dimension $D$.

The skip logic ensures that missing or padded entries do not pollute the learned representations.

2.2.3. SetTransformer Over Rows for Each Column.

Columns are treated as small sets of row values: for each feature $m$, we consider $\{\tilde{x}_{b,1,m}, \ldots, \tilde{x}_{b,n,m}\}$ and apply a stack of SetTransformer-like attention blocks along the row dimension. The implementation supports two attention types:

Linear attention, with a configurable feature map (e.g., ELU, identity, or hedgehog), enabling $\mathcal{O}(n)$ complexity and stable behavior for long sequences.
Induced self-attention, where a small set of learned inducing points mediate interactions between rows, reducing complexity from $\mathcal{O}(n^2)$ to approximately $\mathcal{O}(n)$.

A key design choice is that support/query leakage is controlled at this stage. During training, the SetTransformer can be restricted to attend only to support rows; at inference, it can optionally embed support and query jointly depending on a configuration flag.

2.2.4. Feature-Wise Weight–Bias Parameterization.

Rather than outputting a single vector per column, the column embedder produces feature-specific affine parameters. For each $(b, n, m)$, two linear heads produce:

$ W_{b,n,m} \in \mathbb{R}^D, \quad b_{b,n,m} \in \mathbb{R}^D $,

and the final per-entry embedding is

$ e_{b,n,m} = x_{b,n,m} \cdot W_{b,n,m} + b_{b,n,m}. $

This yields an embedding tensor $ E \in \mathbb{R}^{B \times n \times m' \times D} $. The learned weights and biases capture column-wise statistics, scaling behavior, and type information, and allow the model to normalize and warp distributions in a data-driven way.

2.3. Biaxial Row Encoder: Multi-Scale Feature Interactions

The column embeddings $E \in \mathbb{R}^{B \times n \times m' \times d}$ encode per-cell feature representations. The row encoder aggregates these into row-level summaries suitable for in-context learning.

The original TabICL row encoder tf_row applies a single stack of Transformer encoder layers directly over the feature dimension. While effective, this design treats all features as a flat sequence and relies on one homogeneous attention pattern to model: (i) local relations within semantically related feature subsets; (ii) long-range interactions between distant columns; and (iii) global aggregation into a single row representation.

In high-dimensional, heterogeneous tables, this uniform treatment leads to several limitations:

Lack of explicit grouping: Tabular features often form natural groups (e.g., demographics, vitals, lab tests), but a single attention layer has no explicit mechanism to focus attention within these groups before reasoning across them.
Insufficient multi-scale structure: Capturing both local and global patterns with the same attention pattern can force the model into compromises, especially when some dependencies are short-range (within groups) while others are long-range (across groups).
Representation bottleneck: Collapsing a row into a single vector after a flat attention stack can obscure distinct modes of information (e.g., rare but important features vs. common global signals).

To overcome this, Orion-Bix replaces the flat row encoder with a biaxial attention module that applies multiple, specialized feature-space attentions and aggregates their outputs through multi-CLS tokens, enabling structured, multi-scale row representations.

2.3.1 Input Reshaping.

Starting from the column embeddings $E \in \mathbb{R}^{B \times n \times m \times d}$, we reshape the representation to treat each row’s features as a sequence:

\[ E \;\rightarrow\; X' \in \mathbb{R}^{(B \cdot n) \times m' \times d}, \]

where each element $X'_i \in \mathbb{R}^{m' \times d}$ represents the feature-embedding sequence corresponding to a single row.

2.3.2. Bi-Axial Attention Block.

Each BiAxialAttentionBlock applies four attention modules in sequence over the feature axis:

Standard attention: full cross-feature self-attention capturing generic dependencies:
\[ X_1 = \mathrm{Attn}_{\mathrm{std}}(X_0, X_0, X_0), \]
Grouped attention: local attention within $G$ feature groups of size $\lfloor m' / G \rfloor$. Within each group, we apply local attention:
\[ X_2 = \mathrm{Attn}_{\mathrm{group}}(X_1), \]

This sharpens local interactions between features that naturally cluster.
Hierarchical attention: coarse-scale interactions between two large partitions of the feature set $X_2 = [X_2^{(1)}, X_2^{(2)}]$:
\[ \hat{X}_2^{(1)} = \mathrm{Attn}_{\mathrm{hier}}(X_2^{(1)}, X_2^{(2)}, X_2^{(2)}), \] \[ \hat{X}_2^{(2)} = \mathrm{Attn}_{\mathrm{hier}}(X_2^{(2)}, \hat{X}_2^{(1)}, \hat{X}_2^{(1)}). \]

This enforces coarse-scale interactions between large blocks of features, useful when long-range dependencies span different column subsets.
Relational attention: a second full self-attention over already structured features:
\[ X_3 = \mathrm{Attn}_{\mathrm{rel}}(\hat{X}_2, \hat{X}_2, \hat{X}_2), \]

This performs another full self-attention over features, now operating on representations that already encode both local and hierarchical structure.

2.3.3. Multi-CLS Aggregation.

After the four feature-attention passes, each block introduces $N_{\mathrm{CLS}}$ learnable class tokens:

\[ \mathrm{CLS} \in \mathbb{R}^{N_{\mathrm{CLS}} \times D}. \]

For each row (each element of the batch), these CLS tokens are broadcast and used as queries in a dedicated CLSAttention module:

\[ \mathrm{CLS}' = \mathrm{CLSAttn}(\mathrm{CLS}, X_3), \]

where CLS tokens attend over the feature sequence, followed by a small feed-forward network with residual connections and normalization. The output of one BiAxialAttentionBlock is thus a set of CLS summaries per row:

\[ \mathrm{CLS}' \in \mathbb{R}^{(B \cdot n) \times N_{\mathrm{CLS}} \times D}. \]

Stacking $L_{\mathrm{row}}$ such blocks refines these CLS tokens iteratively. Finally, we flatten the CLS dimension:

\[ R' \in \mathbb{R}^{(B \cdot n) \times (N_{\mathrm{CLS}} \cdot D)}, \]

and reshape back to per-task shape

\[ R \in \mathbb{R}^{B \times T \times D_{\mathrm{row}}}, \qquad D_{\mathrm{row}} = N_{\mathrm{CLS}} \cdot D, \]

which serves as the input to the in-context learner.

This biaxial design allows Orion-Bix to jointly model:

local interactions within feature groups,
structured cross-group dependencies at a coarse level, and
global relational patterns across all features,

while avoiding the information bottleneck of a single pooled vector via multiple CLS tokens.

2.4. In-Context Learner

The final module, denoted ICLearning, performs in-context inference: given row representations and support labels, it predicts labels for the query rows.

2.4.1. Label Injection.

Given row encodings $ R \in \mathbb{R}^{B \times n \times D_{\mathrm{row}}} $ and support labels $ y_{\mathrm{train}} \in \{0, \ldots, C - 1\}^{B \times n_{\mathrm{train}}} $, we embed labels into the same space as rows:

\[ \ell_{b,n} = \mathrm{LabelEmbed}(y_{\mathrm{train}, b, n}) \in \mathbb{R}^{D_{\mathrm{row}}}, \]

where LabelEmbed is implemented as a one-hot encoding followed by a linear projection (the OneHotAndLinear layer). These label embeddings are added to the support row representations:

\[ R_{b,t} \leftarrow R_{b,t} + \ell_{b,t}, \qquad \text{for } t < n_{\mathrm{train}}. \]

This explicitly conditions the row features on the observed labels while leaving query rows unlabeled.

2.4.2. ICL Encoder with Split Attention Mask.

The label-conditioned sequence $ R $ is processed by an encoder (a stack of attention blocks) over the row dimension. To respect the support/query structure, we use an integer mask $ n_{\mathrm{train}} $ that implements:

The first $ n_{\mathrm{train}} $ positions (support rows) can attend only to each other.
The remaining $ n_{\mathrm{test}} $ positions (query rows) can attend to all support rows but not to other query rows.

Formally, for a row index $ t $ and key index $ s $,

\[ \mathrm{mask}(t, s) = \begin{cases} 0, & t < n_{\mathrm{train}},\; s < n_{\mathrm{train}}, \\ 0, & t \ge n_{\mathrm{train}},\; s < n_{\mathrm{train}}, \\ -\infty, & \text{otherwise}. \end{cases} \]

This mask is applied inside the attention kernels (both standard and linear attention variants) to enforce the desired information flow. Support representations cannot depend on queries, and each query prediction depends only on support rows and its own features.

2.4.3. Decoder and Hierarchical Classification.

The encoder outputs $ Z \in \mathbb{R}^{B \times n \times D_{\mathrm{row}}} $. A small MLP decoder maps each row to class logits:

\[ \hat{y}_{b,t} = f_{\mathrm{dec}}(Z_{b,t}) \in \mathbb{R}^{C_{\max}}, \]

where $ C_{\max} $ is the maximum number of classes the model is trained to handle natively.

Real tasks may have $ C > C_{\max} $. Orion-Bix transparently handles such cases via a hierarchical classification tree:

For a task with $ C $ unique labels, we recursively group the labels into at most $ C_{\max} $ super-classes at each level, forming internal nodes and leaf nodes.
Internal nodes predict group assignments; leaf nodes specialize in fine-grained labels.
Predictions at test time traverse this tree, combining probabilities using the chain rule to produce calibrated probabilities over all $ C $ original classes.

When $ C \le C_{\max} $, Orion-Bix uses the flat decoder directly.

2.5. Meta-Training with Synthetic Episodic Data

TabICL is trained on synthetic tabular data drawn from a configurable prior, optimizing a global objective that encourages robust in-context behaviour across a wide range of tasks. However, the original training schedule treats each synthetic dataset as a single supervised task, without explicitly structuring the optimization around support/query episodes. This leads to two limitations:

Implicit few-shot signal: The model learns to infer from sequence prefixes, but the few-shot structure (support vs. query) is only partially enforced by attention masks; it is not the primary unit of optimization.‍
Uncontrolled support selection: Support rows are not explicitly chosen to be informative or diverse with respect to query rows; many updates are driven by redundant or suboptimal supports.

Orion-BiX adopts an explicit meta-learning perspective on top of a similar synthetic data prior:

Synthetic tables are generated from a broad prior over feature counts, label spaces, sequence lengths, and distributions, and treated as a pool of potential tasks.
Episode Construction: An EpisodeGenerator converts these tables into many small episodes, each defined by a support set, a query set, and a per-task feature-count d. Episodes are formed either by random splits or by kNN-based support selection that explicitly chooses support rows that are both relevant and diverse with respect to queries.
A MetaLearningDataset yields episodes in manageable chunks, and the MetaLearningTrainer processes thousands of episodes per update via micro-batching, gradient accumulation, mixed precision, and (optionally) distributed data parallelism.

This explicit meta-learning formulation brings several advantages over the original TabICL training:

Stronger alignment with the test-time objective: The model is always optimized in the exact regime in which it will be evaluated: given a small support set and a query set, infer query labels purely from in-context information.‍
Better use of synthetic diversity: Each synthetic table can yield many distinct episodes with different support/query splits. This increases the effective number of tasks seen during training and exposes the model to a broader spectrum of few-shot situations
‍Support-set quality control: kNN-based episode construction selects support examples that are both close to and diverse for the query set, reducing the proportion of updates driven by uninformative or redundant supports.‍
Stability and scalability: Micro-batching, gradient accumulation, and AMP allow Orion-Bix to train on large episode counts without exhaustively increasing memory usage, while DDP and checkpoint management make multi-GPU training robust.

In summary, Orion-Bix keeps the core strengths of TabICL [tabicl]—column-wise SetTransformer embeddings and a label-aware in-context Transformer—but introduces a biaxial row encoder that better matches the structure of tabular feature spaces and a meta-learning training regime that more directly optimizes for few-shot in-context performance.

2.6. Inference Pipeline and Practical Interface

For practical deployment on real-world tabular datasets, Orion-BiX is wrapped in a scikit-learn–compatible classifier that automates preprocessing and uses an ensemble of transformed views.

2.6.1. Preprocessing and Feature Engineering.

Given an input table X, the wrapper:

detects numerical and categorical features and converts all columns to numeric form;
imputes missing numerical values (e.g., with medians) and handles categorical missingness consistently;
optionally applies one of several normalization schemes (none, power transform, quantile normalization, robust scaling);
clips outliers beyond a configurable z-score threshold;
applies feature shuffling strategies (none, circular shift, random, Latin patterns) to build diverse column orders.

2.6.2. Ensemble of Transformed Views.

The preprocessor constructs multiple transformed “views” of the dataset, each corresponding to a choice of normalization method and feature permutation. For each view:

a support/query split is formed (e.g., using a subset of training points as the support set),
the transformed table is passed through the Orion-BiX model, which produces logits or probabilities for the query rows.

Across views:

logits are re-aligned to correct for any class shifts induced by permutations,
predictions are averaged over ensemble members,
an optional temperature-scaled softmax converts logits to probabilities.

This ensemble scheme improves stability and robustness, especially for datasets with skewed distributions or strong feature-order effects.

Summary.

In summary, Orion-BiX combines:

a column embedder that learns distribution-aware feature embeddings,
a biaxial row encoder that models feature interactions at several structural scales and compresses them into multiple CLS summaries, and
a label-aware in-context learner that uses masked attention and hierarchical classification to handle few-shot tasks with arbitrary label spaces,

all trained via episodic meta-learning on diverse synthetic tables and exposed through an inference pipeline that integrates seamlessly with standard tabular workflows.

3. Experimental Evaluation

We evaluate Orion-Bix against TabICL and other baselines to highlight the benefits of biaxial attention and meta-learning. Our experiments focus on three aspects: (i) domain-specific performance on datasets with natural feature structure, (ii) support set quality to assess robustness, and (iii) few-shot learning curves across varying support sizes. These analyses target the improvements motivating Orion-Bix: structured multi-scale feature interactions and enhanced few-shot adaptation.

3.1. Evaluation Setup

3.1.1. Datasets.

We construct domain-specific evaluation suites by grouping datasets from public benchmarks (e.g., TALENT [talent] and OpenML-CC18 [openmlcc18]) according to application domain, allowing assessment in contexts where feature structure and heterogeneity are most relevant. The domains are: Medical Finance. These domains exhibit the high-dimensional, structured, and multi-scale features that Orion-Bix is designed to handle.

3.1.2. Evaluation Metrics.

For each dataset, we report overall classification accuracy (ACC), class-weighted F1 to account for imbalance, and mean rank across datasets within each domain based on accuracy. All models use official train/test splits unless noted.

3.2. Domain-Specific Performance

Table 1:Domain-specific performance for Medical, Finance and Energy datasets from the benchmark suites. Formatting: Bold = 1st place; underlined= 2nd place within each group.

ModelsMedic

Models	Medical			Finance
Models	Rank	ACC	F1	Rank	ACC	F1
XGBoost	6.32	0.7834	0.7669	6.62	0.7958	0.7885
RandomForest	6.38	0.7779	0.7752	7.32	0.8052	0.8001
CatBoost	6.36	0.7784	0.7594	5.82	0.8117	0.8015
LightGBM	5.32	0.7949	0.7614	6.17	0.8095	0.7974
TabICL	5.54	0.7819	0.7696	6.60	0.8125	0.7942
Orion-Bix	4.10	0.7893	0.7759	5.39	0.8206	0.8125
TabPFN	5.04	0.7984	0.7857	7.17	0.8094	0.7919
Mitra	10.77	0.3935	0.2863	13.67	0.5340	0.4250
ContextTab	8.66	0.6681	0.6129	11.25	0.7430	0.6834
TabDPT	6.86	0.7764	0.7641	8.00	0.8080	0.7960

Formatting: 1st place; 2nd place within each group.

Table 1 reports results on Medical, Finance, and Energy domains. Orion-Bix achieves the best mean rank in Medical (4.10) and Finance (5.39), outperforming TabICL and showing competitive performance against gradient-boosted baselines.

Medical. Orion-Bix attains rank 4.10 (vs. TabICL 5.54), accuracy 0.7893, and F1 0.7759. While TabPFN slightly exceeds accuracy (0.7984), Orion-Bix’s superior ranking indicates more consistent performance across diverse datasets. The gains reflect the benefit of biaxial attention on grouped and hierarchical features typical in medical data.‍
Finance. Orion-Bix leads with rank 5.39 (vs. TabICL 6.60), accuracy 0.8206 (vs. 0.8125), and F1 0.8125 (vs. 0.7942). Finance datasets feature strong hierarchies and multi-scale dependencies, where biaxial attention better captures structured relationships, notably improving F1 by +1.83 points over TabICL.

These results highlight two design benefits of Orion-Bix:

Biaxial attention for structured features: Gains are largest in domains with natural feature groups and hierarchies, confirming that modeling local (grouped), coarse-scale (hierarchical), and global (relational) interactions improves representations over a single attention stack.‍
Multi-CLS aggregation: Maintaining multiple CLS tokens preserves distinct feature aspects, capturing both rare patterns and common global signals, leading to more robust predictions.

‍

‍

3.3 Few-Shot Performance

We assess performance as a function of support set size, $k \in \{5, 10, 20, 32, 64, 128\}$, to test whether Orion-Bix’s meta-learning benefits are most pronounced in the very few-shot regime.

Protocol.

For each dataset and support size $k$, $k$ examples are sampled from the training set (all classes represented) to form the support set. Models are evaluated on the full test set, repeated across multiple seeds, and results are averaged. Accuracy is then plotted against support size.

Results.

Figure 2 shows:

Orion-Bix attains the best accuracy for very low shot sizes ($k \leq 32$), outperforming all three baselines by 2–4 points, indicating that meta-training on diverse episodes improves data efficiency.
Accuracy increases monotonically with $k$ for all models; gains are largest from $k = 5$ to $k = 32$, while improvements beyond $k = 64$ are small, indicating diminishing returns from additional support examples.
As $k$ grows, the gap narrows: on TALENT, Orion-Bix stays clearly best across all $k$, while on OpenML-CC18, TabPFN continues improving and slightly surpasses Orion-Bix at $k = 128$. In both cases, TabPFN and Orion-Bix remain well ahead of TabICL and TabDPT. Overall, Orion-Bix is strongest in the few-shot regime.

3.4 Support Set Quality Analysis

We study how support set selection affects few-shot performance by comparing random vs. diverse (kNN-based) sampling at inference. This tests whether Orion-Bix’s meta-learning, which constructs diverse support/query episodes during pre-training, improves robustness to support quality.

Figure 3:Accuracy of Orion-Bix and TabICL on OpenML and TALENT benchmarks. Bars indicate support selection strategies, averaged over datasets for each model.

Protocol.

Each test dataset is split 80/20 into train/test sets. A support set of size k=32 is drawn using: (i) Uniform sampling, ensuring all classes are represented, and (ii) Diverse (kNN-based), selecting examples close to test queries and maximally diverse. Models are evaluated with standard ICL on the same support/query splits, repeated across multiple seeds for statistical significance.

Motivation.

This protocol measures inference-time robustness: diverse selection simulates a scenario where practitioners can actively choose informative support examples, while random selection represents default, arbitrary support sampling.

Results.

Figure 3 shows that: (1) Uniform sampling slightly outperforms kNN-based selection (by 3–5 points), indicating few-shot performance is not highly sensitive to support curation; (2) Orion-Bix consistently exceeds TabICL under both schemes, with relative ordering unchanged, suggesting meta-training improves overall in-context generalization rather than depending on a specific support selection heuristic.

4 Conclusion

In summary, Orion-Bix preserves the core strengths of TabICL while adding a biaxial row encoder and an episodic meta-learning schedule that directly targets few-shot in-context performance.

‍

Read the paper: https://arxiv.org/abs/2512.00181