Revisiting Weak-to-Strong Generalization in Theory and Practice: Reverse KL vs. Forward KL

TL;DR

This paper explores the use of reverse KL divergence for weak-to-strong generalization in language models, showing it outperforms forward KL in theory and practice by focusing on high-confidence predictions.

Contribution

It introduces a theoretically grounded approach replacing forward KL with reverse KL divergence, providing tighter bounds and practical improvements in model training.

Findings

Reverse KL guarantees outperform weak supervisors when fine-tuning

Reverse KL and reverse cross-entropy outperform forward KL in experiments

Theoretically, reverse KL offers comparable or better guarantees than forward KL

Abstract

As large language models advance toward superhuman performance, ensuring their alignment with human values and abilities grows increasingly complex. Weak-to-strong generalization offers a promising approach by leveraging predictions from weaker models to guide stronger systems, but its effectiveness could be constrained by the inherent noise and inaccuracies in these weak predictions. To address this, we propose a theoretically grounded approach that replaces forward KL divergence-whose mass-covering behavior risks overfitting to imperfect weak signals-with reverse KL divergence. Reverse KL divergence's zero-forcing effect prioritizes high-confidence predictions, effectively mitigating the influence of unreliable weak supervision. Theoretically, we extend existing bounds and derive tighter lower bounds for both forward and reverse KL divergence, establishing that reverse KL achieves at least comparable guarantees to forward KL. Notably, when a sufficiently pre-trained strong model is fine-tuned on the last linear layer, reverse KL guarantees that it outperforms its weak supervisor by the magnitude of their disagreement. Empirically, we demonstrate that reverse KL and reverse cross-entropy enable strong models to successfully outperform those trained with forward KL and standard cross-entropy across most settings, highlighting the practical advantages of these reverse losses.