CAD-GPT2 — Context-Aware Decoding for GPT-2

A minimal, self-contained implementation of Context-Aware Decoding (CAD) on top of Hugging Face's GPT-2, following Shi et al., "Trusting Your Evidence: Hallucinate Less with Context-aware Decoding" (2023).

Standard language models often ignore the context you give them and fall back on what they memorized during pre-training — which is exactly how they end up confidently contradicting a document you just handed them. CAD counteracts this at decoding time only, with no retraining, by contrasting the model's predictions with and without the context and amplifying whatever the context actually changed.

The idea in one line

At each generation step, CAD adjusts the next-token logits so the model is pushed toward tokens that the context makes more likely relative to the model's prior:

adjusted_logits = (1 + α) · logits(y | context, input) − α · logits(y | input)

• logits(y | context, input) — the model conditioned on the context and the prompt.
• logits(y | input) — the same model conditioned on the prompt only (its context-free prior).
• α ≥ 0 — the contrastive strength. α = 0 recovers ordinary greedy/sampling decoding; larger α makes the model lean harder on the context.

The difference term subtracts out what the model would have said anyway, leaving the context's contribution — and then over-weights it. The result is generation that "trusts its evidence" and hallucinates less on knowledge-conflict tasks.

What's in this repo

File	Description
CAD_GPT2.ipynb	The full implementation and a runnable demo.

The core is a single class, CADGPT2LMHeadModel, which subclasses GPT2LMHeadModel and overrides just enough of the generation machinery to apply the contrastive adjustment:

• set_context(text) — tokenizes and stores the context to condition on for the next generation.
• prepare_inputs_for_generation(...) — on the first step, prepends the context tokens to the input and fixes up the attention mask and position IDs so the context is attended to correctly.
• forward(...) — runs the model twice (with and without context), combines the two logit distributions using the CAD formula above, and returns the adjusted logits. After the first step it defers to the cached-KV fast path.

Because it plugs into the standard .generate() API, you get sampling, beam search, temperature, etc. for free.

Usage

from transformers import AutoTokenizer

# Load the CAD-enabled model (see the notebook for the class definition)
model = CADGPT2LMHeadModel.from_pretrained("gpt2", alpha=0.5)

# Give it a context that conflicts with what GPT-2 "knows"
context = "Argentina won the World Cup in 1978, 1986, and 2022."
model.set_context(context)

question = "How many times has Argentina won the World Cup?"
input_ids = model.tokenizer(question, return_tensors="pt").input_ids

output_ids = model.generate(input_ids=input_ids, max_length=80, do_sample=True)
print(model.tokenizer.decode(output_ids[0], skip_special_tokens=True))

Set alpha=0 to compare against a vanilla GPT-2 baseline on the same prompt.

Requirements

pip install torch transformers

Any recent transformers with the GPT-2 head (GPT2LMHeadModel) will work. A GPU is optional — GPT-2 small runs fine on CPU for short generations.

Notes & limitations

• The contrastive forward pass does an extra model call per step (context vs. no-context), so generation is roughly 2× the compute of plain decoding. This is inherent to CAD.
• set_context conditions a single generation; call it again before the next one.
• Best α is task-dependent. The paper reports gains across summarization and knowledge-conflict QA; values around 0.3–0.5 are a reasonable starting point.
• Built on GPT-2 for clarity — the same override pattern transfers to any *LMHeadModel.

Reference

Weijia Shi, Xiaochuang Han, Mike Lewis, Yulia Tsvetkov, Luke Zettlemoyer, Scott Wen-tau Yih. Trusting Your Evidence: Hallucinate Less with Context-aware Decoding. arXiv:2305.14739, 2023.