Adaptive Token Biaser in LLMs

• Adaptive Token Biaser (ATBias) is a decoding-time mechanism that selectively biases tokens associated with key entities to update in-context knowledge.
• It employs targeted logit biasing with n-gram and Jaccard similarity matching to boost new factual information while minimizing computational overhead.
• Empirical evaluations show ATBias outperforms global perturbation methods in stubborn knowledge scenarios with near-baseline latency.

The Adaptive Token Biaser (ATBias) is a decoding-time mechanism for LLMs designed for efficient, high-fidelity in-context knowledge editing. Unlike prior approaches that globally perturb the generation process, ATBias selectively biases the generation probability of tokens closely associated with knowledge-bearing entities. This yields accurate knowledge updates, preserves fluency, and does so with minimal computational overhead compared to previous contrastive-decoding schemes. ATBias can be integrated into any left-to-right LLM decoding pipeline as a "drop-in" wrapper and is shown to outperform established methods in various model-editing benchmarks, especially in challenging or "stubborn" knowledge scenarios (Bi et al., 2024).

1. Motivation and Context

Pretrained LLMs encode vast quantities of world knowledge parametrically, but this information becomes outdated as real-world facts change. Directly updating these models via retraining or fine-tuning is not scalable. In-context editing (ICE) emerged as a practical solution, whereby newly edited facts are included in the prompt, enabling the model to generate content in accordance with the latest information without altering the underlying parameters. While ICE is efficient, its effectiveness diminishes for hard-to-edit ("stubborn") facts because standard decoding is dominated by the original, outdated parametric memory.

Prior approaches such as DeCK [Bi et al. 2024] enhanced ICE by re-scoring all tokens of generated sequences using both pre- and post-edit prompts, thereby amplifying the influence of edits. However, this results in substantial latency increases (up to 2×), degrades grammaticality due to over-perturbation of non-entity tokens, and is computationally inefficient (Bi et al., 2024).

2. Core Mechanism

ATBias departs from global perturbation strategies by targeting only those tokens directly related to key entities—the constituents of either new or outdated knowledge. At every decoding step $t$, ATBias operates as follows:

1. Vocabulary Filtering: The candidate set is reduced to tokens with both high rank ($k$-largest probabilities) and sufficiently high probability (within fraction $\alpha$ of the max). Formally, $V_\text{head} = V_\text{prob} \cap V_\text{rank}$, where $V_\text{prob}$ and $V_\text{rank}$ are probability and rank filters, respectively.
2. String Decoding and Entity Matching: Each candidate's byte-level string is decoded, and n-gram + Jaccard similarity is computed between this string and each entity in the "new" ($E_\text{new}$) or "parametric" ($E_\text{para}$) entity sets.
3. Logit Biasing: Tokens matching new entities receive a positive bias ($\Delta^+$), those matching parametric entities receive a negative bias ($\Delta^-$). The adjusted logit for candidate $k$0 is $k$1, where $k$2 aggregates the positive and negative similarities, weighted by coefficients $k$3, $k$4 and the mean candidate probability $k$5.
4. Renormalization and Sampling: The softmax is recomputed over the filtered set, yielding updated generation probabilities.

This selective reweighting ensures that only entity-bearing tokens are influenced by editing constraints, preventing collateral disturbances to function words, syntax tokens, or unrelated content and significantly reducing compute requirements (Bi et al., 2024).

3. Mathematical Formulation

Let $k$6 be the model's logits at position $k$7 and $k$8. The vocabulary candidate set is filtered:

$k$9

where $\alpha$0 is the $\alpha$1 largest probability. The final candidate set is $\alpha$2.

For string-based entity matching:

$\alpha$3

where $\alpha$4 denotes the set of all character $\alpha$5-grams for token $\alpha$6 and entity $\alpha$7.

The logit is adjusted:

$\alpha$8

where $\alpha$9. Only $V_\text{head} = V_\text{prob} \cap V_\text{rank}$0 are assigned nonzero probabilities after the renormalized softmax.

4. Algorithmic Workflow

The per-step ATBias procedure is:

