Tuning Language Models by Proxy

Thank you for your thoughtful suggestions! We are grateful that you recognize the simplicity and effectiveness of the method. We hope that the following discussion addresses your concerns.

When would this method be needed?

You are definitely right that currently, there are not many situations where model producers provide logit distributions but not model parameters.

Our response is three-fold:

1. While proxy-tuning enables tuning of black box models (and this is our main motivation), it has many benefits even for white box models. For example, it has been useful for a distributed learning setup where the expert is tuned on data that must stay on-device, and therefore cannot be used to tune the larger base model. Proxy-tuning also enables a “tune once, proxy-tune many” setup, where hundreds of models can be improved for the training cost of tuning one model. Even with access to parameters of the base model, tuning extremely large models requires much more resources than proxy-tuning with small (anti-)experts (we used TPUs to finetune the 70B model, while three A100 GPUs is enough to proxy-tune the 70B model). Finally, the strength of tuning is controllable in proxy-tuning (as we show in §6.2), which is not true for direct tuning.
2. Future work may be able to indirectly obtain model logits. For instance, a recent paper (Carlini et al., 2024) reconstructed the entire logit distribution from GPT-3.5-turbo through multiple queries (something we did not know!). Although OpenAI changed the functionality of the API in response, the paper shows that it is still possible to recover the complete logit vector at an increased cost.
3. Very wishfully, we hope that proxy-tuning incentivizes model producers to provide logits in the future, because of the greater user customization it enables. Thus, even though the current setting is not very common now, perhaps it will be more common in the future! Following your suggestion, we will provide more discussion of this in the paper.

The paper would be better without the "explanation" for TruthfulQA results

We really appreciate your suggestion. We definitely intend for it to be a hypothesis, not an explanation. We will rephrase it to say “Direct tuning has been shown to sometimes hurt performance on knowledge-intensive tasks, and it is possible that decoding-time algorithms provide an avenue for better knowledge preservation.”

Thank you for your thoughtful questions, which we address below.

Discussion of past work

DExperts and VDD (and many other methods) use logit arithmetic to steer the generation in a desirable way, whereas we show that it is possible to achieve the effect of finetuning at decoding-time. We believe that this is a surprising finding.

Moreover, we have found that proxy-tuning is equivalent to tuning $\mathcal M$ with the implicit reward model underlying tuning of $\mathcal M^-$ (see response to w3N3)! Thus, proxy-tuning has some special theoretical properties.

Why use the anti-expert?

Following your suggestion, we experiment with an ablation without the anti-expert, which we will add to the next revision of the paper. Specifically, we additively combine $\mathcal M$ and $\mathcal M^+$ with a hyperparameter $\alpha$ applied to $\mathcal M^+$. We use 200 examples sampled from the full test set. We find that proxy-tuning consistently outperforms the ablation without the anti-expert. For AlpacaFarm, the best setting $\alpha=1.0$ gives a win rate of 81.5% compared to 90% for proxy-tuning; for GSM, the best setting $\alpha=0.4$ gives 26% accuracy compared to 26.4%. Note that $\alpha$ is very task-sensitive!

We conclude that proxy-tuning works out-of-the-box, while ablating the anti-expert can give strong performance with a task-sensitive hyperparameter search.

Step by step?

Yes, we adjust the output logits at every time step.

If LMs do not provide output probabilities, does the method work?

Our method does require access to output probabilities or logits. However, a recent paper (Carlini et al., 2024) showed that even when these logits are not provided, they can be reverse engineered — in fact, the authors recovered the entire logit distribution from GPT-3.5-turbo using many queries! Moreover, proxy-tuning may be applicable even with partial access to output logits, as shown in §7.

Example generations from Llama2-13B-chat

See https://postimg.cc/FdrGs1Jx.

Discussion of TruthfulQA findings

For TruthfulQA in the open-ended setting, there are two interesting findings:

1. The 13B-chat slightly outperforms 70B-chat (by 0.8%). We suspect that the benchmark is saturated with scale.
2. The proxy-tuned models outperforms their directly-tuned counterparts. As TruthfulQA is a knowledge-intensive task, we hypothesize that decoding-time algorithms can better preserve knowledge than direct tuning.

Thank you for your insightful questions, and we are grateful that you recognize the effectiveness of the approach.

Inference cost

Proxy-tuning does incur a greater inference-time cost, which we quantify in §C.1. We note that the increased runtime is due to a sequential execution of the models in proxy-tuning; in practice it can be greatly accelerated by deploying on multiple GPUs in parallel that communicate with each other.

