Self-Supervised Alignment with Mutual Information: Learning to Follow Principles without Preference Labels

Thank you for your positive evaluation of our work!

As mentioned in our responses to our first reviewer LYCL, SAMI, like other alignment methods such as DPO, suffers from over-optimization (e.g., non-coherent outputs if trained for too long). We regularize against this “forgetting” by always starting from the initial model during finetuning, using data generated from an intermediate model. While this is a limitation, we do not add additional complexity through our regularization. In fact, by not having a reference model + KL divergence, we reduce algorithmic complexity as we only need to load one model into memory during training. We will make sure to better clarify this relationship and flag the over-optimization issue more clearly in our limitations.

Regarding your other questions and concerns:

• We agree that evaluation on more complex domains is an important limitation and plan to address this in future work.

• Regarding length bias concerns: While HH-RLHF suffered from length bias, our results on TL;DR have shown that length bias can be regularized against simply by stating that responses should be concise. Fig. 3 shows that sequence lengths decreased over iterations as a result of including a conciseness principle in the constitution. As such, while the original objective suffers from a length bias, this can be regularized against by including conciseness as a part of the constitution. Moreover, we have now included updated results with a more principled length correction for our HH-RLHF experiments (see response to reviewer eqUp).

• As mentioned in our response to reviewer LYCL, we only need a small amount of data at each iteration, which is why the computational/resource requirements for training a model are rather low. We will include ablations with respect to principles and potential simplifications of our objective in our revision.

• Regarding regularization: Please see our response to LYCL.

• We have not yet extended to multi-turn interactions; however, related work has looked at multi-turn interactions using self-improvement techniques and shown promising results (Andukuri et al., 2024). In future extensions, we plan to combine ideas from this multi-turn setting with our mutual information objective.

References

• Zelikman, E., Wu, Y., Mu, J., & Goodman, N. (2022). Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35, 15476-15488.
• Andukuri, C., Fränken, J. P., Gerstenberg, T., & Goodman, N. D. (2024). Star-gate: Teaching language models to ask clarifying questions. Conference o

Thank you for your positive evaluation of our work. We completely agree that comparing to other baselines than instruct-finetuned models and base models directly is an important limitation of our work and we plan to address this in future extensions.

Regarding your specific questions:

• We agree that re-running experiments using simplified versions of our contrastive loss (e.g., one-sided vs. two-sided) are important ablations. We will include these in our revision.
• We agree that our first two experiments with HH-RLHF and TL;DR only involved a small number of carefully curated examples. As such, our third experiment with Llama3-70B involves a larger selection of diverse principles (e.g., talk like a pirate or use emojis; see Fig. 5 and Section A.13). We provide examples in Section A.9 (see e.g., at the bottom of p. 23)
• As mentioned in our responses to other reviewers, we are planning to evaluate additional domains requiring more diverse principles, such as roleplaying personas directly (e.g., in MT-bench) in future versions.
• Regarding your comment "I don't understand 'by using a small number of gradient updates during earlier iterations' in line 145. How does this regularize distribution shift?": Thank you for pointing this out. We will revise this statement to be more precise in our revision. Specifically, we regularize in two ways: First, we always train the base model (i.e., the initial model) using data generated from each intermediate model. As such, the same model is never trained twice, and an intermediate model is always using new data for training the next iteration of a model. Second, by only taking a small number of gradient steps, we stay as close as possible to the initial model to avoid forgetting (see also our third response to reviewer LYCL).
• Regarding checkpoint selection: Following previous work (Zelikman et al., 2022), we start with a small number of examples during the first iteration and linearly increase the number of training examples at each iteration. We fix both the number of iterations and the number of examples in advance (see also our third response to reviewer LYCL). We chose 3 iterations simply because we have to run GPT-win rates evaluations to get win rates after each iteration and need to generate new data, so each iteration is resource-intensive. Since related works (e.g., Andukuri et al., 2024) have observed a ceiling after or around three iterations, we followed this approach. We will make sure to point this limitation out more carefully in our revision.

References:

• Zelikman, E., Wu, Y., Mu, J., & Goodman, N. (2022). Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35, 15476-15488.
• Andukuri, C., Fränken, J. P., Gerstenberg, T., & Goodman, N. D. (2024). Star-gate: Teaching language models to ask clarifying questions. Conference on Language Modeling Research.

Thank you for your positive evaluation of our work and providing the additional length correction reference. We will make sure to report both statistical significance in addition to confidence intervals in our revision. Moreover, we will revise figures for clarity as requested.

Please see our attached pdf for updated win rates from Experiment 1 (HH-RLHF Dialogue) based on fitting a logistic regression model to remove length effects. For fitting the logistic model, we have followed the most recent evaluation standard from Length-Controlled Alpaca Evals (Dubois et al., 2024) which similar to Stiennon et al. (2022), inputs the length of each response as well as the result and trains a classifier to predict the counterfactual: “what would the preference be if the model’s output had the same length as the baseline”. Results for HH-RLHF after applying this length correction show the same pattern as our previous results albeit being more conservative. Thank you again for pointing this out, we will make sure to update our figures and text to reflect this change!

Regarding your other concerns and specific questions:

