Rethinking Cross-Modal Fine-Tuning: Optimizing the Interaction Between Feature Alignment and Target Fitting

AI summary

Overview

Research area: Cross-modal transfer learning / fine-tuning of pre-trained foundation models to new data modalities, with a supporting theoretical analysis drawn from optimal transport and generalization-bound theory.

Technical level: Advanced. The paper builds on Wasserstein distances, probabilistic transport plans, conditional entropy, and Lipschitz conditioning of predictors; the algorithm itself is described at a moderate level, but the theory is mathematically dense.

Scope: The paper derives a generalization bound for cross-modal fine-tuning that decomposes target error into feature alignment, a new "feature-label distortion" term, and target fitting, then uses that bound to design a two-stage fine-tuning algorithm (RECRAFT) evaluated on NAS-Bench-360 and PDEBench.

What This Paper Is About

When a model pre-trained on one kind of data (say, text or images) is adapted to a completely different modality (say, protein sequences or PDE simulations), the new data must be mapped into the pre-trained model's representation space. Aligning those representations is not enough on its own: forcing alignment can distort the relationship between features and labels, and uncalibrated alignment plus fine-tuning can hurt target performance. This paper asks how feature alignment and target fitting should interact, answers with a provable generalization bound, and turns that bound into a working algorithm.

Key Contributions

1. A generalization bound for cross-modal fine-tuning. The paper decomposes the generalized target error into four terms: the source task error (a fixed overhead), a feature alignment (FA) term measured as a Wasserstein distance between source and target feature distributions, a feature-label distortion (FLD) term defined as the minimum entropy over valid transport plans mapping source feature-label conditionals to target ones, and a target fitting (TF) term measuring how well the target predictor follows that transport.

2. The concept of feature-label distortion. This is introduced as a way to quantify the complexity of the probabilistic transport map between source and target feature-label predictive distributions under a given target representation. The paper states this is the first generalization bound capturing both pre-trained model quality and the interaction between feature alignment and target fitting.

3. The RECRAFT algorithm ("REthinking CRoss-ModAl Fine-Tuning"), a two-stage procedure derived from the bound. Stage 1 learns the target feature map by minimizing a tractable surrogate for FA plus FLD; Stage 2 learns a target predictor, parameterized as a learnable transport over source labels, by minimizing TF. The two-stage split is adopted because directly minimizing the bound is unstable — changes to the representation simultaneously reshape the transport landscape on which the predictor is optimized, creating a moving target.

4. Empirical evaluation on two cross-modal benchmarks. RECRAFT achieves the best performance on 7 of 8 PDEBench tasks (average rank 1.25) and the lowest prediction error on 8 of 10 NAS-Bench-360 tasks (best average rank of 1.3 among tested baselines).

Main Findings

• Alignment alone is not enough. The paper argues, and its Figure 2 t-SNE visualization illustrates, that naive fine-tuning (NFT) shows no alignment, while minimizing FA alone produces exhaustive alignment — both give suboptimal predictive error. Minimizing FA + FLD instead yields selective alignment and the best performance of the three settings.

• Target error is bounded by source error plus three interaction terms. Formally, err_τ(φ) ≤ err_s(θ) + FA(φ,θ) + E[FLD(u) + TF(u)] over the target feature distribution. The source loss is described as a fixed overhead that is often negligible for foundation models. The paper reports (Appendix B, Fig. 5) that an empirical inspection shows the bound is sufficiently tight.

• PDEBench results. RECRAFT achieves the best average rank of 1.25, with 7 first-place and 1 second-place finishes across the 8 tasks. For comparison, MoNA has average rank 1.875 (3 first, 4 second), PARE 3.000 (1 first, 5 second), ORCA 3.250 (0 first, 1 second), and NFT 5.000. RECRAFT matches the best reported values on Darcy 2D (0.079) and Diffusion-Sorption 1D (1.6E-3), and posts the best values on Advection 1D (0.0078), Burgers 1D (0.0108), Shallow Water 2D (5.4E-3), Diffusion-Reaction 2D (0.817), Diffusion-Reaction 1D (2.8E-3), and Navier-Stokes 1D (0.050). The paper notes MoNA results are quoted from its own paper because code was unavailable.

• NAS-Bench-360 results. RECRAFT attains the lowest prediction error on 8 of 10 tasks and second lowest on 1 task, with the best average rank of 1.3. Baselines compared include hand-designed solutions, general-purpose models without fine-tuning (Perceiver IO), NAS methods without pre-trained knowledge transfer (DASH), and fine-tuning methods (NFT, ORCA, PARE, MoNA). Ablations comparing NFT, FA-only, and RECRAFT (Table 7, Appendix E) are reported to favor RECRAFT on all tasks.

• A practical hyperparameter range. The Lipschitz constant ω, which sets τ_δ, is treated as a hyperparameter chosen as the smallest value that does not degrade source predictive performance on a proxy dataset. The paper reports this optimal value ranges between 0.3 and 0.5 across different source models under the Euclidean metric.

