Methodology

We present here a methodology for biasing model’s decoding with respect to a list of token ids. To preserve the written fluency, we use model uncertainty to scale the intervention strength on the model’s probability so that when the model is certain about what to produce the intervention is minimum. To do that, we employ Shannon’s Entropy $H(p)$ to quantify the model uncertainty:

$$H(\mathbf{p})=-\sum _i{p}_i\log p_i$$

where $\mathbf{p}$ is a vector of probability the size of the model’s vocabulary $|V_{model}|$ obtained by soft-maxing the last layer model’s logits $\mathbf{z}$, i.e.:

$$\mathbf{p}=\text{SoftMax}(\mathbf{z})$$

We use the entropy to define a smooth gating mechanism $\varphi (\mathbf{p})\in [0,1]$:

$$\varphi (\mathbf{p})=\sigma (\frac{H(\mathbf{p})-\tau }{s})$$

where $\tau$ acts as a threshold for the minimum entropy needed for the intervention to start, $s$ is a smoothing factor, and $\sigma$ is the sigmoid function.

Then we identified two possible scenarios, one where we want to encourage (positive biasing) the production of a list of ids, and one where we want to discourage (negative biasing) the production of a list of ids.

Positive Biasing

Given a vocabulary of green tokens $G$ we want to bias towards, we employ a soft-mixture of probability gated by our entropy-gate $\varphi$:

$${\mathbf{p}}^{\prime }=(1-\alpha )\mathbf{p}+\alpha \mathbf{q}$$

where $\mathbf{p}$ are the softmaxed logits produced at a certain decoding step by the model, $\mathbf{q}$ is a biasing prior distribution that allocates probability to each $g\in G$, and $\alpha$ represents the intervention strength calculates as the maximum allowed strength ${\alpha }_{max}$ gated by the entropy gate $\varphi$:

$$\alpha =\varphi (\mathbf{p}){\alpha }_{max}$$

Choosing a prior $\mathbf{q}$

Uniform $\mathbf{q}$

$$q_i=\{\begin{array}{ll}\frac{1}{|G|}& i\in G\\ 0,& \text{otherwise}\end{array}$$

This spreads the steering mass equally across all the green tokens in the selected vocabulary . Works well if is small, otherwise each token in $G$ gets a tiny bump in mass, this has then little effect on the generation, especially if greedy decoding.

Re-normalizing model distribution on $G$

$$q_i=\{\begin{array}{ll}\frac{p_i}{{\sum }_{j\in G}{p}_j}& i\in G\\ 0,& \text{otherwise}\end{array}$$

This re-normalizes the distribution of probability that the model assigned to the set of green tokens in $G$, over themselves. This is like if the model only had the green tokens to choose from.

This is “smarter” than uniform, because increases the $\alpha$ of vocabulary tokens dynamically based on how likely the model already wanted to generate them. However, if no vocabulary tokens are in the high-probability region it can have very little effect in biasing.

Top-K Green Tokens

$$q_i=\{\begin{array}{ll}\frac{1}{|K|}& i\in \text{Top-}K\text{ green tokens}\\ 0,& \text{otherwise}\end{array}$$

Here we want to extract the model’s probability for the green tokens, pick the top-$K$ candidates and create a $q$ prior with uniform distribution over just those $K$.

This is similar to the uniform distribution method, but instead of steering all the green tokens in $G$, we just steer the top-$K$ candidates. This reduces the number of tokens where $\alpha$ is spreaded to.

Negative Biasing

To avoid normalization artifacts or negative probabilities, negative biasing is performed directly in logit space. Given the model’s logits $\mathbf{z}$, we apply a penalty:

$${\mathbf{z}}^{\prime }=\mathbf{z}-\lambda \mathbf{r}$$

where $\mathbf{r}$ is a vector of size $|V_{model}|$ that shapes where the penalty is applied with respect to a red token vocabulary $R$ we want to discourage, and $\lambda$ is a penalty term calculated as the maximum penalty ${\lambda }_{max}$ gated by $\varphi (\mathbf{p})$:

$$\lambda =\varphi (\mathbf{p})\cdot {\lambda }_{max}$$

How to choose $\mathbf{r}$

Uniform

$$r_i=\{\begin{array}{ll}\frac{1}{|R|}& \text{if }i\in R\\ 0& \text{otherwise}\end{array}$$

Unnormalized Uniform

$$r_i=\{\begin{array}{ll}1& \text{if }i\in R\\ 0& \text{otherwise}\end{array}$$

Top-K

$$r_i=\{\begin{array}{ll}\frac{1}{K}& \text{if }i\in \text{TopK}(R)\\ 0& \text{otherwise}\end{array}$$

Unnormalized Top-K

$$r_i=\{\begin{array}{ll}1& \text{if }i\in \text{TopK}(R)\\ 0& \text{otherwise}\end{array}$$

Re-Normalizing model Distribution in $R$

$$r_i=\{\begin{array}{ll}\frac{p_i}{{\sum }_{j\in R}{p}_j}& \text{if }i\in R\\ 0& \text{otherwise}\end{array}$$

the effect is that the total penalty $\lambda$ is conserved over $R$ and allocated more towards those red tokens that the model currently prefers.