Alignment of Large Language Models with Constrained Learning

Re: Weakness

Our problem formulation and algorithm are not limited to single-constraint alignment tasks and are readily scalable to multi-constraint settings. We have designed and conducted additional experiments on a multi-constraint task using a new dataset with two distinct constraints (harmlessness and humor). Specifically, we fine-tune a LLaMA2-7B model on the Anthropic HH-RLHF dataset [1] to obtain a reference policy. We then align the model using our method under varying constraint levels. We use the helpfulness and harmlessness reward models from [6] for evaluating helpfulness and harmlessness, respectively, and the humor–no-humor model as the reward model for humor. We demonstrate that our iterative method satisfies the constraints while maximizing the primary utility on this task.

For implementation, similar to our main experiment, we randomly select 1000 prompts from the full HH-RLHF dataset and sample 128 responses per prompt in the dual sub-gradient step. In the policy optimization step, we train the model using DPO with 10K preference pairs randomly selected from the same dataset. We evaluate our trained models on a randomly sampled set of 2000 prompts, each with 2 responses, from the test split of the HH-RLHF dataset.

The table above shows the results for threshold $\mathbf{b}=[0.5,0.5]$, i.e., an expected improvement of 0.5 in both harmlessness and humor scores. We observe that both constraints are satisfied, while the helpfulness score improves by 9.12%.

We are happy to report additional results for other values of $\mathbf{b}$ during the discussion period and will include this multi-constraint task in the final version of our paper.

Regarding the concern about using the 7B-scale model, we note that our theoretical results and algorithm are applicable to LLMs of any size. The 7B-scale models are sufficient to address the questions we focus on, namely constraint satisfaction and trade-offs among properties (see Section 4.2). Moreover, using 7B models is a common practice in the alignment literature [2,3,4,5], and adhering to this scale facilitates fair and direct comparison with prior work.

We appreciate the reviewer's positive evaluation of our work, and believe that the additional experiments and discussions help address the concern regarding the scope of our experiments.

References:

[1] Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback

[2] Safe RLHF: Safe Reinforcement Learning from Human Feedback

[3] One-Shot Safety Alignment for Large Language Models via Optimal Dualization

[4] Stepwise Alignment for Constrained Language Model Policy Optimization

[5] Bounded Rationality for LLMs: Satisficing Alignment at Inference-Time

[6] Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment

Re: Weakness 1 & 2

We agree with the reviewer that alternating Lagrangian maximization and dual updates is a well-known framework in constrained optimization. Rather than proposing a new optimization framework, our contribution lies in demonstrating how to adapt and analyze this framework in the context of LLM alignment, resulting in an iterative, dualization-based alignment strategy that we believe is fundamental to constrained alignment problems. Importantly, the non-convexity of LLM parametrization introduces unique challenges in this development.

Compared to Huang et al. (NeurIPS 2024), our method introduces two key algorithmic innovations. First, while Huang et al. (NeurIPS 2024) uses an exact optimal dual variable, which is hard to evaluate in practice (see it in Huang et al. (NeurIPS 2024)), our method removes this assumption and instead employs a dual-based descent procedure to approximately find an optimal dual variable in a tractable manner. This demonstrates our first algorithmic novelty. Second, motivated by our theoretical result showing that smaller optimal dual variables lead to smaller optimality gaps, we introduce warm-start initialization and post-processing techniques (see Algorithm 1 and Appendix D) to improve convergence performance. This demonstrates our second algorithmic novelty. These practical designs not only clearly generalize the method proposed in Huang et al. (NeurIPS 2024), but also significantly improve constraint satisfaction for practically aligned LLMs (see Figure 2).

To the best of our knowledge, our work is the first to analyze the duality gap in solving constrained alignment problems and to provide theoretical guarantees on the suboptimality of practically aligned LLMs induced by dual iterations. These insights deepen our understanding of the trade-off between constraint satisfaction and performance when addressing constraints on pretrained LLMs. Our work fills a theoretical gap in the constrained alignment literature, including the work of Huang et al. (NeurIPS 2024), which focuses on simple algorithms without offering a duality and optimality analysis.

