Relating Misfit to Gain in Weak-to-Strong Generalization Beyond the Squared Loss

Thank you so much for reading our paper, and for your comments. We address your concerns ahead:

...recent references, concurrent works

As the reviewer notes, the concurrent works listed were released after the ICML submission (and some fairly recently). Nevertheless, we will be sure to cite these works and include some discussion in our final revision, since they do provide complementary perspectives.

...use of KL divergence in the opposite direction...this deviates from the standard choice...more discussion on this point

Thank you for bringing this up, we want to make sure that the ideas can be understood since they are counterintuitive! Our understanding of the differences between forward and reverse KL is based on this nice blogpost.

Roughly we should think of forward KL as being mass-seeking: It prioritizes learning a distribution that covers all possibilities dictated by the teacher. On the other hand, reverse KL is mode-seeking: it prioritizes learning a distribution that captures the most frequent behavior in the teacher. One can notice this behavior in how the two loss functions handle a student that disagrees with the teacher and predicts 0% probability for an event: forward KL will be $+\infty$ while reverse KL will exclude that event from the loss. Conversely if the teacher predicts 0% probability for an event and the student disagrees, then forward KL will disregard that event, reverse KL will become $+\infty$.

In the context of weak-to-strong generalization (WTSG), it is plausible that the teacher makes errors. Thus, we do not want our student to be mass-seeking. The student should be free to disagree with the teacher. It should be mode-seeking as this is likely where most of the signal from the teacher comes from. We will include this discussion (and possibly visualizations) on the reverse KL in the paper.

The fact that the misfit term in the Pythagorean inequality manifests in the reverse direction, shows that the reverse KL is in fact a better-suited objective function for WTSG in classification—we show that minimizing reverse KL divergence in the WTS training procedure leads to provable and quantifiable WTSG, whereas minimizing forward KL (i.e., the standard cross-entropy) may not necessarily yield these guarantees. Our result thus shows that there is a certain “directionality” to the correlation between misfit and gain when one goes beyond the squared loss. In fact, we also see (see Rebuttal to 93MC) that the reverse KL setup yields better WTSG experimentally when compared to the standard setup---to this extent, we view the use of reverse KL as a feature rather than a limitation.

...realizability assumption

We would like to emphasize that realizability of the ground truth is not a critical feature of Theorems 4.1 and 4.3. In particular, note that in these theorems, the weak model has no assumptions placed on it. Thus, we could just as easily take $g$ in Theorems 4.1/4.3 to be the optimal strong model that approximates (with respect to any metric) the possibly non-realizable target function, and we would obtain a misfit-gain inequality with respect to this $g$ (see also Section 5.5, where we do precisely this). This perspective is in line with that taken in Burns et al (2023) where they care about performance gap recovered (PGR), which in some sense, does away with the realizability assumption.

As you mention, a result similar to Theorem 2 in Charikar et al (2024), which incorporates the "distance" of the target function from the strong class into the bound, would be interesting. This needs a little more care in the KL-divergence setting since we don’t have a triangle inequality. We left this result out of our paper to focus more on our Theorem 4.3, which generalizes the misfit-gain inequality to non-convex classes of strong models. However, in the reference [3] that you mention, this generalization is worked out.

...lower-bounding misfit by gain

Lower-bounding the misfit is tricky because we don’t have a triangle inequality for KL divergence. However, in reference [3], such a lower bound was proven, albeit at the cost of a potentially large constant. Without some regularization assumptions, we suspect it would be hard in general to control this constant.

Thank you so much for reading our paper, and for your comments. We are glad that you like our work! We address your concerns ahead:

...submission is not self-contained

While the high-level geometrical framework of viewing weak-to-strong generalization (WTSG) as “projections onto a convex space” is certainly inspired from the work of Charikar et al. 2024 (as we also generously attribute throughout our paper), we do want to emphasize that our theory in terms of Bregman divergences is a strict generalization of the results in Charikar et al. 2024. As such, their result about regression with the squared loss can be deduced as a special case of our Theorem 4.1 (as also stated in Lines 214-216). We do also set up all the preliminaries about Bregman Divergences and their geometry (Section 3) for deriving Theorem 4.1. Nevertheless, the feedback is well-taken; we will try and include a brief discussion of the Charikar et al. 2024 result in Section 3.1.