1. Compute raw token logits and probabilistic filter.
2. Identify top-$V_\text{head} = V_\text{prob} \cap V_\text{rank}$1 and high-probability tokens to form $V_\text{head} = V_\text{prob} \cap V_\text{rank}$2.
3. Calculate $V_\text{head} = V_\text{prob} \cap V_\text{rank}$3 over $V_\text{head} = V_\text{prob} \cap V_\text{rank}$4.
4. For each candidate $V_\text{head} = V_\text{prob} \cap V_\text{rank}$5:
   • Decode the token.
   • For all $V_\text{head} = V_\text{prob} \cap V_\text{rank}$6 and $V_\text{head} = V_\text{prob} \cap V_\text{rank}$7, compute $V_\text{head} = V_\text{prob} \cap V_\text{rank}$8-gram+Jaccard similarity.
   • Apply positive or negative bias to $V_\text{head} = V_\text{prob} \cap V_\text{rank}$9 if matches are found.
5. Renormalize logits over $V_\text{prob}$0.
6. Sample or select the next token.

All entity matching, bias application, and softmax are restricted to a small set $V_\text{prob}$1 (typically $V_\text{prob}$2), conserving computational resources (Bi et al., 2024).

5. Implementation and Efficiency

ATBias is deployed as a decoding wrapper, introducing no changes to LLM architecture or weights. Entity extraction for new and old facts is performed offline and cached. At inference, only lookups and $V_\text{prob}$3 per-candidate computations are required. Empirical measurements show that latency is nearly indistinguishable from vanilla decoding (1.01× baseline for LLaMA2-7B-chat; 36.19 ms/token vs. 36.03 ms/token), whereas prior methods such as DeCK introduce ≈2× latency (69.99 ms/token) (Bi et al., 2024).

Integration requires no modification of ICE methods, model internals, or external prompts, making ATBias compatible with a wide range of LLMs and editing pipelines.

6. Empirical Evaluation

ATBias was validated on standard knowledge-editing benchmarks:

• CounterFact: One-hop editing (≈3k samples).
• MQuAKE-3k, MQuAKE-2002, MQuAKE-hard: Multi-hop counterfactual QA.
• Challenging ("stubborn") subsets where ICE methods fail on >67% of instances.

Results indicate consistent state-of-the-art performance:

Model & Task	Baseline (IKE/MeLLo)	+DeCK (%)	+ATBias (%)	Δ (ATBias - DeCK)
LLaMA2-7B-chat, MQuAKE-3k, IKE	91.3	91.3	93.1	+1.8
LLaMA2-13B-chat, MQuAKE-3k, IKE	84.6	84.6	89.7	+5.1
LLaMA2-7B-chat, MQuAKE-3k, MeLLo	43.1	–	54.3	+32.3

On "stubborn" cases, ATBias yields 28–57% accuracy, outperforming baselines and prior model-editing methods (ROME, DeCK). Ablations confirm performance is optimal at $V_\text{prob}$4 for n-grams, $V_\text{prob}$5, $V_\text{prob}$6, $V_\text{prob}$7, and $V_\text{prob}$8. Using both probability and rank filtering improves results by 10–20 points compared to either alone (Bi et al., 2024).

7. Discussion and Extensions

ATBias's principal strength is the precision with which logit manipulation is focused on knowledge-bearing tokens, avoiding unnecessary disturbance of language generation mechanics. The method is computationally negligible in terms of overhead, flexible as a decoding-time intervention, and compatible with any ICE-based knowledge editing pipeline.

Limitations include dependence on robust entity extraction and string matching, with decreased efficacy for rare entities, highly granular tokenizations, or non-Latin scripts. Only limited model and pipeline generalization has been tested (LLaMA2, Mistral-7B, two ICE pipelines), so broad external validation remains necessary.

Potential generalizations entail adapting the n-gram/Jaccard matcher to subword or byte-level similarity for multilingual models, expanding targeting to relation or event tokens, or combining ATBias with retrieval-based token gating for even sharper focus.

In summary, ATBias establishes a highly effective, minimally intrusive paradigm for LLM editing, achieving up to +32.3 percentage point gains over prior ICE decoding methods in multi-hop knowledge-intensive QA while maintaining essentially baseline-level generation latency (Bi et al., 2024).