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

Let a classification task contain $n_{\mathrm{tr}}$ labeled context rows and $n_{\mathrm{te}}$ unlabeled query rows, with $n=n_{\mathrm{tr}}+n_{\mathrm{te}}$. After numerical preprocessing, the model receives

where, during pretraining, tasks are sampled with $m\leq 100$ active features and $C\leq 10$ classes. Query labels are used only to compute the pretraining loss and are not provided to the model.

Following TabICLv2 (Qu et al., 2026), the input projection uses circular feature grouping with offsets $(0,1,3)$:

Grouping is applied to active features before padding. A shared linear projection produces the initial cell token,

Only context rows receive label embeddings. Padded feature positions are masked throughout the sample-level encoder.

For an unordered token set $S\in\mathbb{R}^{N\times d}$, let $U\in\mathbb{R}^{K\times d}$ denote the seed vectors obtained by adding an MLP embedding of each normalized slot coordinate to its learned seed vector. Their fixed indices define the latent order, while the input tokens carry no positional encoding. The set read is

The seed slots aggregate information from the input set through cross-attention. Since the input tokens carry no positional information, Eq. (19) is invariant to permutations of $S$.

The ordered latent slots are mixed by a pre-normalized Mamba-2 residual block (Dao and Gu, 2024),

When token-aligned outputs are needed, the latent sequence is read back to the original set:

The readback is permutation equivariant with respect to $S$. Cross-attention provides set-to-slot aggregation and readback, while Mamba-2 mixes the ordered latent slots.

For each feature $j$, the variable-level encoder reads only the context-row tokens:

The resulting $K_{v}$ latent slots are processed by one Mamba-2 residual block and read back to the corresponding tokens of all rows. This set-read–Mamba-2–readback operation constitutes one variable-level block. SOMTab applies $L_{v}=3$ such blocks sequentially, with the output tokens of one block used as the input tokens of the next.

At the sample level, the $m$ updated feature tokens of each row $i$ are augmented with $K_{r}=4$ shared learned row-summary tokens.

The resulting $m+K_{r}$ tokens are first mapped to $K_{s}=4$ ordered latent slots. The slots are then processed sequentially by $L_{s}=3$ Mamba-2 residual blocks, with each block receiving the output of the preceding block. The final slots are denoted by $Z_{i}^{\mathrm{sample}}\in\mathbb{R}^{4\times d}$.

The row representation is obtained by flattening these slots:

Table 2 summarizes the main intermediate tensors.

Before the final reader, SOMTab adds a second label embedding to the context row representations:

Let $G=[g_{1};\ldots;g_{n}]$ and $G_{\mathrm{ctx}}=G_{1:n_{\mathrm{tr}}}$. Each reader block applies

All rows act as queries, while keys and values are restricted to the labeled context. Context rows interact with one another, and each query row attends independently to the context. After 12 reader blocks, a two-layer MLP maps the query representations to class logits.

Algorithm 1 summarizes the generation procedure. Sequence length, feature count, class count, and context size are determined by the training curriculum. Within each task, related scalar and categorical hyperparameters share task-specific distributions. For a scalar parameter, we first draw

use $(st,s(1-t))$ as beta-distribution parameters, and map the resulting draw to the required linear or logarithmic interval. For categorical parameters, task-specific probabilities are obtained by normalizing log-normal weights whose log-standard deviation is sampled from $[0.2,2.0]$. Sharing these distributions within a task induces correlated generator settings.

For each topologically valid pair $i<j$, the edge $i\rightarrow j$ is sampled as

where

The global, outgoing, and incoming scales are sampled log-uniformly from $[0.2,1.5]$, $[0.3,2.0]$, and $[0.3,2.0]$, respectively. The edge-level noise satisfies $\epsilon_{ij}\mid s_{\epsilon}\sim\mathcal{N}(0,s_{\epsilon}^{2})$, where $s_{\epsilon}$ is sampled log-uniformly from $[0.05,0.8]$. Before Bernoulli sampling, logits are clipped to $[-6,6]$ and probabilities to $[0.002,0.995]$.

