Thank you for acknowledging that our work addresses an important problem and for asking the interesting questions. We answer your questions below.

Questions For Authors:

1. I understand this is a paper focusing on theories. But I am curious about the parameterization part in section 4.3. How are we able to constitute the three functions, i.e. the embedding function, the linear map and the injective function? Will they be possible to integrate into current architecture of LLMs? This may be crucial in transferring theory into practice.

This is an important question. Here is one example for how to instantiate each function.

• Function (i): Softmax(NN(x)), where NN(x) is any arbitrary neural network, including fully connected networks or encoder architectures for language modelling such as attention mechanism.
• Function (ii): A linear layer. This can be implemented by any weight matrix which maps from a lower to a higher-than-or-equal-to dimension, since a randomly generated matrix almost surely has full rank.
• Function (iii). This can be composed of injective attention mechanisms, injective linear layers, injective fully-connected networks, and injective softmax layers. We explain them one-by-one: attention mechanisms can be implemented as injective by using the softmax attention, since applying softmax to a matrix will generically make it full rank (see e.g. Bridging the Divide: Reconsidering Softmax and Linear Attention  Han et al 2024). Injective linear layers as explained for function (ii) can be implemented as injective. Injective fully connected layers can be implemented by using a combination of full-rank weight matrices and injective activations such as leaky ReLU; see also references [Furuya et al., Globally injective and bijective neural operators, 2023] and [Puthawala et al, Globally injective relu networks, 2022]. Injective Softmax can be implemented as follows: x → [x, 0] → softmax([x,0]).

2. How will the level of distributional shift between proxy data and true data affect the results of this paper?

This is an interesting question! We just want to point our that distribution shift between proxy and true data can be large, so long as the conditions still hold.

However, it is very interesting to consider the case where certain assumptions hold but others do not, or assumptions hold approximately. We have had a deeper look at our conditions and can relax them in two ways, please see our answer to reviewer 6yLU in ‘Strong conditions; hard-to-verify assumptions’.

The in-between area of approximate but not full compliance with our conditions is indeed interesting to explore but require empirical evaluations, but as this paper focusses on theory, empirical evaluation is left outside the scope. We are currently working on extending the project to empirically test the sensitivity of the conditions in large-scale experiments.

3. Minors: writing issues, e.g. “iff” in condition 1

Please note that this is not a typo, iff is a mathematical shorthand for if and only if. However, we have since relaxed this condition, see our answer above.

Thank you for your insightful comments. We hope to address your main concerns below.

…says “distinct prompts”, but in the stated condition 1, we do not enforce x1≠x2…

Thanks for pointing this out, note that we have also relaxed the conditions 1 and 4:

Condition 1 can be relaxed from iff (i.e. if and only if) to just if, with the following revised statement:

Given $x_1, x_2 \in \mathcal X$ , we have: $\pi^\dagger(\cdot|x_1) = \pi^\dagger(\cdot | x_2) \text{ if } \tilde\pi(\cdot | x_1) = \tilde\pi(\cdot | x_2)$

Moreover, Condition 4 can be shown to be equivalent to

$d_{\mathcal{P_Y}}(\pi^\dagger(\cdot|x_1), \pi^\dagger(\cdot|x_2) ) \leq L d_{\mathcal{P_Y}}(\tilde\pi(\cdot|x_1), \tilde\pi(\cdot|x_2))$.

Notice that this equivalent condition implies Condition 1. This means we can simplify our presentation and combine Conditions 1 and 4 to a single condition.

Given this, perhaps the following version flows better:

‘Our first condition says that if two prompts are mapped to very different response distribution under the true policy then they cannot be mapped to very similar responses under the proxy policy.’

…relevance of this paper’s theory in guiding practice… if the Lipschitz continuity conditions make sense...

Using the medical example in the paper, (the revised) Condition $1$ (subsuming Conditions $1$ and $4$ from) means that if the gold/expert doctor thinks two symptoms are very different (i.e. they map to very different prescriptions), e.g., $|\pi^\dagger(\cdot|x_1) - \pi^\dagger(\cdot|x_2)| = d$, then the proxy/student cannot believe the two symptoms are similar, at least there should be some constant $L$ such that the student should not think that the difference between two prescriptions is less than $d/L$ for any $d$. This essentially requires that the proxy policy has the correct (within a constant scaling factor) idea of ‘distance’ in the output (response distribution) space.

In general any continuously differentiable function is Lipschitz continuous, what makes the difference is the constant $L$. Notice that we did not assume a fixed value for $L$ - it can be seen from the proof of Theorem $3$ that $L$ will impact the lipschitz constant of $\bar\pi$, which then by Theorem $5$ impact the sample complexity of learning with the proposed model parameterisation; this means that if the proxy data is ‘good’ at understanding the distance in the output space, then we need fewer samples to converge.

