PaperlayerView exact revision on arXiv

Source-linked AI summary

Risk-Conditioned Fine-Tuning of Large Language Models

Zixuan Liu, Fangzheng Wu, Brian Summa, Zizhan zheng

arXiv:2609.08064v1cs.LG

TL;DR

LLM safety methods that control expected harmfulness can leave rare severe outputs exposed, while fixed-risk RA-RLHF cannot adapt risk aversion at inference time. The paper trains one policy conditioned on the desired CVaR level and finds that it can closely match risk-specific policies while improving steerability across risk levels.

  • Problem

    Existing safety constraints primarily control average behavior, and fixed-risk RA-RLHF does not provide inference-time adjustment across applications and user preferences.

  • Method

    Risk-conditioned RLHF trains a single policy over risk levels and conditions its responses on both the prompt and desired risk level α.

  • Results

    Across three benchmarks, risk-conditioned policies closely match risk-specific policies while improving steerability, with explicit parameter-level conditioning providing more reliable risk control than prompting.

  • Takeaways & Limitations

    Risk conditioning offers a way to amortize risk-averse alignment across diverse deployment scenarios and user safety requirements without separate risk-specific models.

  • Takeaways & Limitations

    Evaluation mainly relies on learned reward or cost models, which are imperfect proxies for human judgments; human validation is left for future work.

Abstract

from arXiv · show

Large Language Models (LLMs) are increasingly deployed in settings where rare but severe harmful generations can have significant consequences. Existing Risk-Averse RLHF addresses this issue by optimizing Conditional Value-at-Risk (CVaR), but it trains policies for fixed risk levels and therefore cannot adjust the desired degree of risk aversion at inference time. In this paper, we propose risk-conditioned RLHF, a framework that trains a single policy that provides a continuous risk-control interface, enabling users to select different degrees of risk aversion without retraining or deploying multiple risk-specific models. Experiments across multiple benchmarks demonstrate that a single risk-conditioned policy can adapt to different risk levels at inference time, enabling more flexible and risk-aware LLM deployment.

1 Introduction

Risk-conditioned RLHF addresses rare, severe harmful generations that average safety constraints may leave in the low-probability tail. It trains one policy across risk levels so users can adjust risk aversion at inference time without deploying separate policies.

  • Motivation: Expectation-based safety constraints can leave rare, severe harmful generations unresolved because safe responses may offset harmful ones on average.RA-RLHF instead targets the worst α-fraction using CVaR.
  • Motivation: Fixed-risk RA-RLHF cannot adjust risk aversion at inference time, while multiple risk-specific policies require repeated training and storage.Different applications and users may require different degrees of conservativeness.
  • Approach: The proposed framework provides a continuous risk-control interface, allowing inference-time selection of risk aversion without retraining or deploying multiple models.Figure 1 contrasts this capability with prior fixed-risk methods.
  • Approach: Risk-conditioned RLHF conditions one policy on prompt x and risk level α sampled across a deployment interval.The policy produces responses according to π(· | x, α) and is trained over a distribution of risk levels.
  • Conditioning: Explicit parameter-level conditioning provides more reliable risk control than natural-language prompting in the reported experiments.The paper studies prompt-based and parameter-based conditioning mechanisms.
  • Results: Experiments show competitive or stronger tail-risk performance at unseen risk levels and comparable performance to risk-specific policies with better steerability.The paper reports results across multiple benchmarks and model settings.

2 Preliminary

The preliminary section defines standard RLHF, CVaR, and risk-averse RLHF. It distinguishes average-reward optimization from optimization focused on unfavorable response tails.

  • RLHF: Standard RLHF maximizes expected reward while penalizing deviation from a reference policy through KL divergence.The regularized reward combines the reward model with the KL penalty.
  • CVaR: CVaR evaluates unfavorable tail behavior rather than average performance by focusing on outcomes below a risk-dependent threshold.VaRα is the threshold below which approximately an α-fraction of outcomes fall.
  • Risk-Averse RLHF: Risk-averse RLHF maximizes CVaRα of regularized rewards, thereby optimizing performance among the worst α-fraction of responses.The risk level α is fixed in this formulation.

3 Risk-conditioned RLHF

