Learning the Pareto Frontier of Predictive Models under Distribution Shift

Technical Summary: Learning the Pareto Frontier of Predictive Models under Distribution Shift

Problem Statement
Modern machine learning increasingly relies on reusing pretrained foundation models across downstream tasks. However, these models differ not only in performance but also in accessibility: some are "black-box" (providing only API predictions), while others are "white-box" (allowing access to internal representations). When deploying these models to a target domain with distribution shift, no single strategy—zero-shot reuse, fine-tuning, or direct training on target data—is uniformly optimal. The choice depends on the unknown nature of the source-target shift. Analysts face a dilemma: zero-shot models may be convenient but miscalibrated; fine-tuned models may overfit with limited target data; and directly trained models may lack efficiency. Current approaches typically force a choice between these mutually exclusive strategies, potentially discarding complementary information.

Methodology: Frontier Learning
The authors propose Frontier Learning, a framework that treats a library of candidate models (spanning different training histories and access regimes) as complementary sources of information rather than alternatives.

1. Unified Representation Construction:

   • The framework constructs a unified target-domain feature vector, $\Psi(x)$, by concatenating signals from all candidate models.
   • For white-box candidates, it extracts internal representations (e.g., bottleneck or penultimate-layer features, $r_m(x)$).
   • For black-box candidates, it uses prediction outputs ($q_m(x)$), such as logits or probabilities.
   • The resulting frontier representation is $\Psi(x) = [\psi_1(x); \dots; \psi_M(x)]$, where $\psi_m$ is either a representation or a prediction.

2. Target-Supervised Learner:

   • A lightweight, regularized supervised learner $g_\theta$ is trained on the concatenated frontier representation using labeled target-domain data.
   • For regression, a regularized linear model is used; for classification, regularized logistic or multinomial logistic regression.
   • The objective is to minimize regularized empirical risk: $\hat{\theta} = \arg\min_\theta \left[ \frac{1}{|I_{tr}|} \sum \ell(y_i, g_\theta(\Psi(x_i))) + \lambda \Omega(\theta) \right]$.
   • Crucially, the candidate models remain fixed; only the final combination weights are learned.

3. Theoretical Property:

   • The hypothesis class of Frontier Learning contains standard strategies (zero-shot, fine-tuning, direct training) as special cases.
   • Consequently, the method is guaranteed to achieve an empirical risk on the training sample no worse than any individual baseline candidate. This provides a structural basis for cautious model composition without unconditional guarantees of lower test risk.

Key Contributions

1. Formulation: The paper formulates prediction under distribution shift as a "model-frontier" learning problem, motivated by the heterogeneous ecosystem of foundation models (varying in training regime and access mode).
2. Unified Framework: It introduces a method that simultaneously combines internal representations from accessible (white-box) models and prediction outputs from inaccessible (black-box) models.
3. Statistical Guarantee: It establishes a "no-worse" empirical risk property via hypothesis-class containment, demonstrating that standard strategies are recoverable as limiting cases.
4. Empirical Evaluation: The framework is evaluated in simulations and two real-world settings, demonstrating robustness across varying degrees of source-target compatibility.

Experimental Results
The authors evaluated Frontier Learning in three settings:

• Simulations: In a controlled regression experiment varying the number of shared feature dimensions ($d_{share}$) between source and target:

  • No single baseline (zero-shot, fine-tuning, or direct) was uniformly best across all shift regimes.
  • Frontier Learning consistently achieved the lowest Mean Squared Error (MSE) across all settings.
  • Improvements were most significant when individual candidates were weak or captured complementary information (e.g., ~89% MSE reduction when $d_{share}=0$).

• Visual Domain Adaptation (DomainNet/VisDA):

  • Task: Classifying images in the "Painting" domain using models trained on "Real," "Clipart," "Sketch," and direct "Painting" training.
  • Setup: Combined bottleneck representations from white-box models (Direct, Fine-tuned Real, Fine-tuned Clipart) with prediction outputs from a black-box zero-shot model (Sketch).
  • Result: Frontier Learning achieved 69.55% top-1 accuracy, outperforming the strongest individual baseline (Direct Painting at 68.17%) and significantly outperforming the weak zero-shot candidate (31.76%).

• Clinical Mortality Prediction (MIMIC-IV-Note):

  • Task: Predicting in-hospital mortality for the CVICU domain using discharge summaries.
  • Setup: Combined note-level representations from fine-tuned (MICU), zero-shot (SICU, CCU), and directly trained (CVICU) models. All were treated as white-box for representation extraction.
  • Result: Frontier Learning achieved the highest AUPRC (0.6638), improving upon the best individual candidate (Fine-tuning MICU at 0.6494) by approximately 2.2%. It also achieved the highest accuracy and F1 score.

Significance and Claims
The paper claims that in an era of proliferating black-box and white-box foundation models, Frontier Learning offers a practical, statistically grounded alternative to committing to a single pretrained model.

• Robustness: By learning how to combine candidate representations and predictions based on target labels, the framework adapts to the specific shift regime rather than relying on a pre-selected strategy.
• Efficiency: The method adds minimal computational overhead, as it only trains a lightweight learner on fixed candidate outputs.
• Caution: The authors emphasize that the method provides a "no-worse" guarantee on training risk and a structural basis for cautious composition, but it does not unconditionally guarantee lower test risk without proper regularization and validation.
• Paradigm Shift: The work suggests a shift from selecting a single "best" model to learning over a heterogeneous frontier of candidates, enabling more reliable reuse of large-scale AI models under distribution shift.

The authors note limitations, including the dependency on the diversity of the candidate library and the use of simple feature concatenation, suggesting that more sophisticated fusion methods could be a direction for future work.