We have also provided a simple experiment to illustrate a case when all four conditions are satisfied, please see our answer to reviewer pBxw for the description, under ‘Weakness 3’.

…toy 1D bandit..

We have provided a simple experiment. Please see our answer to reviewer pBxw for the description, under ‘Weakness 3’.

…the low dimensional encoding condition (condition 3)…

The low-dimensional encoding is similar to the standard low-dimensional manifold condition typically used in deep learning. For example, the last token embedding in a LLM represents, in a lower dimension, the next-token prediction distribution. The key difference here is that we assume the encoding to be bi-Lipschitz; note that this can be considered a mild condition since all invertible, continuously differentiable functions whose inverse is also continuously differentiable satisfy this condition. In most deep learning tasks, we assume that data can be fitted to a deep neural network, which are continuously differentiable, satisfying the condition; in a typical low-to-high-dimension decoder, the neural network can be implemented as an injective function, and its inverse is also typically continuously differentiable.

…offline model-based RL…most common use of RLHF is: (1) collect a lot of gold reward labels from humans, (2) train a reward model against these labels, which will be our proxy reward model, (3) train an agent against this proxy reward model.

This could be best explained in the context of our simple experiment: the experiment can be readily adapted to the RLHF pipeline you described: 1) collect gold or noisy reward labels from humans, 2) train a biased reward model (biased since temperature is increased), (3) train a proxy policy from the biased reward model using our parameterisation (4) finetune on a small number of additional gold reward labels from humans.

In general, our framework applies to contextual bandits, the subclass of RL problems relevant to current LLM training. The extension to full sequential RL is left outside the scope.

How would the encoder/decoder setup look in real-world LLM cases?

Due to space constraint, we have written the detailed explanation to this question in our response to aNgX. Please kindly refer there, thank you.

Thank you for your encouraging comments and insightful questions. We hope to address your main concerns below.

Weakness 1 & 2: The derived bounds appear to be vacuous at first glance…

This is a good question. Note that in the number of samples bound in Theorem 5 and 6, there are three constants in front of $D$: the lipschitz constants $L_\phi$ and $L_{\bar\pi}$, and the matrix $p$-norm of $\|\tilde{\Theta}\|_p$; the lipschitz contant of a function intuitively sets its smoothness, and matrix norm can be seen as the lipschitz constant of the linear map given by the matrix. The power $D$ is applied the product of all terms inside the brackets (including $1/\epsilon$), and in particular if the product $= 1$ then raising to $D$ would remain it at $1$. Morally, with the same amount of training data, the higher the dimension of the input and output spaces of the hypothesis functions, the more we have to regularise the hypothesis space (i.e. the more smooth we need the functions) to achieve the same generalisation error. So here, the final bound is not necessarily an astronomically large number - it depends on how much we are willing to regularise (i.e. lower the constants of ) the hypothesis space.

Weakness 3: …a toy example or a simple experiment…

We include a simple experiment to demonstrate the efficacy of our method on a common real-world scenario - regularisation in learned rewards from tempered softmax. We set the environment as follows:

• $\mathcal X = \mathbb R^5$
• $\mathcal Y = \{1,2,3\}$, so $\mathcal{P_Y} = \Delta^2$ a two-simplex.
• $D = 1$
• $\pi^\dagger: \mathbb{R}^5 \to \mathbb R \to \Delta^2$
• $\text{logit}(\tilde\pi(y_k|x))=\frac{\big(\log(\pi^\dagger(y_k|x))+\log(\sum_k\pi^\dagger(y_k|x))\big)}{T}$. Temperature $T=5$.

We parametrise the policy model using our proposal, where in particular the final injective component is implemented using a combination of full-rank matrices and leaky-relu activation (which is injective); for more detailed discussion on how to implement the architecture please refer to our response to aNgX. We train a policy model using our proposed parameterisation, first on $8000$ proxy samples $(\tilde x, \tilde y_w, \tilde y_l)$ generated from $\tilde{\pi}$, then only finetune $\bar\pi$ from $\pi^\dagger$ on $35$ true samples $(x^\dagger, y_w^\dagger, y_l^\dagger)$ generated from $\pi^\dagger$. We compare the KL divergences $KL(\pi^\dagger, \pi^\dagger_\theta)$ and $KL(\pi^\dagger, \tilde\pi)$ to see if the learned $\pi^\dagger_\theta$ is robust against distribution shift $\pi^\dagger \mapsto \tilde\pi$. We repeat the experiments $6$ times. The results are the following:

We see that the mean with $\pi^\dagger$(pi_dag) is closer to 0 than $\tilde\pi$(proxy), having learned on only 35 true samples.

Due to the time constraints of the rebuttal, we were not able to fully tune the $\tilde\pi_\theta$ and $\pi^\dagger_\theta$ models. In the final version we will run larger experiments and include also comparisons with fully blackbox models for $\pi^\dagger$.

Line 116 (left col): What is the definition of sequence space ℓ1.

