Kad: A Framework for Proxy-based Test-time Alignment with Knapsack Approximation Deferral

TL;DR

The approach can be described as token-specific cascading method, where the token-specific deferral rule is reduced to 0-1 knapsack problem, and primal and dual approximations of the optimal deferral decision are derived.

Abstract

Several previous works concluded that the largest part of generation capabilities of large language models (LLM) are learned (early) during pre-training. However, LLMs still require further alignment to adhere to downstream task requirements and stylistic preferences, among other desired properties. As LLMs continue to scale in terms of size, the computational cost of alignment procedures increase prohibitively. In this work, we propose a novel approach to circumvent these costs via proxy-based test-time alignment, i.e. using guidance from a small aligned model. Our approach can be described as token-specific cascading method, where the token-specific deferral rule is reduced to 0-1 knapsack problem. In this setting, we derive primal and dual approximations of the optimal deferral decision. We experimentally show the benefits of our method both in task performance and speculative decoding speed.

Figures (4)

• Figure 1: Illustration of token-specific cascading distribution ${\bm{\pi}}^{<\lambda}$ with $\lambda = 0.4$. The left and center plots show the probability mass function (PMF) of ${\bm{p}}$ and ${\bm{q}}^*$, respectively. The right plot shows the PMF of ${\bm{\pi}}^{<\lambda}$, where the blue parts shows the mass coming from ${\bm{p}}$ and the red parts the one coming from ${\bm{q}}^*$ (including rescaling by $\alpha$). Dotted bars show that mass from ${\bm{p}}$ that was rejected by the deferral decision. We can observe the most probable token of ${\bm{p}}$ and the one of ${\bm{q}}^*$ both have high probabilities in the resulting mixture.
• Figure 2: Empirical acceptance rate per model in speculative generation schema. In brown the dual approximation ${\bm{\pi}}^{<\lambda}$ with $\lambda = 0.4$, in blue the primal approximation with $b = 0.9$ and in red Nudging with $\lambda = 0.4$.
• Figure 3: Entropy different $H\left[ \langle {\bm{p}}, {\bm{d}}\rangle^{-1} {\bm{p}} \odot {\bm{d}} \right] - H\left[ \langle {\bm{p}}^*, {\bm{d}}\rangle^{-1} {\bm{p}}^* \odot {\bm{d}} \right]$ for different values of $\lambda$. We report mean differences evaluated on OLMo-2 1B and 13B on MATH500.
• Figure 4: Illustration of the effect of the critical element size on the filled budget; a bigger critical element can leave a bigger unfilled gap.

Theorems & Definitions (11)

• lemma 1
• lemma 2
• lemma 3
• lemma 4
• lemma 5
• lemma 6
• proof
• proof
• proof
• proof