TACTICL: Task-Aware Compression of Tabular ICL Models

TACTICL: Task-Aware Compression of Tabular ICL Models

Abstract Summary:

The paper introduces TACTICL, an automated task-aware compression framework for tabular in-context learning (ICL) models. The framework jointly prunes transformer layers and replaces them with lightweight adapters trained on downstream tasks, blending in-context with in-weight learning. The authors study TACTICL on 47 benchmark datasets and show that up to 85% of layers can be substituted without substantial performance drop on a given downstream task. The framework maintains robustness to data shifts, leaving its in-context ability intact. TACTICL provides a robust framework for exploiting depth-wise redundancy of tabular foundation models by combining task-specific adaptation and structured compression.

Introduction and Motivation:

Tabular foundation models (TFMs) have advanced rapidly, with successive generations achieving state-of-the-art performance while growing in size and inference cost. For example, TabPFNv1 comprises 12M parameters and was limited to tables with 1,000 samples and 100 features; TabPFNv2.5 consists of 10M weights and was built for datasets with up to 50,000 data points and 2,000 features; TabICL2 consists of 27M weights and generalizes effectively to million-scale datasets. Inference costs are still half an order of magnitude larger than lightweight tree-based baselines. Prior work has shown that not all layers contribute equally to task performance, hinting at depth-wise redundancy, and pruning layers offers an opportunity to substantially lower inference costs without sacrificing accuracy.

The authors study whether TFMs can be compressed without collapsing them into pure in-weight learners, specifically whether they can simultaneously prune and adapt a TFM for a downstream task, achieving the benefits of both compression and fine-tuning within a single framework and without sacrificing ICL capability. This is non-trivial because prior work shows that the transition from in-context to in-weights learning is sharp and often irreversible, and fine-tuning a model on a fixed dataset can actively suppress its ICL ability due to low data diversity.

Contributions:

1. A unified framework for simultaneous structured pruning and domain adaptation for state-of-the-art tabular ICL models.

2. The first automated pipeline for explicitly combining in-context learning and dataset-specific memorization to decrease computational complexity while maintaining or even improving predictive performance.

3. A quantitative evaluation on a subset of 47 TabArena datasets, demonstrating that TACTICL can retain the performance of TabPFNv2.5 while reducing the compute cost almost linearly with the number of dropped layers.

Problem Setup:

Given a pretrained model M with L layers and a supervised learning downstream task with training data Dtrain, hold-out (search) data Dsearch, and unlabelled test data Dtest, a compression configuration c ∈ 0,1 L is a binary vector indicating which layers are retained (ci=1) or removed (ci=0). The goal is to identify c* which maximizes the generalization performance f(c) for a fixed compression ratio n/L. Alternatively, instead of removing layers, a layer may be substituted with a lightweight adapter (e.g., a shallow MLP) trained on a subset of Dtrain to approximate the original layer's input-output hidden-state mapping.

Methodology:

TACTICL consists of three components:

• S (Layer Selection): A greedy search over possible layer-replacement configurations, guided by AUC/RMSE and prediction stability evaluated on a held-out search set. The method iteratively drops the layer with the least negative impact on empirical generalization performance.

• P (Adapter Pre-Training): For each dropped layer, a small, lightweight model is trained to approximate the functionality of the layer it replaces via supervised learning on the original in- and output embeddings. The adapter architecture is a two-layer MLP with GELU activations and a skip connection. All adapters are pretrained once, independently and in parallel, before the search begins.

• F (Fine-Tuning): Optionally, the compressed model is fine-tuned to align layers and adapters and recover residual performance lost during compression. A few epochs of AdamW update the weights of the adapter models (with all other weights frozen), using a mixture of pretraining data and data from the downstream task.

Experiments and Results:

Section 4.1 (Can Less Be More?): Using TabPFNv2 with L=12 layers, the authors exhaustively evaluate all possible compression configurations. Results show that dropping up to 3 layers consistently yields a non-negative performance change for half of the datasets, confirming redundancy. Substitution results are even more positive, showing that up to 9 layers can be substituted with non-negative performance changes in some settings. Shapley value analysis reveals that layer importance varies greatly between removing and substituting layers as well as across datasets, motivating task-aware compression.