Re: Weakness 3.1, 3.2 & Question 1

Our problem formulation and algorithm are not limited to single-constraint alignment tasks and are readily scalable to multi-constraint settings. To address the concerns, we have designed and conducted additional experiments on a multi-constraint task using a new dataset with two distinct constraints (harmlessness and humor). Specifically, we fine-tune a LLaMA2-7B model on the Anthropic HH-RLHF dataset [1] to obtain a reference policy. We then align the model using our method under varying constraint levels. We use the helpfulness and harmlessness reward models from [6] for evaluating helpfulness and harmlessness, respectively, and the humor–no-humor model as the reward model for humor. We demonstrate that our iterative method satisfies the constraints while maximizing the primary utility on this task.

For implementation, similar to our main experiment, we randomly select 1000 prompts from the full HH-RLHF dataset and sample 128 responses per prompt in the dual sub-gradient step. In the policy optimization step, we train the model using DPO with 10K preference pairs randomly selected from the same dataset. We evaluate our trained models on a randomly sampled set of 2000 prompts, each with 2 responses, from the test split of the HH-RLHF dataset.

The table above shows the results for threshold $\mathbf{b}=[0.5,0.5]$, i.e., an expected improvement of 0.5 in both harmlessness and humor scores. We observe that both constraints are satisfied, while the helpfulness score improves by 9.12%.

We are happy to report additional results for other values of $\mathbf{b}$ during the discussion period and will include this multi-constraint task in the final version of our paper.

Re: Weakness 3.3 & Question 2

First, we would like to clarify some key differences between the recent constrained alignment techniques and our work. While many of them (such as Stepwise Alignment [4]) study trade-offs among multiple objectives or properties, our primary objective is to impose constraints on LLMs, i.e., to find an optimal feasible model. It is well-known that Pareto optimality in multi-objective optimization is fundamentally different from optimality under constraints. Additionally, the prior multi-objective alignment works typically do not address constraint satisfaction explicitly, as we do in Figure 2a.

To compare trade-offs, we evaluate our method against the beaver-7b-v1.0 model from Safe-RLHF [2], as well as the best reported models, SACPO and P-SACPO, from Stepwise Alignment [4]. For evaluation, we use the GPT-4o-mini model as the judge and the same prompt dataset used in the GPT-4 evaluation from the Safe-RLHF paper. We observe that, at constraint levels $\mathcal{b}=8,9$, our method generally achieves higher win rates across both helpfulness and safety criteria (see the table below).

Re: Question 3, 4

Please refer to Appendix D.2 for a detailed runtime comparisons of our main experiment. The computational cost of our method is comparable to that of the one-shot method, as each iteration requires only a single epoch of DPO. Improvements shown in Figure 2 are typically observed within four iterations. The additional runtime of the multi-shot method only comes from dual variable estimation, which can be efficiently accelerated through engineering techniques.

Regarding the concern about using the 7B-scale model, we note that our theoretical results and algorithm are applicable to LLMs of any size. The 7B-scale models are sufficient to address the questions we focus on, namely constraint satisfaction and trade-offs among properties (see Section 4.2). Moreover, using 7B models is a common practice in the alignment literature [2,3,4,5], and adhering to this scale facilitates fair and direct comparison with prior work. For the robustness and generality aspects, we conduct evaluations using model-based methods (Section 4), GPT-based methods (Appendix E.2), and evaluate on an adversarial dataset (Appendix E.3). We will also include the additional multi-constraint experiment mentioned earlier in the final version of our paper.

We appreciate the reviewer’s comments, and believe that the additional experiments and discussions provide further insight to open your evaluation of our work.

References:

[1] Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback

[2] Safe RLHF: Safe Reinforcement Learning from Human Feedback

[3] One-Shot Safety Alignment for Large Language Models via Optimal Dualization

[4] Stepwise Alignment for Constrained Language Model Policy Optimization

