Weak-to-Strong Generalization under Distribution Shifts

TL;DR

This paper addresses weak-to-strong generalization under distribution shifts, where weak annotators trained on different data can mislead a strong model during fine-tuning. It proposes RAVEN, a robust framework that learns adaptive weights for weak models and jointly trains the strong model, augmented by an easy-sample guided initialization to bootstrap learning. Across image and text classification and preference alignment, RAVEN delivers substantial out-of-distribution gains (e.g., ~55–57% PGR improvements and notable accuracy gains) while still performing well in-distribution, and it tends to assign higher weights to more reliable weak annotators without GT access. The work combines theoretical insights and extensive experiments to show that robust ensemble supervision can be automatically leveraged to approach ground-truth performance under distribution shifts, with practical benefits for scalable alignment of high-capability models.

Abstract

As future superhuman models become increasingly complex, accurately supervising their behavior may exceed human capabilities. Recent works have demonstrated that in such scenarios, weak models can effectively supervise strong models, a phenomenon known as weak-to-strong generalization. However, we find that naive weak-to-strong generalization fails under distribution shifts, often leading to worse performance of the strong model than its weak supervisors. To address this, we propose RAVEN, a robust weak-to-strong generalization framework that dynamically learns the optimal combinations of weak models in addition to parameters of the strong model. We demonstrate the effectiveness of RAVEN on image classification, text classification, and preference alignment tasks. RAVEN outperforms alternative baselines by over 30% on out-of-distribution tasks while matching or surpassing existing methods on in-distribution tasks. Moreover, our results show that RAVEN assigns higher weights to more accurate weak models, demonstrating its ability to automatically identify trustworthy supervision.

Figures (13)

• Figure 1: An illustration of weak-to-strong generalization under distribution shifts. Left: The original W2S generalization framework burnsweak. Right: Our extended framework that incorporates distribution shift. In this scenario, the weak model is not only limited in understanding but also unfamiliar with the fine-tuning distribution, making it less reliable for annotating data to supervise the strong model.
• Figure 2: Performance variation with distribution shifts. The Cyan bar indicates positive PGR of InD, gray represents positive PGR for OOD, and red denotes negative PGR. PGR is reported as a percentage. Detailed results can be found in Appendix \ref{['sec:appendix:motivation']}.
• Figure 3: Across different datasets, we observe substantially higher performance variance in OOD scenarios than in InD scenarios. The variance is computed using $20$ weak models that were initialized with different random seeds and trained on datasets from the WILDS benchmark koh2021wilds.
• Figure 4: Left: Correlation between weak model accuracy $Acc(P_{trg}; f_{w}^{\text{src}})$ and its W2S accuracy $Acc(P_{trg}; f_{s}^{\text{pseudo}})$. Right: The weights $\Theta_{W}$ assigned to weak models across two runs of RAVEN.
• Figure 5: Effect of the number of weak models.

Theorems & Definitions (1)

• Remark 1