Section 4.2 (Is Greedy Good Enough?): The authors evaluate several proxy metrics for layer importance and find that none of the structural proxies (cosine similarity, CKA similarity, MLP final losses, residual contributions) reliably predict the performance impact of layer removal or substitution. The only exception is AUC measured on the search set, which retains meaningful predictive performance. Greedy search recovers >90% of the oracle AUC for configurations of up to 8 substituted layers on TabPFNv2. For TabPFNv2.5, TACTICL allows replacing up to 15 layers with negligible performance degradation, and continues to outperform distillation up to 17 layers. The drop/substitute-from-last baseline underperforms substantially, confirming that a searched or learned selection of which layers to remove is necessary.

Section 4.3 (Can It Still Learn In-context?): The authors evaluate the compressed model under three perturbations (feature order permutation, label flip, random features) and find that performance remains stable under perturbations. However, cross-dataset evaluation shows negative transfer for most compressed models, indicating that TACTICL indeed adapts the model to a task by pruning away layers not needed for inference on the target task.

Section 4.4 (Do We Need Fine-tuning?): Fine-tuning the compressed model (with all weights frozen except adapter parameters) improves performance and can potentially recover full model performance. This effect is more pronounced at lower compression rates, where improved performance is observed for some datasets.

Section 4.5 (Does TACTICL Speed Up Prediction?): For classification, inference is 1.25×, 1.63×, and 2.83× faster at compression budgets of 25%, 50%, and 85%, with essentially no AUC cost at 25–50% (median ratio ≥ 0.995) and only a 4% drop at 85%. Regression shows similar speed-ups (1.25×, 1.61×, 2.66×) but degrades more sharply, with RMSE increasing by roughly 3%, 16%, and 61%. For larger synthetic datasets, speed-ups may reach 1.4×, 2.0×, and 5.4× at the three budgets.

Conclusion:

TACTICL achieves up to 5.4× inference speedup, with an average of 2.78×, while maintaining competitive performance. The authors find that layer-importance metrics commonly used in the literature fail to identify good compression configurations, and no single pruning configuration is universally optimal. Fine-tuning at low compression rates can surpass the performance of the original, unpruned model, suggesting TACTICL also serves as an efficient mechanism for domain adaptation.

Limitations:

The current framework does not provide a principled criterion for determining the tradeoff between compression level and performance. The empirical evaluation is limited to tasks within the pre-training constraints of TabPFNv2.5.

Improvements to AI Systems:

1. Task-Aware Dynamic Layer Pruning with Adapter Substitution:

The improved AI system can automatically identify and remove redundant transformer layers for a specific downstream task, replacing them with lightweight, task-trained adapters. This reduces inference compute by up to 5.4× (average 2.78×) while maintaining or improving predictive accuracy, enabling deployment of large tabular foundation models on resource-constrained edge devices or real-time applications.

2. Hybrid In-Context + In-Weight Learning:

The system can blend in-context learning (ICL) with dataset-specific memorization without collapsing into a pure in-weight learner. It preserves robustness to data shifts (e.g., feature permutations, label flips, random features) while adapting to the target task, making it suitable for dynamic environments where input distributions may change after deployment.

3. Greedy Search-Based Compression with Empirical Validation:

The system uses a greedy, AUC-guided layer-selection algorithm that outperforms structural proxies (e.g., cosine similarity, CKA) and fixed pruning heuristics. It recovers >90% of the oracle performance for up to 8 substituted layers, ensuring reliable compression without exhaustive search—enabling rapid, automated model optimization for new tasks.

4. Parallel Adapter Pre-Training for Fast Compression:

The system pre-trains all adapters independently and in parallel before the search, reducing the time to compress a model from hours to minutes. This allows on-the-fly customization of a single pretrained model for multiple tasks, each with its own compressed variant, without retraining the base model.

5. Fine-Tuning for Performance Recovery and Domain Adaptation:

The system can optionally fine-tune only the adapter weights (freezing all other parameters) using a mix of pretraining and downstream data. At low compression rates, this recovers or even surpasses the original model’s performance, effectively turning compression into a form of efficient domain adaptation—useful for specialized fields like medical diagnostics or financial forecasting where data is scarce but task-specific.

6. Scalable Inference for Large Tabular Datasets:

For datasets with millions of rows, the system achieves speed-ups up to 5.4× at 85% compression, enabling real-time inference on streaming tabular data (e.g., fraud detection, recommendation systems) where the original model would be too slow.

7. Task-Specific Model Versioning:

The system can generate multiple compressed variants of the same base model, each optimized for a different task (e.g., classification vs. regression, or different domains). This allows a single pretrained model to serve diverse applications efficiently, reducing storage and maintenance costs in multi-tenant AI services.