[5] Bounded Rationality for LLMs: Satisficing Alignment at Inference-Time

[6] Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment

Re: Weakness 1

Thank you for the suggestion. We have conducted additional experiments on a multi-constraint task involving two distinct constraints (harmlessness and humor) motivated by the Helpful Assistant task in [1]. Specifically, we fine-tune a LLaMA2-7B model on the Anthropic HH-RLHF dataset [4] to obtain a reference policy. We then align the model using our method under varying constraint levels. We use the helpfulness and harmlessness reward models from [6] for evaluating helpfulness and harmlessness, respectively, and the humor–no-humor model as the reward model for humor. We demonstrate that our iterative method satisfies the constraints while maximizing the primary utility.

For implementation, similar to our main experiment, we randomly select 1000 prompts from the full HH-RLHF dataset and sample 128 responses per prompt in the dual sub-gradient step. In the policy optimization step, we train the model using DPO with 10K preference pairs randomly selected from the same dataset. We evaluate our trained models on a randomly sampled set of 2000 prompts, each with 2 responses, from the test split of the HH-RLHF dataset.

The table above shows the results for threshold $\mathbf{b}=[0.5,0.5]$, i.e., an expected improvement of 0.5 in both harmlessness and humor scores. We observe that both constraints are satisfied, while the helpfulness score improves by 9.12%.

We are happy to report additional results for other values of $\mathbf{b}$ during the discussion period and will include this multi-constraint task in the final version of our paper.

For additional benchmarks related to our main experiment, we have evaluated our method against the beaver-7b-v1.0 model from Safe-RLHF, as well as the best reported models, SACPO and P-SACPO, from [5]. For evaluation, we use the gpt-4o-mini model as the judge and the prompt dataset used in the GPT-4 evaluation from the Safe-RLHF paper. We observe that for constraint levels $\mathbf{b}=8,9$, our method in general achieves higher win rates in both helpfulness and safety criteria (see the table below).

Re: Weakness 2

We would like to clarify two key differences between these papers and our work. First, we have different alignment objectives: the mentioned papers focus on exploring trade-offs among multiple objectives or properties, whereas our primary objective is to impose constraints on LLM models, i.e., to find an optimal feasible LLM model. It is well established that the Pareto optimality of multi-objective optimization is inherently different from optimality under constraints. Additionally, the multi-objective alignment papers typically do not suggest constrained satisfaction as we do in Figure 2(a).

Second, while those works focus on inference-time alignment, our approach addresses offline alignment. Thus, our method and the referenced approaches are orthogonal and each applicable to different tasks and settings.

Last but not least, our method is complementary to inference-time alignment. For instance, our iterative dualization approach can be directly applied to the constrained inference-time alignment to ensure constrained satisfaction, an extension we leave for future work.

We appreciate the reviewer for bringing up these decoding-time approaches. We will incorporate the comparisons in the final version of the paper.

Re: Weakness 3

Our problem formulation and algorithm are not limited to single-constraint alignment and are readily scalable to multi-constraint settings. To demonstrate this, we will include the additional multi-constraint experiment in our response to Weakness 1 in the final version of our paper.

We appreciate the reviewer's positive evaluation of our work, and believe that the additional experiments and discussions help address the concern regarding the evaluation of our method.

References:

[1] Decoding-time language model alignment with multiple objectives.

[2] Bounded Rationality for LLMs: Satisficing Alignment at Inference-Time.

[3] Pad: Personalized alignment of llms at decoding-time.

[4] Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback

[5] Stepwise Alignment for Constrained Language Model Policy Optimization

[6] Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment

Re: Weekness 1

Our problem formulation and algorithm are not limited to single-constraint alignment tasks and are readily extensible to multi-constraint settings. To address the reviewer’s concern, we have designed and conducted additional experiments on a multi-constraint task using a new dataset with two distinct constraints (harmlessness and humor). Specifically, we fine-tune a LLaMA2-7B model on the Anthropic HH-RLHF dataset [1] to obtain a reference policy. We then align the model using our method under varying constraint levels. We use the helpfulness and harmlessness reward models from [11] for evaluating helpfulness and harmlessness, respectively, and the humor–no-humor model as the reward model for humor. We demonstrate that our iterative method satisfies the constraints while maximizing the primary utility on this task.