...conclusion section

We will make sure to elaborate more on the future directions mentioned in the conclusion section.

Could you provide more intuition or visualization around the dual space projections used in your geometric arguments?

Thank you for bringing this up. It can indeed be difficult to visualize the geometry of KL-divergence since it can behave quite differently from our familiar Euclidean geometry. Frank Nielsen has a very nice article discussing this very topic.

Since we are most familiar with Euclidean geometry, it is best to appeal to those native intuitions when understanding dual space projections. More specifically, our geometric arguments require us to generalize the notion of angles so we can discuss the Pythagorean theorem in a more general setting. Bregman divergence theory gives us such a generalization, but it comes at a critical cost: asymmetry. In Euclidean geometry, we can measure angles between any two vectors by taking an inner-product. In Bregman geometry, we can only measure angles between a primal vector and a dual vector. But once we designate one vector as primal and one as dual, much of our traditional intuition from Euclidean geometry applies.

When technically manipulating these expressions, we like to imagine that we are in Euclidean geometry and have access to inner products. The only restriction is we cannot swap terms in inner products or in squared $L_2$ distances. If proofs go through after taking into account these restrictions, then it is very likely they will hold in the general Bregman setting. We would be happy to include some more intuition and visualization regarding this in our final revision.

Thank you so much for reading our paper, and for your comments. We are glad that you like our work! We address your concerns ahead:

1...additional insights beyond confirming the correlation between misfit and gain.

2..intuitive interpretation for reverse KL...

We address both 1) and 2) in our response to Reviewer 7oRJ. Please also see our response to Reviewer fGrp for some intuition on the geometry induced by Bregman Divergences.

3...impractical aspects...If the experiments were conducted in a more standard manner, would the same trends and conclusions still hold?

We thank the reviewer for prompting us to do this comparison; we believe the obtained results lend more conceptual support to our theory. For the experiments in Section 5.2, we ran a comparison where we perform linear-probing for the strong model in the “standard” manner. That is, in the weak-to-strong training phase, instead of finetuning a convex combination of $k=100$ linear layers on the reverse KL objective, we finetune these on the forward cross-entropy objective XE(weak, strong) (i.e., the natural direction), and compare the findings to the numbers presented in Section 5.2. Note that the only difference between the settings now is the objective that the linear layers are finetuned on. The table below shows the comparison for the test accuracy of the strong model, as well the XE loss between the ground truth and the strong model.

The comparison is indeed quite interesting—we can see that the strong model that is finetuned on the reverse KL objective shows better final test accuracy (i.e., better WTSG) for nearly all the datasets! Namely, we do not really see significant performance degradation (in fact, we see improvement in nearly all cases) with reverse KL, when compared to the standard setup.