Methodology in Plain English

The researchers start from a formal problem setup: a source model with an embedder θ and prediction map p_s, and a target dataset from a different modality with its own embedder φ and predictor p_τ. They define generalized source and target errors, and feature distributions as the push-forwards of input distributions under each embedder.

They then prove that target error is bounded by the source error plus three measurable quantities: how far apart the source and target feature distributions are (feature alignment), how much entropy is required to transport source label conditionals into target label conditionals (feature-label distortion), and how closely the learned target predictor tracks the entropy-minimizing transport (target fitting).

Because these terms are not directly computable, the authors build surrogates. Feature alignment is approximated by a Wasserstein-1 distance under the Euclidean metric, scaled by ω, where ω is tuned so that the source cross-entropy prediction is roughly ω-Lipschitz. Feature-label distortion is approximated by a conditional entropy H[Z' | Z], which the researchers estimate empirically: for each target point they sample a pseudo source label from p_s under the target feature map φ, count co-occurrences of (z, z'), and compute entropies from those counts — avoiding the need for the inaccessible oracle transport. The two surrogates are summed and minimized to learn φ.

In Stage 2, the target predictor is written as an expectation over a learnable transport Λ_u(z' | z; φ), and the transport parameters are learned by minimizing the negative log-likelihood of the target labels under that predictor. The paper argues that aligning p_τ with the target conditional in turn reduces the gap between the learned and oracle transports, lowering TF.

Experiments follow the ORCA protocol: RoBERTa for 1D tasks and Swin Transformers for 2D tasks, with CoNLL-2003 and CIFAR-10 as proxy datasets; learning rates, epochs, and optimizers are also taken from ORCA.

Why This Matters

Impact on research. Prior cross-modal methods such as ORCA, PARE, and MoNA combined alignment and fine-tuning heuristically or through bi-level designs without a theoretical link to target generalization. This paper supplies that link, identifies feature-label distortion as a previously overlooked quantity, and shows that measuring it changes which representations are preferred. That reframes alignment from "match the distributions as closely as possible" to "match them selectively, in a way that preserves feature-label semantics."

Real-world applications (drawn from domains the paper cites):

• Genomics: using gene expression profiles to enrich representations for tissue image data.
• Protein structure prediction: transferring vision or language models to protein data.
• Scientific simulation: adapting models to PDE-derived data across families such as Darcy flow, Navier-Stokes, advection, Burgers, and diffusion-reaction.
• Signal and sensor domains: cosmic ray signals, human gestures, ECG, and audio (FSD50K).

Industry relevance. Teams that want to reuse an existing large pre-trained model rather than train a new one for a niche modality stand to benefit most, especially where labeled target data is scarce — the paper explicitly notes that large feature-label distortion causes target fitting to overfit when fine-tuning data is limited. The method requires only a proxy dataset from the source modality, not the original pre-training data.

Future Directions

• Closing the optimization gap. The authors decompose the bound into two stages specifically because joint minimization is unstable, with the representation continually moving the predictor's optimization target. Whether a stable joint or alternating scheme could outperform the two-stage surrogate is left open.

• Estimating the oracle transport more precisely. Feature-label distortion is surrogated by an empirical conditional entropy computed from pseudo source labels. The paper does not report how sensitive RECRAFT is to the quality of this estimate.

• Reducing dependence on the ω hyperparameter. ω must be tuned between 0.3 and 0.5 per source model using a proxy dataset; a more automatic way to select it is not reported.

• Broadening evaluation. Results are reported on NAS-Bench-360 and PDEBench, and MoNA's numbers are quoted rather than reproduced because its code is unavailable. Independent replication of the comparison across more modality pairs is not reported.

Target Audience

Machine learning researchers working on transfer learning, domain adaptation, and foundation-model adaptation to new modalities; theoretically inclined readers interested in generalization bounds built on optimal transport and information-theoretic quantities; and practitioners in scientific computing, genomics, or signal processing who need to fine-tune pre-trained models onto data types the original model never saw, particularly when target labels are scarce.

Authors’ abstract

Adapting pre-trained models to unseen feature modalities has become increasingly important due to the growing need for cross-disciplinary knowledge integration. A key challenge here is how to align the representation of new modalities with the most relevant parts of the pre-trained model's representation space to enable accurate knowledge transfer. This requires combining feature alignment with target fine-tuning, but uncalibrated combinations can exacerbate misalignment between the source and target feature-label structures and reduce target generalization. Existing work, however, lacks a theoretical understanding of this critical interaction between feature alignment and target fitting. To bridge this gap, we develop a principled framework that establishes a provable generalization bound on the target error, which explains the interaction between feature alignment and target fitting through a novel concept of feature-label distortion. This bound offers actionable insights into how this interaction should be optimized for practical algorithm design. The resulting approach achieves significantly improved performance over state-of-the-art methods across a wide range of benchmark datasets.