Can activation alignment close the context gap in tabular in-context learning?

Every prediction reprocesses the whole context

Tabular foundation models such as TabPFN-3 and TabFM predict by in-context learning: a labeled training set is supplied as context alongside an unlabeled query, and the model produces a prediction in a single forward pass without weight updates. Unlike tree-based models, where training and inference are decoupled, these models must carry all training examples through every forward pass, so each prediction is expensive.

Shrinking the context lowers that cost but degrades accuracy, because predictive quality depends heavily on how many in-context examples are available. The paper frames this as a sample-efficiency bottleneck and asks a narrower question than prior compression work: instead of compressing the input or pruning layers, can a model that has seen the full context teach a data-constrained copy of itself how to behave by correcting its internal representations?

Aligning student activations to teacher activations

The method, activation alignment, treats the same frozen model in two roles. A teacher conditions on the full context of size N_full; a student conditions on a stratified random subset of size N_student = floor(alpha * N_full), preserving the label distribution. Because the student sees fewer labeled examples, its intermediate representations carry less task information.

A linear aligner f(h) = Wh + b is fitted to map the student's layer-k activations toward the teacher's. Rather than predicting teacher activations directly, it predicts the residual between teacher and student activations, and its weights are initialized near zero so that corrections start near zero and the student's original representations are preserved at the outset. At inference the corrected activation is the student activation plus the aligner's output, applied at a single transformer layer with the backbone left frozen.

Training the aligner without labels or GPUs

The aligner is trained on synthetic unlabeled query vectors generated from the training features using TabPFN's unsupervised generative capabilities, sampling from the learned joint multivariate distribution so that feature correlations are captured; the experiments use 1,000 synthetic samples. The same TabPFN-generated queries are used for TabFM, since the aligner probes the representation space rather than depending on a model-specific input distribution.

Training minimizes mean squared error between the aligner's predicted residual and the actual teacher-minus-student activation difference, using AdamW with a learning rate of 1e-3 for TabPFN and 1e-4 for TabFM, plus early stopping with patience of 10 epochs on an 80/20 split of activation pairs. For both models, which each comprise 24 layers, activations are extracted and aligned at the final layer. TabPFN uses an 8-estimator ensemble with a separate aligner trained independently per estimator, while TabFM uses a single estimator. The student inherits the teacher's fitted preprocessors so that both operate in the same input representation space, avoiding representation drift from differing permutations, category codes, or unstable quantile estimates on small subsets.

The practical consequence is cost: fitting a single linear layer on pre-extracted activations requires no GPU and converges in seconds to minutes on commodity hardware, in contrast to context-compression or layer-pruning alternatives that demand substantial GPU training.

Gains hold across every context budget

Evaluation covers 38 classification datasets from the TabArena benchmark with TabPFN-3 and TabFM, using context fractions alpha from 0.1 to 0.9 and 5 repeats. Statistical significance is assessed with a two-sided paired Wilcoxon signed-rank test across datasets; for each context fraction the test pairs the 5-repeat mean scores of aligned models against the comparator, and the overall test pairs per-dataset mean scores averaged across all context fractions to avoid treating multiple fractions of the same dataset as independent observations.

Against the unaligned baseline student, the aligned TabPFN student wins 82.2% of evaluated (dataset, context fraction) pairs, with gains statistically significant overall (p < 0.0001) and at every individual context fraction. The aligned TabFM student wins 79.8% overall against its baseline (p < 0.0001), also maintaining significance at every evaluated context fraction. The paper reports that across both models and all evaluated context budgets the aligned student outperforms the unaligned baseline in 81% of cases.

Two different measures of what alignment buys

The paper reports two quantities that are easy to conflate. The effective sample fraction E_alpha describes how much baseline context the aligned student behaves like: at alpha = 0.1, the median is 0.235 for TabPFN and 0.200 for TabFM, with a mean near 0.33 across both architectures. In other words, without processing additional in-context tokens at inference, the aligned student reaches accuracy comparable to an unaligned model given two to four times as many labeled examples, and median sample efficiency stays strictly above parity at all evaluated budgets.

The relative gap closed is a different measure: the share of the teacher-student performance deficit that alignment recovers. Across all evaluated context fractions the median gap closed stays above 34% for TabPFN and above 24% for TabFM, and in the lowest-data regime at alpha = 0.1 it reaches a median of 48.3% and 44.7% respectively. A reader comparing these numbers should note that one is expressed relative to baseline context size and the other relative to the teacher's advantage; they answer different questions and are not interchangeable.

Partial-context students beat full-data XGBoost

The paper also compares against XGBoost trained on 100% of the training data, while the aligned student is restricted to 10% to 50% of context examples. Despite that handicap, the aligned student beats full-data XGBoost in 81.1% of conditions for TabPFN and 80.3% for TabFM (p < 0.0001 overall for both), with statistically significant advantages across all context budgets at alpha >= 0.2.

In a three-way ranking of aligned student, baseline student, and full-data XGBoost, the aligned student holds the top average rank across all context budgets: 1.37 for TabPFN and 1.40 for TabFM, versus 2.19 and 2.13 for the baseline student and 2.44 and 2.47 for XGBoost. The authors read this as evidence that combining tabular in-context learning with activation alignment can surpass fully trained gradient boosting even under strict context budgets.

Where alignment does not help

The evaluation is limited to classification tasks, although the aligner operates on intermediate transformer activations with unlabeled synthetic queries and no classification-specific assumptions; the authors expect it to translate to regression but note this has not been empirically verified. Empirical coverage spans two tabular foundation models, TabPFN-3 and TabFM, and generalization to other tabular in-context learning architectures remains untested.

Alignment is also not universally beneficial. Five of the 38 datasets (MIC, SDSS17, credit_card_clients_default, kddcup09_appetency, and polish_companies_bankruptcy) show a negative mean effective-sample gain for both models, and the paper states it does not currently have a characterization separating these cases. A further scope condition is built into the setup: the teacher must have access to the full context, so the method addresses inference-time compute and memory bottlenecks rather than genuinely small datasets where abundant data never existed. The authors suggest extending the framework by conditioning the aligner on summary statistics or domain priors, or training across diverse external datasets where data is plentiful.

The priority is low-data regimes

The decision-relevant pattern in this paper is not that alignment helps everywhere, but where it helps most and what it costs. The largest recoveries occur at the smallest context budgets: at alpha = 0.1 the median gap closed reaches 48.3% for TabPFN and 44.7% for TabFM, while the median gap closed across all budgets stays above 34% and 24% respectively. A reader weighing whether to adopt this should therefore look first at the low-context end of their own workload, where a student restricted to 10% of the context behaves like an unaligned model with two to four times as much data.

The comparison that matters most is against the full-data tree baseline: the aligned student wins roughly 80% of conditions against XGBoost trained on 100% of the data while itself seeing only 10% to 50% of context examples, and it holds the best average rank in the three-way comparison at every context budget. The counterweight is the five datasets with negative mean effective-sample gain, for which the paper offers no separating characterization, so the practical threshold is to check per-dataset behavior rather than assume the aggregate win rate transfers.

Sources used in this answer