For implementation, similar to our main experiment, we randomly select 1000 prompts from the full HH-RLHF dataset and sample 128 responses per prompt in the dual sub-gradient step. In the policy optimization step, we train the model using DPO with 10K preference pairs randomly selected from the same dataset. We evaluate our trained models on a randomly sampled set of 2000 prompts, each with 2 responses, from the test split of the HH-RLHF dataset.

The table above shows the results for threshold $\mathbf{b}=[0.5,0.5]$, i.e., an expected improvement of 0.5 in both harmlessness and humor scores. We observe that both constraints are satisfied, while the helpfulness score improves by 9.12%.

We are happy to report additional results for other values of $\mathbf{b}$ during the discussion period and will include this multi-constraint task in the final version of our paper.

Regarding the concern about using the 7B-scale model, we note that our theoretical results and algorithm are applicable to LLMs of any size. The 7B-scale models are sufficient to address the questions we focus on, namely constraint satisfaction and trade-offs among properties (see Section 4.2). Moreover, using 7B models is a common practice in the alignment literature [2,3,4,5], and adhering to this scale facilitates fair and direct comparison with prior work.

Re: Weakness 2

The assumptions used in our analysis are intended to reflect realistic LLM training settings. First, Assumption 1 requires that the constraint functions be linearly independent over $m$ responses. This is easily satisfied in practice, as the response space of LLMs has a much higher dimensionality than the number of constraints $m$. Assumption 3 assumes that a generation distribution can be approximated by a practical LLM (i.e., a transformer network). Given the expressiveness of transformer (e.g., a universal approximation theory [6]), it is practical to assume the finiteness of the parametrization gap. We also note that KL divergence is widely used in regularized alignment problems to encourage the aligned model to remain close to a reference model. Similar assumptions have been adopted in prior work (e.g., [7]). In practice, the KL divergence can be evaluated to empirically verify the closeness between two LLMs (e.g., [1, 8]).

In our optimization assumptions, we introduce a set of feasibility assumptions in LLM policy space (Assumption 2) and parameter space (Assumptions 4 and 6). The rationale of Assumptions 2 and 4 is the same as the feasibility in constrained optimization: a set of constraints can be satisfied by some points in the policy and parameter spaces. Empirically, we have evaluated the feasibility of aligned LLMs using different constraint thresholds (see Figure 2(a)), which is not provided by the previous constrained alignment studies. We use Assumption 6 for a technical reason as follows. It is known that the primal iterates induced by a dual-based method can only converge to an optimal one on average (e.g., see [9]), which is not desirable for LLMs. To show convergence of the primal iterates, the dual function has to be smooth, at least locally at some optimal dual variables. Thus, it is necessary to strengthen the feasibility with respect to optimal dual variables in Assumption 6. Last but not least, Assumption 5 assumes the local strong convexity of the dual function, which is indeed a consequence of Assumption 1 at an optimal dual variable. Empirically, it has been observed that the dual function is locally strongly convex (e.g., see Figure 1 of [3]). Therefore, the local strong convexity is a reasonable assumption in practice.

We will incorporate these clarifications into the final version of our paper.

Re: Weakness 3

Among our key contributions, we show how to adapt and analyze the dual sub-gradient method in the context of LLMs, where the non-convexity of LLM parametrization presents a key challenge. Algorithmically, we instantiate the dual sub-gradient method as an iterative alignment algorithm, introducing warm-start initialization and post-processing (see Algorithm 1 and Appendix D) to improve convergence performance. A key algorithmic novelty is that we remove the evaluation of an optimal dual variable, which is often computed in a biased manner in practice (see [3]), and we introduce an iterative method to approximately find an optimal dual variable in an unbiased way. This technique generalizes the previous non-iterative method as a special case (e.g., [3]) and leads to significantly improved constraint satisfaction (see Figure 2).

