Closing the Context Gap: Activation Alignment for Tabular In-Context Learning

1 Introduction

• Tabular foundation models’ dependence on context size makes sample efficiency a practical bottleneck, especially when data, memory, or computation is constrained.Their attention cost grows quadratically with context size, and every prediction processes the training examples again.Source: §1 · p. 1
• Activation alignment trains a lightweight linear module to map reduced-context student activations toward those of a full-context teacher.It operates at a single transformer layer without model fine-tuning and can be trained without a GPU in seconds to minutes on commodity hardware.Source: §1 · p. 2
• Across 38 TabArena classification tasks, aligned students outperform unaligned baselines with statistically significant gains in 81% of cases across both models and all evaluated context budgets.With 10% of training context, aligned models achieve parity with unaligned students using two to four times as much data.Source: §1 · p. 2

2.1 Tabular Foundation Models and In-Context Learning

• Tabular foundation models perform in-context prediction in one forward pass, and their performance depends heavily on the number of context examples.TabPFN-3 and TabFM are examples of this approach; TabPFN-3 processes tabular data with a unified token per example, while TabFM uses row-level tokenization.Source: §2.1 · p. 2

2.2 Compressing Tabular Foundation Models

• Prior efficiency methods compress context, optimize soft prompts, prune layers, synthesize smaller networks, or retrieve local contexts, often using substantial GPU training or end-to-end optimization.The aligner instead operates directly on pre-extracted representations and converges on ordinary CPU cores in seconds to minutes.Source: §2.2 · pp. 2–3

2.3 Feature-Based Knowledge Distillation

• Unlike feature-based distillation across architectures, this method transfers information across data regimes within the same frozen architecture using synthetic unlabeled queries.The aligner remains deployed at inference, whereas feature-based distillation typically discards its learned regressor after updating student weights.Source: §2.3 · p. 3

2.4 Activation Steering

• Activation steering supports modifying hidden states during inference; unlike Tuned Lens, this aligner writes corrected representations back into the forward pass to improve predictions.Tuned Lens decodes intermediate states into the model’s final output distribution for interpretability.Source: §2.4 · p. 3

3.1 Problem Formulation

• The teacher uses the full context and the student a stratified subset, while both use the same frozen model; the goal is to recover some predictive quality through activation transformation.The student subset size is Nstudent = ⌊α·Nfull⌋.Source: §3.1 · p. 3

3.2 Pipeline Overview

• The pipeline generates unlabeled queries, extracts teacher and student activations, fits a linear aligner with MSE, then applies it to student activations at inference.Figure 1 depicts synthetic queries passed through the model under both contexts for fitting, followed by applying the trained aligner to the student at test time.Source: §3.2 · p. 4, Figure 1 · p. 4

3.3 Synthetic Data Generation

• Synthetic queries are sampled from a TabPFN-learned joint distribution of training features, capturing their correlations without requiring additional labeled data.The experiments use 1,000 synthetic samples; for TabFM they use the same TabPFN-generated queries to probe the representation space.Source: §3.3 · p. 4

3.4 Aligner Training

• The aligner is a linear map over layer-k activations and predicts a residual correction rather than the teacher activation directly.Near-zero initialization makes its initial corrections small, preserving the student’s original representations at the start of training.Source: §3.4 · pp. 4–5
• For both 24-layer models, activations are aligned at the final layer, with MSE training and early stopping on a held-out activation-pair split.AdamW learning rates are 10^-3 for TabPFN and 10^-4 for TabFM, with patience of 10 epochs.Source: §3.4 · p. 5
• TabPFN uses a separate aligner per ensemble estimator, and students inherit teacher preprocessors to keep their input representations consistent.The TabPFN ensemble has eight estimators; TabFM uses one estimator.Source: §3.4 · p. 5

4 Experimental Setup

• The benchmark contains 38 binary and multiclass classification datasets from TabArena-v0.1, evaluated with TabPFN-3 and TabFM.TabPFN-3 uses an eight-estimator ensemble and TabFM a single estimator.Source: §4 · pp. 5–6
• The primary analysis tests student fractions from 0.1 to 0.5, with five repeated experiments per dataset and ROC-AUC or negative log-loss as appropriate.The full fraction spectrum from 0.1 to 0.9 is also evaluated; higher scores are better for both metrics.Source: §4 · p. 5
• Comparisons include baseline and aligned students, full-context teachers, and full-data XGBoost; effective sample fraction estimates the unaligned fraction needed to match an aligned score.Alignment metrics are evaluated only where the teacher beats the unaligned baseline, excluding five TabPFN slices and two TabFM slices among 190 each.Source: §4 · pp. 5–6

