Task Matrices: Linear Maps for Cross-Model Finetuning Transfer

AI summary

Overview

Research area: Machine learning — parameter-efficient model adaptation, transfer learning, and the interpretability of transformer representation spaces.

Technical level: Advanced. The paper assumes familiarity with transformer internals (hidden states, attention and feed-forward blocks, classification heads) and with least-squares regression over matrix-valued targets.

Scope: The paper proposes and empirically tests "task matrices" — learned linear maps from a pretrained model's intermediate-layer representation to a finetuned model's final-layer representation — across vision and text transformers.

What This Paper Is About

Practitioners who want to adapt a foundation model to a downstream task usually finetune it, which is expensive, or train a linear probe, which is cheap but weaker. The paper asks whether the representational change caused by finetuning is approximately linear, and if so, whether a single learned matrix can recover much of the finetuned model's performance using only the base model plus a small amount of data. It tests this linearity assumption on both vision models (CLIP ViT B-32, DeiT, DINOv3) and text models (RoBERTa, all-MiniLM-L12-v2).

Key Contributions

1. Introduces the task matrix, defined as an N_embed × N_embed linear transformation from a base model representation to a finetuned representation, constructed by least-squares regression over pairs of base and finetuned hidden states from a dataset D.

2. Demonstrates cross-model linear structure empirically for vision and text transformer architectures, showing that base models augmented with task matrices surpass linear probes and sometimes approach finetuned accuracy.

3. Shows data efficiency and multi-task generalization, reporting that task matrices can be built from as little as 1% of training data (and in some cases from fewer than 10 samples), and that a single joint task matrix can serve multiple datasets with accuracy dropping from 92% to 81% when going from 1 to 8 classification datasets.

4. Provides ablations and a public implementation, showing that a finetuned classifier head reading directly from the base model performs worse than a task matrix on all datasets and settings, and that a frozen, randomly initialized decoder head still yields strong results.

Main Findings

• Task matrices beat linear probes on all seven text datasets. For RoBERTa-large, task matrices outperformed linear probes across Emotion, HANS, BLiMP, Trec-6, SNLI, ATIS, and Banking77. The abstract states task matrices "beat linear probes by as much as 80% on challenging multi-class datasets."

• Selected RoBERTa results (n=5, 95% CI). Emotion: linear probe 58.9 ± 1.5, task matrix 66.0 ± 2.8, finetuned 91.4 ± 0.8. HANS: 81.3 ± 0.8 vs 96.8 ± 0.2 vs 100.0 ± 0.0. BLiMP: 66.3 ± 1.5 vs 75.3 ± 1.0 vs 83.0 ± 1.0. Trec-6: 79.8 ± 1.5 vs 84.9 ± 1.2 vs 95.1 ± 1.6. SNLI: 71.2 ± 0.5 vs 76.3 ± 1.8 vs 88.7 ± 0.6. ATIS: 89.3 ± 0.2 vs 95.5 ± 0.3 vs 97.8 ± 0.3. Banking77: 64.3 ± 3.0 vs 83.2 ± 1.4 vs 92.0 ± 0.7.

• Vision results are close to finetuned accuracy but not uniformly better than probes. On CLIP ViT-B/32, task matrices beat linear probes on EuroSAT (96.2 vs 95.9), GTSRB (87.2 vs 86.8), MNIST (99.03 vs 98.7), SUN397 (74.8 vs 73.8), and SVHN (66.7 vs 66.6), while probes were ahead on DTD (77.2 vs 75.7), RESISC45 (91.7 vs 89.1), and Stanford Cars (79.9 vs 79.7). Finetuned baselines were higher on every vision dataset, with the largest gaps on GTSRB (98.7 vs 87.2) and SVHN (96.4 vs 66.7).

• Layer position behaves differently in vision and text. Vision task matrices show an upward trend in performance across layers, indicating increasing decodability in later layers. Text task matrices often perform best at intermediate layers — for example layer 18 for RoBERTa-SNLI and layer 7 for RoBERTa-BLiMP — and SNLI performance decreases in later layers. The highest-class-count vision datasets, SUN397 (397 classes) and Stanford Cars (196 classes), showed low decodability in early and middle layers that rose rapidly only at late layers.

• Strong data-scarce results. With training data limited to a 20% split, task matrices showed an 82% improvement on ATIS and 81% improvement on Trec-6 over linear probes. RoBERTa ATIS: probe 42.7 ± 2.4, task matrix 77.7 ± 1.9, finetuned 91.9 ± 1.6. Banking77: probe 14.1 ± 1.5, task matrix 46.5 ± 1.1, finetuned 79.0 ± 1.1. On CLIP at 20% training data, results stayed close to the full-training results from the main vision table.

• Multi-task matrices retain most accuracy. A single joint task matrix across the vision datasets dropped from 92% to 81% accuracy as the number of classification datasets rose from 1 to 8.

• Ablations confirm the transformation matters. Using a finetuned classifier head to read directly from the base model performed worse than task matrices on all datasets across all settings (for example, RoBERTa full train on Banking77: 1.7 ± 0.2, and on BLiMP: 2.4 ± 0.8). With a randomly initialized frozen decoder head, vision task matrices still closely matched earlier results — for example MNIST 99 ± 0.08 and GTSRB 86.9 ± 0.4 — establishing that the task matrix alone can be sufficient.

• Model dependence. On all-MiniLM-L12-v2, task matrices beat probes on HANS (82.3 vs 76.0), BLiMP (50.0 vs 38.1), Trec-6 (84.7 vs 75.1), and SNLI (64.5 vs 56.7), but probes were better on Emotion (66.8 vs 63.5), ATIS (94.9 vs 91.9), and Banking77 (90.9 vs 88.3). On DINOv3 ViT-B/16, task matrices were generally lower than linear probes; the authors attribute this to the strength of the pretrained backbone and the lack of middle-layer enrichment in vision models.