Theoretically, we characterize the duality gap of the constrained LLM alignment problem and quantify the optimality of the LLM policies induced by the dual iterations, which, to the best of our knowledge, have not been established in the existing literature. Our assumptions are tailored to the structure of LLM alignment, bridging the gap between constrained optimization and constrained alignment.

Re: Weakness 4

The primary goal of this work is to understand how the parameterization gap affects the performance of practically aligned LLMs, rather than to develop methods for minimizing it. Drawing from the universal approximation theory of transformers [6], we observe that there always exists a transformer (e.g., a transformer block with an attention layer with 2 heads of size 1 each, and a feed forward layer with 4 hidden nodes) capable of achieving a small parameterization gap. In theory, a zero parametrization gap is commonly assumed in the development of LLM alignment methods (e.g., DPO [10]), and this assumption has led to strong practical alignment performance. In practice, most LLMs (e.g., GPT-3 small: 12 layers with 12 heads of size 768 each, and a feedforward layer with 3072 hidden notes) are significantly overparametrized, which helps ensure that the parameterization gap remains small. Therefore, our theoretical results are prescriptive for aligning practical LLMs under constraints. Last but not least, the parametrization gap can be estimated in practice using the KL divergence estimates (e.g., [1, 8]) to compare an aligned model against a much larger model. We will add a discussion on this in the final version.

Re: Question 1

We note that practical LLMs are typically highly overparametrized, making it reasonable to expect that increasing model size generally reduces the parametrization gap due to improved expressiveness. To estimate the parametrization gap for a given model class, one can treat a much larger model class (e.g., 70B) as a proxy for the ideal model and estimate the KL divergences between the model and a few sampled proxy models (or a reference model) using standard estimation techniques in [1, 8]. Since the KL divergence serves as an upper bound on the 1-norm distance, it provides a useful estimate of the parametrization gap for a given model class. We will clarify these points in the final version.

Re: Question 2

The systematic bias observed in the one-shot empirical results can be attributed to the parametrization gap and optimality of the dual variables (see Theorem 4). In our multi-shot empirical results, we demonstrate that this bias can be significantly reduced by using instance-dependent dual variables (see Theorem 3). We will include a discussion of this point in the final version.

Re: Question 3

The computational cost of our iterative method is comparable to that of the one-shot method, as each iteration requires only a single epoch of DPO. Improvements shown in Fig. 2 are typically observed within four iterations. The additional runtime of the multi-shot method only comes from dual variable estimation, which can be efficiently accelerated through engineering techniques. Please refer to Appendix D.2 for detailed runtime comparisons.

We appreciate the reviewer’s comments, and believe that the additional experiments and discussions provide further insight to open your evaluation of our work.

References:

[1] Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback

[2] Safe RLHF: Safe Reinforcement Learning from Human Feedback

[3] One-Shot Safety Alignment for Large Language Models via Optimal Dualization

[4] Stepwise Alignment for Constrained Language Model Policy Optimization

[5] Bounded Rationality for LLMs: Satisficing Alignment at Inference-Time

[6] Are transformers universal approximators of sequence-to-sequence functions?

[7] Theoretical Analysis of KL-regularized RLHF with Multiple Reference Models

[8] Scaling Laws for Reward Model Overoptimization

[9] Approximate Primal Solutions and Rate Analysis for Dual Subgradient Methods

[10] Direct Preference Optimization: Your Language Model is Secretly a Reward Model

[11] Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment

Dear Reviewer GuGo, As the author-reviewer discussion window is nearing its end, we would greatly appreciate it if you could let us know whether our updates address your concerns. We have included an additional experiment on a different task with multiple constraints, a new dataset, and a new model to further demonstrate the empirical scope of our work, and we have responded to all of your questions. We are happy to engage in further discussion or provide additional clarification if you have any remaining questions.