The framework trains a policy conditioned on prompts and risk levels, with stochastic policy-gradient optimization and mechanisms that inject risk information into the model. Its analysis connects training performance on a risk grid to behavior across the full risk interval.

  • Framework: Risk-conditioned RLHF augments the policy input with α and trains over sampled risk levels to create a continuous risk-control interface.The policy is optimized for the worst α-fraction of responses at each sampled risk level.
  • Policy Gradient: The proposed optimization parameterizes the objective with policy and threshold-network parameters and derives gradients for both components.The gradients are estimated from stochastic samples rather than computed exactly.
  • Policy Gradient: Stochastic training samples prompts, risk levels, and completions, then estimates and updates the threshold network and policy gradients.The estimation error decreases as batch size B and completions per prompt N increase.
  • Theory: Theorem 2 states that Algorithm 1 converges to a nonsmooth stationary point with stationarity error O(T^-1/2) under a constant step size.The guarantee assumes the conditions stated in the paper.
  • Theory: Theorem 3 attributes unseen-risk-level error to optimization error on observed levels and grid-coverage error 2Lh.The grid-coverage component decreases as the training risk grid becomes denser.
  • Conditioning Mechanisms: Parameter conditioning shares most parameters across risk levels while a smaller conditioned subset and gating network map α to model behavior.The conditioned policy has parameter count O(K|S| + |SC| + |Sgate|).
  • Conditioning Mechanisms: The paper studies logit-, attention-, and prompt-conditioned policies, with attention conditioning allowing α to influence selected intermediate token interactions.Prompt conditioning appends the target risk level to the input and requires no additional model parameters.

4 Experiments

Experiments evaluate conditioning mechanisms, computational overhead, benchmark performance, steerability, cross-evaluation, and ablations across three benchmarks. The results support parameter-based risk conditioning as an effective and efficient approach for adapting to unseen and observed CVaR risk levels.

  • Experiment Setup: Experiments compare risk-conditioning mechanisms, baselines, and ablations across IMDB-Gen, RealToxicityPrompts-Gen, and Safe-RLHF using task-specific reward or cost metrics.The study includes Pythia-70M, Pythia-2.8B, and Llama-3.1-8B-Instruct, with higher metric values indicating better performance.
  • Conditioning Mechanisms: Parameter-based conditioning outperforms prompt-based conditioning, while attention conditioning only slightly exceeds logit conditioning across the evaluated CVaR risk levels.Training uses risk levels {0.1, 0.3, 0.5, 0.7, 0.9} and holds out {0.2, 0.4, 0.6, 0.8} to assess interpolation.
  • Conditioning Mechanisms: Conditioned policies add negligible parameter overhead, with peak GPU memory and per-update training time remaining close to RA-RLHF.The attention-conditioned policy is selected as the default for the remaining experiments.
  • Core Benchmarking Results: The full Risk-conditioned LM remains close to oracle variants, indicating modest degradation from shared training across risk levels, while Risk-conditioned-Oracle usually outperforms RA-RLHF-Oracle.Risk-conditioned-Oracle isolates the gradient-based CVaR optimization from sharing one conditioned policy across risk levels.
  • Core Benchmarking Results: A single risk-conditioned policy interpolates effectively across held-out risk levels, outperforms RA-RLHF-Mix in most cases, and exceeds prompting and logit mixing for reliable risk control.Additional evaluations report smooth, stable, and overall monotonic changes in worst-tail behavior across denser risk levels.
  • Ablations: Increasing conditioned parameter sets from K = 1 to K = 5 substantially improves held-out risk performance, but larger capacities provide limited or negative additional benefit.Increasing K from 5 to 16 adds parameters from 0.74M to 2.37M for only 0.39% average-score improvement, while K = 32 slightly degrades performance.

5 Conclusion

Risk-conditioned RLHF trains one language model to adapt to different CVaR risk levels at inference time. Across three benchmarks, it closely matches risk-specific policies while improving steerability.

  • A single risk-conditioned policy adapts to different CVaR risk levels at inference time, unlike fixed-risk RA-RLHF policies.
  • Across three benchmarks, risk-conditioned policies closely match risk-specific policies while improving steerability.
  • Risk conditioning may amortize risk-averse alignment across diverse deployment scenarios and user safety requirements.

Limitations

The evaluation relies mainly on imperfect reward or cost models, and the deployment procedure for choosing α remains unresolved. Human validation and better α calibration are left for future work.

  • Evaluation mainly relies on reward or cost models, which are imperfect proxies for human judgments and may permit reward hacking or superficial improvements.The authors propose human scoring of worst-tail responses as additional validation.
  • The framework does not fully solve how users or system designers should choose α for a specific deployment domain.The authors characterize α selection as a nontrivial, application-specific calibration problem.
  • No universally accepted operational procedure for choosing α exists across CVaR application domains.The paper describes regulation, sensitivity analysis, and estimation from sequential choice data as possible approaches, not a deployment-specific LLM rule.
  • Calibrating α for LLM safety applications is identified as an important, underexplored direction beyond the scope of this paper.