• Double descent appears in the regression. Appendix A reports double descent: accuracy rises with the number of images used, then declines sharply near the full embedding dimension (768 in CLIP), where a unique exact solution exists when the number of input samples equals the embedding dimension. Ridge regression was tested and did not improve on the least-squares baseline.

• Full CLIP with a frozen text encoder. With the vision model replaced by models from Tang et al. (2024) and the text encoder frozen (n=1), task matrices reached 96.51 on EuroSAT, 97.34 on MNIST, and 84.82 on GTSRB, against finetuned values of 98.66, 99.32, and 97.69 respectively.

Methodology in Plain English

The researchers take a model that has been finetuned on a specific dataset and a copy of the same model that has not been finetuned. For a set of training examples, they record two vectors: the hidden representation at a chosen intermediate layer of the base model, and the final-layer representation of the finetuned model. Both are extracted from the [CLS] token, which serves as the unified representation for vision and text.

They then solve a least-squares regression that finds the matrix W that best maps base representations to finetuned representations. At test time, they take the base model's intermediate representation for a new input, multiply it by W, and feed the result into the finetuned model's classification head (or, in one ablation, a frozen randomly initialized head). No weights of either backbone are changed.

The comparison baselines are a linear probe (retraining only the decoder head with cross-entropy loss on the base model's final-layer embeddings) and the fully finetuned model. Finetuning and linear probes were trained until no further improvement, typically many more epochs than the task matrix, which learns from at most one iteration of the training data. Vision experiments used the CLIP ViT B-32 vision tower finetuned end-to-end on the vision component, plus DeiT-tiny-patch16-224 and DINOv3 ViT-B/16. Text experiments used RoBERTa-large and all-MiniLM-L12-v2. Results are mostly averaged over five independent runs with 95% confidence intervals. Datasets were filtered to those with substantial performance gaps between base and finetuned models.

Why This Matters

The work suggests that the change finetuning makes to a transformer's internal state can be captured by a single linear map, which reframes domain adaptation as a representational alignment problem rather than a weight-update problem. This connects interpretability findings about linear structure inside models to a practical adaptation method. The abstract notes the approach is useful "when storing or releasing finetuned models is impractical or commercially infeasible."

Real-world applications implied by the paper's framing and results:

• Ship a base model plus small matrices instead of many full finetuned checkpoints, when storage, distribution, or licensing of finetuned weights is a barrier.
• Domain-specific classification with scarce labels, since task matrices showed the largest relative gains over probes in the 20%-data setting (for example ATIS and Trec-6).
• Multi-task serving from one adapter, since a single joint task matrix retained 81% average accuracy across 8 vision datasets.
• Vision deployment scenarios such as traffic sign recognition (GTSRB), satellite imagery (EuroSAT, RESISC45), texture and scene recognition (DTD, SUN397), and digit recognition (MNIST, SVHN), which are the datasets tested.

Industry relevance: The method's cost profile is attractive — training a task matrix is a least-squares fit over a d × d matrix, comparable in runtime to fitting a linear probe, and construction worked with as little as 1% of training data. The paper's own caveats matter for deployment: gains over probes are not uniform on vision tasks, the approach underperformed probes on DINOv3 ViT-B/16, and some vision datasets (GTSRB, SVHN) show large gaps to the finetuned model.

Future Directions

• Better estimators than least squares. The authors report experimenting with Ridge regression without improvement and explicitly name approximation techniques from Hernandez et al. (2023) as a promising avenue, along with noting that W* can be learned from very few samples in some cases.
• Understanding the double-descent behavior. The sharp accuracy decline near the embedding dimension (768 in CLIP) is described but not resolved; characterizing when full-data construction is preferable to intermediate sample counts remains open.
• Explaining architecture dependence. Why task matrices beat probes on RoBERTa-large and all-MiniLM-L12-v2 for some datasets but trail probes on DINOv3 ViT-B/16 is attributed to backbone strength and the absence of middle-layer enrichment in vision models — a hypothesis rather than a demonstrated mechanism.
• Extending beyond the tested backbone families. The paper evaluates CLIP ViT-B/32, DeiT-tiny-patch16-224, DINOv3 ViT-B/16, RoBERTa-large, and all-MiniLM-L12-v2; generalization to other architectures and to larger-scale multimodal settings is not established.

Target Audience

Researchers and engineers working on parameter-efficient adaptation, transfer learning, or mechanistic interpretability of transformers, particularly those interested in replacing weight-level finetuning with representation-level interventions. The paper is also relevant to practitioners who need many task-specific variants of a foundation model but face storage, compute, or redistribution constraints. Readers without a background in transformer hidden states and linear algebra will find the approach difficult to follow from the paper alone; the public repository at https://github.com/DarrinOBrien/Algoverse-AI-ADDP supports replication.

Authors’ abstract

Results in interpretability suggest that large vision and language models learn implicit linear encodings when models are biased by in-context prompting. However, the existence of similar linear representations in more general adaptation regimes has not yet been demonstrated. In this work, we develop the concept of a task matrix, a linear transformation from a base to finetuned embedding state. We demonstrate that for vision and text models and ten different datasets, a base model augmented with a task matrix achieves results surpassing linear probes, sometimes approaching finetuned levels. Our results validate the existence of cross-layer linear encodings between pretrained and finetuned architectures. Moreover, we show that a data-based approximation for such encodings is both efficient and generalizable to multiple domains. We make our implementation publicly available.