When the modular component is active, the number of modules is sampled uniformly from $\{2,\ldots,5\}$. The module-pair bias includes a Gaussian term scaled by a draw from $[0.2,1.0]$ and a within-module diagonal bias sampled from $[0.5,2.5]$. The low-rank dimension is sampled uniformly from $\{2,\ldots,\min(8,L)\}$, with weight sampled log-uniformly from $[0.1,1.5]$.

Table 3 lists the graph-family mixture, and Figure 5 shows representative samples. After edge sampling, parent sets are truncated to bound the input dimension of the random node functions. Observed features are drawn with replacement from a latent pool that includes nodes sharing ancestry with the target, allowing correlated and occasionally duplicated columns.

The tail regime is sampled once per task:

The mixture contains three regimes:

1. 1.

   Independent Student-$t$ (0.70). The three degrees of freedom are sampled independently as $\nu_{g},\nu_{\mathrm{out}},\nu_{\mathrm{in}}\sim\mathcal{U}(1,10)$.

2. 2.

   Hierarchical Student-$t$ (0.20). A shared base degree of freedom is sampled from

   $$\nu_{0}\sim\begin{cases}\mathcal{U}(1,2),&0.35,\\ \mathcal{U}(2,5),&0.40,\\ \mathcal{U}(5,10),&0.25.\end{cases}$$

   (33)

   Each realized value is $\nu_{q}=\operatorname{clip}(\nu_{0}e^{\delta_{q}},1,10)$ with $\delta_{q}\sim\mathcal{N}(0,0.15^{2})$ and $q\in\{g,\mathrm{out},\mathrm{in}\}$.

3. 3.

   Symmetric Pareto (0.10). With $\alpha\sim\mathcal{U}(1.5,5.0)$, $U\sim\mathcal{U}(0,1)$, and an equiprobable sign $s\in\{-1,+1\}$, a raw draw is

   $$r=s\left(U^{-1/\alpha}-1\right).$$

   (34)

   Vector-valued draws are normalized by the $0.75$ quantile of $|r|$ before applying their sampled scale.

In every regime, $g$ is clipped to $[-5,5]$, and the incoming and outgoing corrections are clipped to $[-6,6]$. Figure 6 shows representative survival curves before clipping.

The dimension of each latent node is sampled logarithmically from 1 to 32. Root states are drawn from one of four families: independent Gaussian, uniform on $[-1,1]^{d}$, uniform in the unit ball, or a random-covariance model. With probability $0.35$, a random linear map and activation are applied to the root state.

For a non-root node with parent states $\{z_{p}:p\in\operatorname{Pa}(j)\}$, the parents are either concatenated and passed through a single function, or processed separately and combined by a normalized sum, product, maximum, or log-sum-exp. The single-function form is used for one parent and is sampled with probability $0.5$ otherwise. Table 4 lists the function families. Their probabilities are sampled at the task level rather than fixed to a uniform distribution.

The activation library comprises GELU, ReLU, $\tanh$, sine, cosine, sigmoid, and absolute value. After propagation, node states are standardized, dimension-wise reweighted, normalized by their mean sample norm, and rescaled by a factor sampled log-uniformly from $[0.1,10]$.

The observation model samples a task-level categorical rate

and converts each selected latent feature to categorical form independently with probability $q_{\mathrm{cat}}$. Numerical features are one-dimensional latent projections. With probability $0.35$, a numerical feature is min–max scaled and transformed by a Kumaraswamy distribution whose two shape parameters are sampled log-uniformly from $[0.2,5]$.

Categorical features have cardinality between 2 and 9 and use one of two converters. The nearest-center converter uses at most 12 latent dimensions and an $\ell_{p}$ distance with log-uniform $p\in[0.5,4]$. The softmax converter uses a standardized random projection, category biases, and a separation parameter sampled log-uniformly from $[0.1,10]$. An additional numerical-to-categorical conversion is applied with task-level probability $0.2$, followed by feature standardization, outlier clipping, random feature permutation, and zero padding.

One to three target nodes are combined by the same random function mechanism to produce a standardized continuous target. The class count is set to $C=2$ with probability $1/2$ and is otherwise sampled uniformly from $\{2,\ldots,10\}$. The resulting marginal probability is $5/9$ for binary tasks and $1/18$ for each class count from 3 to 10. The continuous target is discretized using either randomly selected observed-value thresholds or Gaussian thresholds, and the class identities are then permuted.