Ethical Considerations

The paper positions risk-conditioned RLHF as a risk-level control method for LLM alignment, distinct from fixed-risk optimization, prompt-level control, and reward-weight conditioning. Its analysis also relies on sampling, boundedness, and regularity assumptions.

  • Positioning: The approach brings risk-conditioned policies from standard reinforcement learning into LLM alignment.
  • Positioning: Unlike fixed-risk policy optimization and prompt-level risk control, the method conditions one policy on the target risk level and interpolates to unseen levels without retraining.
  • Positioning: Unlike reward-weight conditioning methods, the framework varies risk sensitivity under a unified alignment procedure.
  • Optimization procedure: The algorithm uses sampled prompts, risk levels, and completions to estimate threshold-network and policy gradients before updating parameters.
  • Analysis assumptions: The stochastic estimators are analyzed under assumptions including i.i.d. prompt-risk sampling, i.i.d. completions, and bounded regularized returns.
  • Optimization analysis: Under the stated assumptions, the stochastic gradient estimators are unbiased and have mean-squared-error bounds.

B.3 Analysis of Algorithm 1

The analysis treats Algorithm 1 as stochastic optimization of a generally nonsmooth hard-hinge CVaR objective. Under boundedness, smoothness, sampling, and weak-convexity assumptions, it establishes valid generalized gradients and convergence toward nonsmooth stationarity.

  • Objective and gradients: Algorithm 1 is analyzed as stochastic subgradient optimization because the hard-hinge CVaR objective is generally nonsmooth.
  • Convergence analysis: The Moreau-envelope gradient serves as the stationarity measure for the weakly convex nonsmooth objective.
  • Assumptions: The analysis assumes bounded threshold-network and regularized-return gradients, bounded iterates, and weak convexity of the negated objective.
  • Objective and gradients: The Clarke subdifferential justifies the indicator-based update as a valid stochastic generalized gradient at the hinge's nondifferentiable point.
  • Convergence analysis: The convergence theorem uses a shared learning rate for the policy and threshold blocks, while block-specific analysis scales descent by the smallest block step size.
  • Convergence analysis: With a constant step size γ = Θ(T^-1/2), Algorithm 1 approaches nonsmooth stationarity at rate O(T^-1/2).

B.4 Analysis of Logit-Mixing LM

Logit mixing combines fixed-risk policies through normalized geometric interpolation, limiting intermediate behavior to responses supported by both endpoints. Risk-conditioned policies instead learn a nonlinear mapping from risk level to behavior, allowing intermediate risk-sensitive responses during training.

  • Logit mixing linearly interpolates endpoint logits, producing a normalized geometric mixture of independently trained fixed-risk policies.The normalization constant converts the interpolated logits into a probability distribution.
  • The resulting mixture cannot assign substantial probability to response regions weakly supported by both endpoint policies.This structural limitation means intermediate behaviors absent from both endpoints are unlikely to be recovered.
  • A risk-conditioned policy can learn high probability for desired intermediate behaviors through nonlinear conditioning on α rather than endpoint interpolation.The risk level is injected through a learned gating mechanism, giving the policy greater flexibility over intermediate behaviors.
  • When the balanced response has probability δ under both endpoints, logit mixing assigns it probability δ Zτ(x), which remains small when δ is small.The endpoint-supported responses instead receive unnormalized masses (1 −δ)1−τδτ and δ1−τ(1 −δ)τ.
  • Attention conditioning injects α into internal attention computations, whereas logit conditioning modifies only the final output layer.The attention-conditioned variant applies risk-dependent updates to selected attention projections before final token logits are produced.
  • Prompt conditioning encodes α and its interpretation as a natural-language prefix, and the authors report improved consistency across risk levels.For α = 0.2, the prefix explains that lower α means stricter safety and describes the worst-20% tail objective.

C.2 Implementation Details for Algorithm 1

Algorithm 1 samples risk levels from a broad grid, leaving intermediate values unseen during training to evaluate steerability and interpolation. The implementation uses a lightweight threshold network conditioned on prompt representations and α.

  • Training samples risk levels uniformly from a grid covering the deployment range, with intermediate values held out for evaluation.This setup tests whether the learned policy can steer and interpolate across unseen risk levels.
  • The threshold network combines mean-pooled prompt embeddings with α and maps them through a two-layer MLP to produce a scalar output.The MLP uses hidden size 256 and tanh activations.
  • Tables 5 and 6 report the hyperparameters used for Algorithm 1 and PPO training.

