Risk-Conditioned Fine-Tuning of Large Language Models

Introduction

Large Language Models (LLMs) ([?], [?], [53]; have demonstrated remarkable capabilities across a wide range of domains, including summarization ([?], [?], [50]), conversational assistance ([61], [?]), and complex reasoning ([2], [33], [8]). As LLMs are increasingly deployed in real-world applications with broad societal impact, it is crucial to ensure that their responses do not contain harmful or toxic content, such as discrimination ([34], [?], [26]), or violate social norms ([3], [32], [4]). To this end, recent work ([23], [56], [?]) has extended the standard fine-tuning framework of Reinforcement Learning from Human Feedback (RLHF) by incorporating safety constraints that limit the expected harmfulness of model outputs. However, expectation-based safety constraints primarily control average behavior. As a result, highly safe responses can offset harmful ones in expectation, leaving the low-probability tail of generations still vulnerable to rare but severe harmful outputs.

To address this issue, Risk-Averse RLHF (RA-RLHF) ([7]) introduces risk aversion into LLM fine-tuning. RA-RLHF adopts Conditional Value-at-Risk (CVaR) ([?], [39]) to directly optimize rare high-risk generations rather than average response quality, making it particularly suitable for LLM safety alignment in high-stakes applications, such as medical advice ([?], [60]), legal assistance ([48]), and disaster management ([37], [9], [29]), where even a small probability of severe harmful output may be unacceptable. The degree of risk aversion in CVaR is controlled by the risk level $\alpha\in(0,1)$, which determines the fraction of worst-case outcomes used to evaluate the policy. A smaller $\alpha$ concentrates on more extreme rare failures. In contrast, a larger $\alpha$ considers a broader portion of the output distribution and behaves closer to average harm, which is appropriate for lower-risk applications such as casual conversation ([61], [?]) or creative writing ([?], [27]), where overly conservative behavior can unnecessarily reduce helpfulness or diversity. Therefore, the choice of risk level plays a central role in determining the behavior of the aligned policy.

However, selecting an appropriate $\alpha$ is challenging because there is unlikely to be a universal risk level that works well across all applications and users ([?]). Different deployment scenarios may require different degrees of conservativeness, and individual users may also have different preferences over the trade-off between safety and utility ([12], [13], [16]). Existing methods, such as RA-RLHF, train the policy for a fixed risk level, and therefore do not provide a mechanism for adjusting the desired degree of risk aversion at inference time. A naive solution is to train and deploy multiple policies, one for each target risk level. However, this requires repeated training and storing multiple model instances, which is computationally expensive and may be impractical in resource-limited settings ([?], [36]).

In this paper, we propose risk-conditioned RLHF (Figure 1), which trains a single policy that can be steered across a continuum of risk levels within a deployment interval at inference time. To do this, we condition the policy $\pi$ on both the prompt $x$ and the desired risk level $\alpha\in[\alpha_{\min}, \alpha_{\max}]$, producing responses $Y \sim\pi(\cdot\mid x, \alpha)$. We formulate this as a risk-conditioned RLFH optimization problem, where training is performed over a distribution of risk levels $p(\alpha)$. As a result, the learned policy provides a continuous risk-control interface: users can select the desired degree of risk aversion at inference time without retraining or deploying multiple risk-specific models. To optimize this objective, we propose a risk-conditioned policy gradient algorithm (Algorithm 1), provide convergence analysis (Theorem 2), and present additional analysis showing that the resulting policy yields a uniform approximation of the risk frontier (Theorem 3). To instantiate the risk-conditioned policy, we further study how the risk level $\alpha$ should be injected into the LLM (Figure 2). Inspired by recent work on multi-objective fine-tuning ([?], [?]), we consider both prompt-based conditioning, which represents $\alpha$ as part of the input text, and parameter-based conditioning, which injects $\alpha$ directly into selected model parameters. Empirically, we find that explicit parameter-level conditioning provides more reliable risk control than natural-language prompting. Experiments across multiple benchmarks show that the proposed risk-conditioned policy can closely match the performance of policies trained specifically for individual risk levels, while using only a single deployable model. More importantly, the learned policy remains steerable at risk levels not observed within the training interval, achieving competitive or stronger tail-risk performance than baselines such as inference-time prompting, multiple fixed-risk policies and logit-mixing policy (Section 4.3).

In summary, our contributions are: 1. We introduce the risk-conditioned RLHF framework, where a single language model is trained to adapt to different CVaR risk levels at inference time. This formulation avoids training and deploying separate policies for different target risk levels while retaining an explicit risk interpretation. 2. We propose a risk-conditioned policy gradient algorithm for the proposed framework, provide convergence analysis, and show that the resulting policy leads to uniform approximation over the risk frontier. 3. Extensive experiments with Pythia-70M (Section 4), Pythia-2.8B model (Appendix E.1), and Llama-3.1-8B-Instruct (Appendix E.2) show that our risk-conditioned policy achieves performance comparable to risk-specific policies trained at individual risk levels, while offering better inference-time steerability.

Preliminary

Reinforcement Learning from Human Feedback (RLHF). RLHF is a widely used technique for aligning LLMs with human preferences and typically consists of three stages [61]. The first stage is supervised fine-tuning (SFT), where a LLM is fine-tuned on a high-quality dataset. In the second stage, the SFT model is prompted with $x \in\mathcal{X}$, where $\mathcal{X}$ is a finite context space, and generates multiple responses $y_i \in\mathcal{Y}$, where $\mathcal{Y}$ is a finite completion space. These responses are then presented to human annotators, who provide preference labels. A reward model $r(x,y)$ is subsequently trained from these preference comparisons. The third stage is policy optimization, where the learned reward model provides feedback for further fine-tuning the SFT model. In particular, let $\pi\in\Delta_{\mathcal{Y}}^{\mathcal{X}}$ denote an LLM policy that maps each prompt $x$ to a discrete probability distribution $\pi(\cdot|x) \in\Delta_{\mathcal{Y}}$, where $\Delta_{\mathcal{Y}}$ is the set of discrete distributions over $\mathcal{Y}$. The standard RLHF objective optimizes a policy $\pi$ to maximize the expected reward while regularizing its deviation from a reference policy $\pi_{\mathrm{ref}}$ through a KL-divergence penalty: $\mathbb{E}_{x\sim D,y\sim\pi(\cdot|x)}[r(x,y)]-\beta\operatorname{KL}(\pi\|\pi_{\mathrm{ref}})$, where $D$ is a dataset of prompts, and $\operatorname{KL}(\pi\|\pi_{\mathrm{ref}})=\mathbb{E}_{x\sim D,y\sim\pi(\cdot|x)}\left[\log\frac{\pi(y|x)}{\pi_{\mathrm{ref}}(y|x)}\right]$. Equivalently, the RLHF objective can be written as $\mathbb{E}_{x\sim D,y\sim\pi(\cdot|x)}\left[r(x,y)-\beta\log\frac{\pi(y|x)}{\pi_{\mathrm{ref}}(y|x)}\right]$. We define $G(x,y):=r(x,y)-\beta\log\frac{\pi(y|x)}{\pi_{\mathrm{ref}}(y|x)}$ as the regularized reward. In this work, we focus only on the third stage.

Conditional Value-at-risk (CVaR). CVaR has recently been introduced as a risk-sensitive criterion for evaluating learned policies [7]. While the expected value in the standard RLHF objective measures the average performance of a policy, CVaR focuses on tail behavior and captures how the policy performs under unfavorable outcomes [18].

Formally, let $Z$ be an integrable random variable. For a risk level $\alpha\in(0,1]$, the value-at-risk (VaR) of $Z$ is defined as: $\operatorname{VaR}_{\alpha}(Z)=\min\{z\mid F(z)\ge\alpha\}$, where $F(z)=\mathbb{P}(Z\le z)$ is the cumulative distribution function (CDF). $\operatorname{VaR}_{\alpha}(Z)$ is the threshold below which approximately an $\alpha$-fraction of outcomes fall. CVaR then measures the average value of $Z$ in this lower tail as $\operatorname{CVaR}_{\alpha}(Z)=\mathbb{E}_{z\sim Z}\{z\mid z\le\operatorname{VaR}_{\alpha}(Z)\}$. A useful variational characterization of CVaR is given by [14, 20]:

where $(\cdot)_{+}=\max(\cdot,0)$. In this formulation, $\eta$ plays the role of a learnable tail threshold, and the penalty term $(\eta-Z)_{+}$ emphasizes samples whose outcomes fall below this threshold. While CVaR is often formulated as minimizing upper-tail costs, we adopt the equivalent reward-maximization formulation appropriate for RLHF and use this variational form for optimization.

Risk-Averse RLHF. To improve the tail performance of RLHF policies, [7] incorporates CVaR into the standard RLHF objective. Instead of maximizing the average regularized reward over all sampled responses, risk-averse RLHF optimizes the average performance over the worst $\alpha$-fraction of responses. Formally, for a fixed risk level $\alpha$, the objective is to find a policy $\pi^{*}$ that solves $\max_{\pi}\mathbb{E}_{x\sim D}\left[\operatorname{CVaR}_{\alpha}(G(x,Y))\right]$, where $Y\sim\pi(\cdot|x)$. For each prompt $x$, $G(x,Y)$ is a random variable induced by sampling a response from the policy, and $\operatorname{CVaR}_{\alpha}(G(x,Y))$ measures the expected regularized reward among the worst $\alpha$-fraction of responses. Thus, the objective encourages the policy to avoid low-reward tail responses. We include a detailed related work on risk-conditioned RL, risk averseness in LLMs, and multi-objective finetuning in Appendix A.

Risk-conditioned RLHF

This section presents our risk-conditioned RLHF framework. We first formalize the risk-conditioned RLHF problem in Section 3.1. We then introduce a risk-conditioned policy gradient algorithm for optimizing the proposed problem in Section 3.2. Finally, in Section 3.3, we describe several practical mechanisms for instantiating risk-conditioned policies by injecting the risk level $\alpha$ into LLMs.

Problem Formulation

Our goal is to learn a single policy that can adapt to different risk levels at inference time. Instead of training a separate risk-averse policy for each fixed $\alpha$, we augment the policy with $\alpha$ as an additional conditioning input together with the prompt $x$. Formally, we define a risk-conditioned policy as $\pi: \mathcal{X} \times(0, 1] \to\Delta_{\mathcal{Y}}$, where $\pi(\cdot\mid x, \alpha)$ denotes the response distribution for prompt $x$ under risk level $\alpha$. Given $Y \sim\pi(\cdot\mid x, \alpha)$, we write the corresponding regularized reward as $G(x,Y;\alpha)$, emphasizing that both the sampled response and the KL-regularized reward are induced by the $\alpha$-conditioned policy. We define the risk-conditioned RLHF objective as

where $p(\alpha)$ is a distribution supported on a risk interval $[\alpha_{\min}, \alpha_{\max}] \subset(0, 1]$. For each sampled $\alpha\in[\alpha_{\min}, \alpha_{\max}]$, the policy $\pi(\cdot\mid x, \alpha)$ is optimized to improve the average regularized reward among the worst $\alpha$-fraction of responses for each prompt. By training over $\alpha\sim p(\alpha)$, the resulting policy learns a continuous risk-control interface, allowing the desired level of risk aversion to be selected at inference time without training.

Risk-conditioned Policy Gradient

A direct way to optimize (2) is to sample risk levels and apply an existing fixed-$\alpha$ risk-averse method to the conditioned policy. For example, RA-RLHF ([7]) estimates the CVaR objective by ranking sampled trajectories according to their rewards and updating the policy using low-reward tail samples. Although this method provides a practical way to estimate risk aversion, it relies on empirical tail selection, which is difficult to characterize the resulting optimization error. To learn a continuous risk-control interface and obtain an analyzable optimization procedure, we instead use the variational form of CVaR in (1). Specifically, we treat $\eta$ as an optimizable tail threshold, which allows us to develop a gradient-based method for jointly updating the risk-conditioned policy and the threshold predictor. Formally, we parameterize the tail threshold by a neural network $\eta_{\omega}(x,\alpha)$ and the risk-conditioned policy by $\pi_{\theta}(\cdot\mid x,\alpha)$. We then define

Under (1), optimizing the risk-conditioned objective in (2) leads to the parameterized optimization problem $\max_{\theta,\omega}\mathcal{J}(\theta,\omega)$. We next derive the gradients of $\mathcal{J}(\theta,\omega)$ with respect to the policy parameters $\theta$ and the threshold parameters $\omega$.

Theorem 1. The gradients of $\mathcal{J}(\theta,\omega)$ for our proposed risk-conditioned RLHF objective can be computed as follows:

where $u_{\theta,\omega}(x,Y,\alpha) = \eta_{\omega}(x,\alpha) - \frac{1}{\alpha}\left(\eta_{\omega}(x,\alpha)-G(x,Y;\alpha)\right)_{+}$ and $\mathbb{1}\{\cdot\}$ is the indicator function.

The proof is deferred to Appendix B.1. In practice, the exact gradient (3) and (4) are unavailable and can only be estimated via stochastic samples. We refer the details to Appendix B.2. Specifically, given a batch $\{(x_b,\alpha_b)\}_{b=1}^{B}$, where $B$ is the batch size, we sample $N$ responses $y_{b,1},\ldots,y_{b,N} \sim\pi_{\theta}(\cdot\mid x_b,\alpha_b)$ for each prompt risk pair. We then construct stochastic estimators $\widehat{g}_{\omega}(\theta,\omega)$ and $\widehat{g}_{\theta}(\theta,\omega)$ to approximate $\nabla_{\omega}\mathcal{J}(\theta,\omega)$ and $\nabla_{\theta}\mathcal{J}(\theta,\omega)$, respectively. Moreover, we show that their estimation errors decrease on the order of $\mathcal{O}\left(\frac{1}{BN}+\frac{1}{B}\right)$, which vanishes as the batch size $B$ and the number of completions $N$ become large.