Degenerate tasks are removed using validity checks that do not require fitting a separate predictor. A task is rejected if the features or target contain non-finite values, if fewer than $\max(1,\min(2,m))$ features have non-negligible variance, or if the largest absolute feature–target marginal correlation is at most $0.015$ on a subsample of at most 2048 rows. Constant columns are removed after classification conversion. The row order is then resampled up to ten times until the context and query partitions contain the same set of at least two classes.

The main model is trained entirely on DCH-TailMix tasks. Stage 1 covers a broad range of task sizes and context fractions. Stage 2 continues from the Stage-1 checkpoint with a lower learning rate. Table 5 summarizes the curriculum.

Stage 2 uses polynomial decay with power two. Each synthetic task contains between 2 and 100 features and at most ten classes, with the class-count distribution specified in Appendix B.5. NumPy and PyTorch use random seed 42 in all training runs.

Training uses a hybrid Muon–AdamW optimizer. Muon is applied to non-embedding matrix parameters with momentum $0.95$ and five Newton–Schulz iterations. Embedding and lower-dimensional parameters use AdamW with $(\beta_{1},\beta_{2})=(0.9,0.95)$ and $\epsilon=10^{-8}$. Weight decay is set to zero, and the global gradient norm is clipped at 1.0.

Training is distributed over four NVIDIA A100 GPUs with 40 GB of memory per GPU. Each global batch contains 64 independently generated tasks. The per-device micro-batch size is one, and gradient accumulation yields the global batch size. Stage 1 requires approximately 17 days, followed by approximately six days for Stage 2 on the same hardware. The final Stage-2 checkpoint is used for evaluation.

For each TALENT classification dataset (Ye et al., 2024), the official training and validation splits form the labeled context, and the official test split forms the query set. Accuracy, Macro F1, and log-loss are computed on the test rows, with equal weight assigned to each dataset in aggregate results.

For model $a$ and dataset $d$, log-loss is normalized across the methods in the common result table:

All methods in an aggregate comparison are evaluated on the same intersection of datasets.

For methods rerun on TALENT, timing includes model fitting, when applicable, and prediction. Dataset loading, external preprocessing, model construction, and checkpoint loading are excluded. Runtime is normalized by the number of labeled context rows and reported per 1000 context examples; aggregate runtime is the median across datasets. Peak GPU memory is the maximum allocated CUDA memory during the same interval.

The controlled context-scaling experiment fixes $n_{\mathrm{te}}=1024$, $m=32$, and $C=10$, and varies $n_{\mathrm{tr}}$ from 512 to 28,672. All foundation models use one estimator, with the default eight-estimator TabICLv2 configuration reported separately. Each point is measured after one warmup run and reports the median of three repetitions on identical synthetic inputs.

The controlled ablation compares synthetic priors and final ICL readers under the same training budget. All variants are initialized randomly, trained for 20K steps, and evaluated every 5K steps. Evaluation uses the same subset of 15 TALENT datasets where SOMTab performs relatively less favorably than TabICLv2 in the main benchmark. This subset is used only for the ablation analysis.

At each evaluated checkpoint, DCH-TailMix achieves higher mean accuracy than the Cauchy and SCM–Tree priors under the same architecture (Hollmann et al., 2025; Qu et al., 2025; Qu et al., 2026). The Mamba Reader variant, which replaces the final attention reader with Mamba, remains less accurate throughout training. These results support the use of Mamba for representation construction and attention for final in-context retrieval.

Table 7 reports per-dataset accuracy on 20 TALENT classification datasets. The final 15 datasets constitute the fixed evaluation subset used in the controlled early-training ablation in Appendix D.3.

We follow the public TabArena evaluation protocol and include the models provided by the benchmark. Figure 9 compares the performance–efficiency trade-off, while Figure 10 reports pairwise comparisons between models.

Figure 10 reports pairwise win rates on TabArena. Each entry shows the fraction of datasets where the model in the row achieves higher predictive performance than the model in the column. This comparison reveals how frequently one model outperforms another across individual datasets rather than relying only on averaged benchmark scores.