Moreover, we also computed the gain and misfit terms in the Pythagorean inequality, where for the reverse KL experiment, we compute the reverse misfit that we propose (i.e., KL(strong, weak), whereas in the standard forward XE experiment, we compute the the misfit as XE(weak, strong) (i.e., the “natural” misfit that one might consider).

We can clearly see that if one were to run the standard WTSG setup, the Pythagorean inequality is completely off (gain and misfit don’t quantitatively align), whereas the reverse misfit is more representative of the gain, as confirmed by our theory as well. Again, this indicates a clear “directionality” in the Pythagorean inequality for WTSG in the classification setting! We would be happy to include this discussion in the final version.

It is plausible that in practice, running standard forward XE(weak, strong) minimization might lead to a minimizer that is in fact close to the minimizer of the reverse KL(strong, weak) that we propose. It would be an interesting future direction to precisely characterize scenarios when this is the case.

Finally, we want to emphasize that in our experiments, we only do weak-to-strong finetuning of the last linear layer. In contrast, the best numbers reported in the original WTSG paper by Burns et al. 2023 correspond to full finetuning (on the forward KL objective) of all the weights of the architecture, and the WTS accuracies are hence naturally better. The primary concern we point out about reverse KL minimization is the non-convexity of the objective, which can possibly cause issues for first-order methods. However, if one is finetuning all the weights in the architecture anyway, this concern is not central, since in such a case, even the standard forward cross-entropy objective becomes non-convex.

Thank you so much for reading our paper, and for your comments. We address your concerns ahead:

...intuition behind key contribution not clearly explained

Thank you for bringing this up, we want to make sure that the ideas can be understood since they are counterintuitive! Our understanding of the differences between forward and reverse KL is based on this nice blogpost. We want to note that, after the ICML deadline, an independent work also highlighted the value of reverse KL divergence.

Roughly we should think of forward KL as being mass-seeking: It prioritizes learning a distribution that covers all possibilities dictated by the teacher. On the other hand, reverse KL is mode-seeking: it prioritizes learning a distribution that captures the most frequent behavior in the teacher. One can notice this behavior in how the two loss functions handle a student that disagrees with the teacher and predicts 0% probability for an event: forward KL will be $+\infty$ while reverse KL will exclude that event from the loss. Conversely if the teacher predicts 0% probability for an event and the student disagrees, then forward KL will disregard that event, reverse KL will become $+\infty$.

In the context of weak-to-strong generalization (WTSG), it is plausible that the teacher makes errors. Thus, we do not want our student to be mass-seeking. The student should be free to disagree with the teacher. It should be mode-seeking as this is likely where most of the signal from the teacher comes from. We will include this discussion (and possibly visualizations) on the reverse KL in the paper.

...strong assumptions like convexity...new insights beyond Charikar et al. 2024...sample complexity, dynamics

While the high-level framework of viewing WTSG as “projections onto a convex space” is certainly inspired from the work of Charikar et al. 2024, we want to highlight that our theory in terms of Bregman divergences is more complete, in that it recovers their results for regression as a special case, and also paves the way for other settings. Furthermore, as described in the paper, several things become more complicated while applying this theory to non-symmetric Bregman divergences. We believe that our technical contributions towards fixing these yield insights that are of interest by themselves.

First, the fact that the misfit term in the Pythagorean inequality manifests in the reverse direction, shows that the reverse KL is a better-suited objective function for WTSG in classification—we show that minimizing reverse KL divergence in the WTS training procedure leads to provable and quantifiable WTSG, whereas minimizing forward KL (i.e., the standard cross-entropy) may not necessarily yield these guarantees. In fact, we also see (see Rebuttal to 93MC) that the reverse KL setup yields better WTSG experimentally when compared to the standard setup. Our result thus shows that there is a certain “directionality” to the correlation between misfit and gain when one goes beyond the squared loss.

Second, our guarantee (Theorem 4.3) for reverse KL minimization, holds even if the function space is non-convex (unlike the work of Charikar et al. 2024). Our solution to handle non-convexity, namely optimizing over a convex combination of $k$ functions so as to approximate the convex hull of the function space—is both practical (as evidenced by the experimental results), and also allows us to derive concrete theoretical guarantees. The error term in the guarantee becomes smaller as $k$ increases, matching the intuition that the convex hull is being approximated better. Importantly, this solution allows us to optimize over convex combinations of functions from an arbitrary (possibly non-convex) function space, and not just linear/logistic heads! This opens the gates for obtaining WTSG guarantees with full-finetuning, and not just linear probing. In particular, the same bound holds even if we finetune all the weights of a convex combination of neural networks—such freedom is not afforded in the theory of Charikar et al. 2024. Our solution requires deriving an insightful connection with existing density estimation methods, and requires non-trivial adaptations to our setting. We view this contribution to be novel, and possibly also as a useful modeling paradigm in other learning settings, especially given that our experimental numbers improve with reasonably large values of $k$.

Finally, sample complexity bounds are a straightforward adaptation from Charikar et al. (2024), and we wanted to keep the paper focused on new contributions. As for learning dynamics, for our initial study, we assume that the $\arg\min$ of the reverse KL can be approximated well (as also suggested by our experiments); nevertheless, this is an interesting future direction.

...what's $\mathcal{F}*$

It is the dual space of $\mathcal{F}$ (see Line 123).