D Additional Experiment Setup

The experiments use IMDB, RealToxicityPrompts, and Safe-RLHF, with held-out evaluation splits and dataset-specific reward or cost models. Evaluation measures CVaR across risk levels from sampled responses.

  • Experiments cover IMDB, RealToxicityPrompts, and Safe-RLHF, using 80% of each dataset for fine-tuning and 20% for evaluation.
  • Dataset-specific models provide sentiment rewards, toxicity costs, or harmfulness costs, with lower toxicity and harmfulness costs indicating safer outputs.The reward is defined as negative cost for RealToxicityPrompts and Safe-RLHF.
  • The implementation uses specified Python, PyTorch, and TRL versions, with RA-RLHF following its official implementation and default hyperparameters.
  • Evaluation samples 64 responses per prompt and computes CVaRα as the average reward among the lowest ⌈α × 64⌉ samples.Reported values are means and standard deviations across five random seeds.
  • The logit-mixing baseline selects the two nearest trained risk levels surrounding each target α and interpolates between their policies.

E.1 Additional Results with Pythia-2.8B

On Pythia-2.8B, parameter-based conditioning generally outperforms prompt conditioning, while attention conditioning slightly exceeds logit conditioning. Across held-out and larger-model evaluations, one conditioned policy remains competitive with risk-specific alternatives and provides stable control over unseen risk levels.

  • Parameter-based conditioning generally outperforms prompt-based conditioning across CVaR risk levels on Pythia-2.8B.Among parameter-based variants, attention conditioning slightly outperforms logit conditioning.
  • The conditioning mechanisms add modest parameters without meaningfully increasing memory usage or training time.This supports using one risk-conditioned policy instead of storing multiple separate risk-specific models.
  • Risk-conditioned-Oracle outperforms RA-RLHF-Oracle in most settings at the larger Pythia-2.8B scale.The full Risk-conditioned LM is slightly behind the oracle variants.
  • Risk-conditioned LM remains close to RA-RLHF-Oracle at unseen held-out risk levels across three benchmarks and generally exceeds RA-RLHF-Mix.It avoids training, storing, and selecting among multiple risk-specific policies.
  • On Llama-3.1-8B-Instruct with Safe-RLHF, the method achieves performance comparable to the oracle while avoiding additional cost.
  • Across varying α, the policy shows an overall smooth and monotonic trend, with stable control across a denser range of unseen values.A few adjacent α values show small nonmonotonic fluctuations.

E.4 Additional Ablation Studies

Additional ablations show that risk control depends on moderate conditioning capacity and sufficiently broad, dense training-grid coverage. Qualitative and calibration evaluations further examine behavior across risk levels and interpolation beyond those observed during training.

  • Risk-Level Evaluation: Calibration, dense risk-control, and qualitative evaluations assess how performance and generations vary across observed and held-out CVaR risk levels.The figures and tables cover Pythia-2.8B and Pythia-70M experiments, including three-benchmark comparisons and Safe-RLHF examples.
  • Number of Conditioned Parameter Sets K: Increasing conditioned parameter sets from K = 1 to K = 5 clearly improves performance on IMDB and RealToxicityPrompts.The evaluation uses unseen risk levels α ∈{0.2, 0.4, 0.6, 0.8}, with other training configurations fixed.
  • Number of Conditioned Parameter Sets K: K = 16 yields modest additional gains, averaging +3.24% on IMDB and +3.48% on RealToxicityPrompts, while extra parameters increase by 220.3% relative to K = 5.Increasing K to 32 instead reduces performance on both datasets despite greater parameter overhead.
  • Number of Conditioned Parameter Sets K: Risk-control benefits improve up to moderate conditioning capacity, then quickly saturate; excessive K may make optimization less stable.The reported trend supports using a small number of conditioned parameters rather than scaling K indefinitely.
  • Training Grid Coverage and Mesh Size: The partial grid {0.1, 0.3, 0.5} performs competitively near its covered region but drops at larger α, especially at α = 0.8.This indicates that limited grid coverage restricts off-grid steerability beyond the trained interval.
  • Training Grid Coverage and Mesh Size: The sparse full-coverage grid {0.1, 0.5, 0.9} improves performance at larger α over the partial grid but remains slightly below the default grid on average.This matches the stated relationship between denser training-grid coverage and lower off-grid approximation error.
Loading 2609.08064v1…