5.2 Sample Efficiency and Distillation Gains

• At α = 0.1, median effective sample fractions are 0.235 (TabPFN) and 0.200 (TabFM); accuracy matches unaligned models given two to four times more examples.These effective-fraction values are reported for the 10% context setting.Source: §5.2 · pp. 6–7

5.3 Comparison with Full-Data XGBoost

• 81.1% (TabPFN) and 80.3% (TabFM) of conditions favor the aligned student over full-data XGBoost; both overall comparisons are significant (p < 0.0001).The advantage remains statistically significant at context fractions α ≥ 0.2.Source: Figure 2 · p. 7, §5.3 · p. 7

5.1 Alignment Substantially Improves Over Baseline Student

• 82.2% (TabPFN) and 79.8% (TabFM) overall win rates favor alignment over the unaligned baseline, with significant gains across every context fraction.Overall gains are significant at p < 0.0001 for both models.Source: §5.1 · p. 6, Figure 2 · p. 7

5.2 Sample Efficiency and Distillation Gains

• Median sample efficiency remains strictly above parity across all evaluated context budgets.Source: §5.2 · p. 7
• At α = 0.1, alignment closes 48.3% (TabPFN) and 44.7% (TabFM) of the teacher–student deficit; medians exceed 34% and 24% across budgets, respectively.Source: §5.2 · p. 7

6.1 Limitations

• Figure 3 reports positive data gains for both models across context fractions and a significant fraction of the teacher–student gap closed.Source: Figure 3 · p. 8
• Evaluation covers classification only; transfer to regression is expected but has not been empirically verified.Source: §6.1 · p. 8
• The empirical evaluation covers TabPFN-3 and TabFM; generalization to a broader family of tabular foundation models remains unvalidated.Source: §6.1 · p. 8
• Five of 38 datasets show negative mean effective-sample gains for both models, and the authors lack a way to identify such datasets in advance.Characterizing which tasks benefit remains an open question.Source: §6.1 · p. 8

6.2 Future Directions

• Future work could test nonlinear aligners, adaptive or multi-layer intervention, hyperparameter optimization, and output-logit distillation objectives.Source: §6.2 · p. 9
• Alignment could be combined with context compression, layer pruning, or retrieval-based context selection, which the study does not compare against.The student evaluated here uses a stratified random subset; correcting representations after retrieval-based selection is proposed as a possible extension.Source: §6.2 · p. 9
• Further mechanistic analysis could characterize representation displacement, explain why linear alignment works, and identify residual nonlinear structure.The results establish that a linear correction recovers a substantial share of the teacher–student gap, but do not characterize its geometry.Source: §6.2 · p. 9
• The current aligner needs offline access to the full dataset, so it does not address settings where only a small dataset exists.Proposed extensions include conditioning on summary statistics or domain priors, or training across external datasets.Source: §6.2 · p. 9

• Activation alignment projects a data-limited student’s intermediate hidden states toward full-context representations at a single layer.The aligner is an affine transformation optimized by MSE on synthetic data, with no external labels; the backbone stays frozen.Source: §7 · p. 10
• Across 38 TabArena classification benchmarks on TabPFN-3 and TabFM, alignment yields broad, statistically significant gains over unaligned baselines.With 10% of training examples as context, aligned students match unaligned models provided two to four times as much data.Source: §7 · p. 10
• Aligned students using partial contexts outperform full-data gradient-boosted trees in over 80% of evaluations.Source: §7 · p. 10
• The findings indicate that simple linear mappings can steer tabular foundation-model representations toward high-data regimes.Source: §7 · p. 10

• The breakdown spans all 38 TabArena classification datasets and context fractions α from 0.1 to 0.9.Source: A Per-Dataset Sample Efficiency… · p. 13
• Table 1 reports per-dataset effective baseline sample fractions Eα and mean gains (Eα − α) for TabPFN and TabFM.Source: Table 1 · p. 14
• Green shading marks Eα > α or positive mean gain, while dashes mark slices excluded by the teacher-superiority filter.Source: Table 1 · p. 14