In addition, as models push the limits of scale and available training data, we believe inference-time methods will be important to pushing model capabilities further. In many cases, users may be willing to wait longer for a better generation.

Assumes a situation with small accessible LM and large inaccessible LM

Note that the small and large pretrained models do not need to be in the same model family, as long as they share the same vocabulary (or at least enough overlap for the task of interest), as we showed for applying Llama-2 models to steer GPT-3.5. Excitingly, very recent work (Minixhofer et al., 2024) developed a method that swaps a LM’s tokenizer with an arbitrary new one, and should be able to alleviate the requirement of shared tokenizers in proxy-tuning!

Moreover, while proxy-tuning enables tuning of black box models (and this is our main motivation), it has use cases outside of black box settings. For example, it has been useful for a distributed learning setup where the expert is tuned on data that must stay on-device, and therefore cannot be used to tune the larger base model. Proxy-tuning also enables a “tune once, proxy-tune many times” setup, where hundreds of models can be improved for the training cost of tuning one model. Even with access to parameters of the base model, tuning extremely large models requires much more resources than proxy-tuning. Finally, the strength of tuning becomes controllable for different use cases (as we show in §6.2), which is not true for direct tuning.

Failure modes

We did not observe any consistent failure modes.

Investigation of different temperature values

We use temperature in the code adaptation experiments (§4), following the same decoding hyperparameters as the Codex paper (Chen et al., 2021) with temperature = 0.8 and top $p$ = 0.95. For a consistent comparison, we did not explore other options. These hyperparameters are meant to encourage diversity in sampling.

Thank you for your insightful comments, and we are grateful that you recognize the effectiveness of the method in many use cases.

Underlying reasons for effectiveness

Since submission, we have developed a better theoretical understanding of proxy-tuning, which we summarize below.

Given a pretrained model $\mathcal M$, the objective for RL with a KL divergence penalty (used by e.g., PPO) is defined by

$\operatorname{argmax}\_{\mathcal{M}^\ast} \underset{y\sim P\_{\mathcal{M}(\cdot\mid x)}}{\mathbb{E}} r(x,y) - \beta \operatorname{KL}( P\_{\mathcal{M}}(\cdot \mid x) \mid\mid P\_{\mathcal{M}^\ast}(\cdot \mid x))\quad\quad\quad\text{(1)}$

This is well-known to have a closed form solution (Korbak et al., 2022),

$P_{M^*}(y \mid x) = \frac{1}{Z} P_{M}(y \mid x) \exp \left(\frac{1}{\beta} r(x,y)\right)\quad\quad\quad\text{(2)}$

More generally, this means any finetuned model $\mathcal M^\ast$ can be viewed as implicitly optimizing an underlying reward $r$ given by

$r(x,y)=\beta \log\frac{P_{\mathcal M^*} (y\mid x)}{P_{\mathcal M} (y\mid x)},\quad\quad\quad\text{(3)}$

as can be seen by substituting Eq (3) into the RHS of Eq. (2) to recover the LHS.

Now, proxy tuning is

$P_{\mathcal M^*}(y \mid x) \propto P_{\mathcal M}(y \mid x) \frac{P_{\mathcal M^+}(y \mid x)}{P_{\mathcal M^-}(y \mid x)}\quad\quad\quad\text{(4)}$

where $\mathcal M^+$ and $\mathcal M^-$ are small tuned and untuned proxies. Thus proxy-tuning can be seen as tuning $\mathcal M$ with the underlying reward for tuning of the small anti-expert, $r(x,y)=\beta \log\frac{P_{\mathcal M^+} (y\mid x)}{P_{\mathcal M^-} (y\mid x)}$. Intuitively, to the extent that the tuning of $\mathcal M$ and of $\mathcal M^-$ correspond to the same underlying reward function, proxy-tuning is equivalent to true finetuning.

Analysis of logits from proxy-tuned and directly-tuned models

Following your suggestion, we compare the proxy-tuned and directly-tuned models in terms of the KL div in their predictions, when conditioning on the same prefix. We used AlpacaFarm prompts and only looked at the prediction for the first time step of generation. We find that the median KL div between the proxy-tuned and directly-tuned models is 0.147, compared to 0.239 between the base (untuned) model and the directly-tuned model. This means that the proxy-tuned probability distributions are indeed more similar to those of the directly-tuned model!