This is the space of (possibly infinite) sequences $\mathbf x$ such that $\sum_i^\infty |x_i| < \infty$. We will include this as a footnote in the paper for completeness.

I believe \ocircle is the composition operator, but defining it somewhere early might avoid confusion.

Yes it is, thanks for pointing it out. We will define it to help with clarity.

Equation 9: $\bar\pi^\dagger_\theta \to \pi^\dagger_\theta$ (let me know if I am wrong) ?

That is right! Thank you for the catch.

Thank you for your thoughtful comments, we hope to address your main concerns below.

Strong conditions; Hard-to-verify assumptions:

First we point out that while our conditions may be stringent, there exist a lot of ground truth -> proxy shifts that do meet our criteria; our results show that you can still get gains even under those shifts. For instance, increased temperature (from tempered softmax) is a common bias in learned (proxy) policy models; this prevalent form of bias satisfies our conditions 1-4.

On the other hand, since the submission, we have managed to relax the conditions in the following ways:

1. Condition 1 can be relaxed from iff (i.e. if and only if) to just if, with the following revised statement:

   Given $x_1, x_2 \in \mathcal X$ , we have: $\pi^\dagger(\cdot|x_1) = \pi^\dagger(\cdot | x_2) \text{ if } \tilde\pi(\cdot | x_1) = \tilde\pi(\cdot | x_2)$

   Moreover, Condition 4 can be shown to be equivalent to

   $d_{\mathcal{P_Y}}(\pi^\dagger(\cdot|x_1), \pi^\dagger(\cdot|x_2) ) \leq L d_{\mathcal{P_Y}}(\tilde\pi(\cdot|x_1), \tilde\pi(\cdot|x_2))$.

   Notice that this equivalent condition implies Condition 1. This means we can simplify our presentation and combine Conditions 1 and 4 to a single condition about distributional distance, conducive to future work using distribution metrics to verify this condition.

2. Instead of requiring that (the revised) Condition 1 (which subsumes the original Conditions 1 and 4) holds for all $x_1, x_2 \in \mathcal X$, we only need to require them for $x_1, x_2$, $P_\mathcal{X}$-almost surely and all proofs go through. This helps future work verify the conditions as one can approximate with samples from $P_{\mathcal X}$.

Additionally, as far as we are aware this is the first work that describes sufficient conditions on proxy preferences that allow them to be used to learn a (different) ground truth reward function. In general this problem is as hard as establishing guarantees for RLHF under distribution shifts. Given the prevalence of reward hacking and its impact (e.g., alignment faking [Greenblatt et al., Alignment faking in large language models, 2024], rare behaviour [Jones et al., Forecasting rare language model behaviors, 2025], sabotage [Benton et al., Sabotage evaluations for frontier models, 2024]) we argue that steps towards a principled framework for proxy rewards is crucial for future LLM models.

Uncertain applicability to reward models vs. human preferences.

Thank you for this. One way to apply our conditions in the real world is as follows. Consider the following example: the true reward function can be described as some (unknown) linear combination of (unknown) functions $r^\dagger(\cdot) := g(f_1(\cdot), \ldots, f_n(\cdot))$ and the proxy reward function is an (unknown) combination a subset of these unknown functions $\tilde{r}(\cdot) := h(f_1(\cdot), \ldots, f_i(\cdot))$, where $i < n$ and $g$ may or may not be equal to $h$. Note that many of the reward hacking examples mentioned in the introduction can be described using this example:

• Oscillator performance (ground truth) and output amplitude & frequency (proxy) [Bird & Layzell, 2002]: This proxy is missing factors that are needed to learn an oscillator of a specific frequency, as the proxy also rewards amplifiers that output noise. This resulted in learning a radio instead of an oscillator.
• Race ranking (ground truth) and player score (proxy) [Clark & Amodei, 2016]: This proxy is missing key properties such as race finish time which led to agents that drove in circles to collect powerups, increasing their player score.
• Student success (ground truth) and accepting admissions offer (proxy) [Golden, 2001]: The percentage of students accepting admissions offers is used by university ranking systems to indicate overall interest and selectivity. However, a university can maximize offer acceptance by rejecting highly-qualified applicants that they believe will attend another university, and so the proxy is clearly missing factors of student success.

In cases such as these, Conditions 1 and 2 generically do not hold. This is because if functions are missing in the proxy reward then the ground truth and proxy policies cannot have the same level sets. There is also no guarantee that the image of the proxy policy contains the image of the ground truth policy. In general, our conditions can be used to invalidate potential proxies, given (incomplete) prior knowledge of the structure of true and proxy policies.

….more empirical or real-world studies….

While the time constraint of the rebuttal did not allow us to run real-world experiments, we provide a simple experiment addressing the case with tempered proxy policy, a common form of bias in practice. Please find it in our response to pBxw (addressing Weakness 3)