Retrieval-Augmented Generation: External Memory for Generation

This witness uses the same finite corpus as the demo. The middle assertion is the core lesson: with $k=2$, the correct current chunk is retrieved, so coverage is true, but the selected answer remains wrong because the stale chunk has more retrieval weight.

import numpy as np

ANSWERS = ["30", "14", "45", "unknown"]
TRUE_ANSWER = "30"

def norm(v):
    v = np.asarray(v, dtype=float)
    return v / np.linalg.norm(v)

query = norm([1.0, 0.0])

docs = [
    {
        "id": "old_policy",
        "text": "2024 handbook: standard returns are accepted within 14 days.",
        "embedding": norm([0.98, 0.20]),
        "fresh": 0,
        "supports_answer": "14",
        "gen": {"30": 0.05, "14": 0.86, "45": 0.02, "unknown": 0.07},
    },
    {
        "id": "current_policy",
        "text": "April 2026 policy: standard returns are accepted within 30 days.",
        "embedding": norm([0.92, 0.39]),
        "fresh": 1,
        "supports_answer": "30",
        "gen": {"30": 0.82, "14": 0.05, "45": 0.03, "unknown": 0.10},
    },
    {
        "id": "forum_exception",
        "text": "A customer once returned a jacket after 45 days during a promotion.",
        "embedding": norm([0.86, -0.51]),
        "fresh": 0,
        "supports_answer": "45",
        "gen": {"30": 0.10, "14": 0.05, "45": 0.75, "unknown": 0.10},
    },
    {
        "id": "shipping",
        "text": "Express shipping arrives in 2 business days.",
        "embedding": norm([0.30, 0.95]),
        "fresh": 1,
        "supports_answer": "unknown",
        "gen": {"30": 0.15, "14": 0.10, "45": 0.05, "unknown": 0.70},
    },
    {
        "id": "warranty",
        "text": "Electronics warranty coverage lasts one year.",
        "embedding": norm([0.10, -0.99]),
        "fresh": 1,
        "supports_answer": "unknown",
        "gen": {"30": 0.12, "14": 0.08, "45": 0.05, "unknown": 0.75},
    },
]

def softmax(xs, tau=1.0):
    xs = np.asarray(xs, dtype=float) / tau
    xs = xs - xs.max()
    ps = np.exp(xs)
    return ps / ps.sum()

def rag(k=2, tau=0.25, fresh_bonus=0.0, noise_sigma=0.0, seed=0):
    rng = np.random.default_rng(seed)
    rows = []
    for d in docs:
        cosine = float(query @ d["embedding"])
        noise = float(rng.normal(0.0, noise_sigma)) if noise_sigma else 0.0
        score = cosine + fresh_bonus * d["fresh"] + noise
        rows.append({**d, "cosine": cosine, "noise": noise, "score": score})

    rows.sort(key=lambda r: (-r["score"], r["id"]))
    top = rows[:k]
    weights = softmax([r["score"] for r in top], tau=tau)

    rag_probs = {a: 0.0 for a in ANSWERS}
    for w, r in zip(weights, top):
        for a in ANSWERS:
            rag_probs[a] += float(w) * r["gen"][a]

    answer = max(ANSWERS, key=lambda a: (rag_probs[a], -ANSWERS.index(a)))
    cited = max(top, key=lambda r: r["gen"][answer] * float(weights[top.index(r)]))

    return {
        "retrieved_ids": [r["id"] for r in top],
        "weights": {r["id"]: float(w) for r, w in zip(top, weights)},
        "coverage": any(r["supports_answer"] == TRUE_ANSWER for r in top),
        "evidence_mass": sum(float(w) for r, w in zip(top, weights) if r["supports_answer"] == TRUE_ANSWER),
        "rag_probs": rag_probs,
        "answer": answer,
        "p_correct": rag_probs[TRUE_ANSWER],
        "correct": answer == TRUE_ANSWER,
        "cited_source": cited["id"],
        "citation_supports_answer": cited["supports_answer"] == answer,
    }

baseline = rag(k=1, tau=0.25)
covered_but_wrong = rag(k=2, tau=0.25)
fresh_rerank = rag(k=2, tau=0.25, fresh_bonus=0.12)

assert baseline["retrieved_ids"] == ["old_policy"]
assert baseline["coverage"] is False
assert baseline["answer"] == "14"
assert baseline["citation_supports_answer"] is True
assert baseline["correct"] is False

assert covered_but_wrong["retrieved_ids"] == ["old_policy", "current_policy"]
assert covered_but_wrong["coverage"] is True
assert covered_but_wrong["answer"] == "14"
assert covered_but_wrong["p_correct"] < 0.50
assert 0.0 < covered_but_wrong["evidence_mass"] < 1.0

assert fresh_rerank["retrieved_ids"][0] == "current_policy"
assert fresh_rerank["coverage"] is True
assert fresh_rerank["answer"] == "30"
assert fresh_rerank["p_correct"] > covered_but_wrong["p_correct"]
assert fresh_rerank["cited_source"] == "current_policy"