• We take only one gradient step on each batch and never train twice on a given data point within a given iteration. Moreover, we alternate between two dataset splits (A, B) between iterations, ensuring that if a model was trained on split A it never sees data points from split A when generating new training data for the next iteration, for which we use split B (and vice versa if a model was trained on split A). As such, at every point in our pipeline (data generation, training, and win-rate evaluation), a data point encountered by a given model is seen for the first time and is never seen again. We then fix the number of iterations across experiments and compute win rates at the end of each iteration to evaluate performance (similar to Zelikman et al., 2022).
• While it is necessary to have principles, the principle writer could be both a human or another language model. Moreover, principles could be sampled from a pre-existing dataset. As such, the principle writer is not strictly necessary for our pipeline. However, we would like to emphasize that we were interested in exploring a setting in which a principle writer might be weaker than the student being finetuned. We believe that—similar to using a weak supervisor model to label data for training a strong student (i.e., weak-to-strong generalization; see e.g., Burns et al.; 2023)—this is an important point to make on its own as it is not unlikely that future models might surpass human users in capabilities while still having to follow instructions human users. As you mention in your next comment, Figures 3-4 show that a small principle writer can indeed be used to align a stronger student, which we believe is a key finding suggesting that small aligned principle writers / models (which act as a stand-in for a human user) can be used to steer strong students. Thank you for pointing this out again, we will make sure to state this more clearly in our results!
• We agree that additional ablations for the usefulness of antitheses are important and will include these future extensions of our work.
• We agree that we need to make our goal—steerability of a language model via constitutions—more explicit. This goal is distinct from aligning a model to a specific constitution (or a specific distribution of labels through RLHF). Instead, it aims to increase steerability more generally.
• Yes, the inputs to SAMI and baselines are the same. Both the original base model and SAMI-finetuned models use a base model template with no additional special tokens except BOS and EOS tokens. Instruct-model based templates use the exact same input, with the only difference being the additional special tokens required by the tokenizer.

References:

• Burns, C., Izmailov, P., Kirchner, J. H., Baker, B., Gao, L., Aschenbrenner, L., ... & Wu, J. (2023). Weak-to-strong generalization: Eliciting strong capabilities with weak supervision. arXiv preprint arXiv:2312.09390.

• Dubois, Y., Galambosi, B., Liang, P., & Hashimoto, T. B. (2024). Length-controlled alpacaeval: A simple way to debias automatic evaluators. arXiv preprint arXiv:2404.04475.

• Zelikman, E., Wu, Y., Mu, J., & Goodman, N. (2022). Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35, 15476-15488.

Thank you for your positive evaluation of our work!

We completely agree that future work should include a wider range of tasks. Widening both the range of tasks and principles is crucial for training a general constitution-following model, and we plan to address this limitation in future extensions of our work. We will ensure that this limitation is carefully addressed in our limitations section.

Regarding your questions:

• SAMI's computational cost is relatively low. We apply only one gradient step per batch, with a batch size of 128. At each iteration, we train on at most 12 batches (1,536 examples). Following Zelikman et al. (2022), we start with a small number of batches and increase the number of batches by a constant factor in later iterations. Consequently, fine-tuning small models like Mistral-7B or Mixtral-8x7B takes no more than 5-10 minutes at a given iteration (using reasonable hardware such as A100 GPUs). While fine-tuning LLaMA-70B is more expensive due to its size, the dataset still remains small (≤1,536 examples at each iteration; see section A.3 for further details). For comparison, the original DPO paper used 170,000 HH-RLHF examples (see p. 7 in Rafailov et al., 2023), an additional SFT stage prior to DPO fine-tuning, and trained for multiple epochs. Importantly, we view SAMI as a complementary approach to other alignment fine-tuning methods, not a replacement. SAMI's goal is to enhance a model's steerability by amplifying the connection between a set of guiding principles and the responses that realize them, and it can in principle be applied at later stages in post-training, e.g., after DPO or instruction finetuning.

• As mentioned above, we have not yet evaluated SAMI on other domains. This is an important limitation that we will carefully address in our limitations section. An important aspect to consider is that training on summarization and evaluating on mathematical reasoning is unlikely to benefit a model. Instead, we anticipate that training on a variety of domains should help the model generalize. For example, jointly training on code generation, mathematical reasoning, summarization, and other domains should increase generalization performance. We note that this limitation is not specific to SAMI but a data limitation that should apply to finetuning/alignment methods more generally.

• To clarify, we regularize by training the initial (i.e., original base) model at each iteration, using data generated from each intermediate model. This approach is standard (see Zelikman et al., 2023; Andukuri et al., 2024) and prevents overfitting. We only train on a small number of examples using a low learning rate to avoid substantial changes to the initial model. The reason for using regularization is that, as with other alignment approaches like RLHF and DPO which require a reference model and KL divergence, a model might exploit the objective to obtain more reward. In our case, this could mean pushing the log probabilities to an identity matrix to maximize mutual information. Thus, regularization does not prevent the model's ability to learn new behaviors but instead prevents it from "forgetting" desirable behaviors due to reward overoptimization.

References:

• Zelikman, E., Wu, Y., Mu, J., & Goodman, N. (2022). Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35, 15476-15488.
• Andukuri, C., Fränken, J. P., Gerstenberg, T., & Goodman, N. D. (2024). Star-gate: Teaching language models to ask clarifying questions. Conference on Language Modeling Research.
• Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., & Finn, C. (2023). Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36.