File size: 10,281 Bytes
1802505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | #!/usr/bin/env python3
"""
needle_score.py -- turn a set of CP_EVAL_PPL CSVs into one tidy retrieval row.
Reads, for one (model, T, cache): the per-(variant, block) CSVs produced by the
sweep, averages mean_nll over the ANSWER SPAN only, and decides significance.
SIGNIFICANCE IS ON RAW NLL, NOT ON THE SCORE
The score's denominator is cache-specific, so a score-space threshold cannot be
compared across caches. Rule:
retrieval present at p_frac iff nll(p_frac) < nll_null - 2*se_combined
with se from the B blocks (the only available spread -- CP_EVAL_PPL emits
count/mean_nll and no variance). `null` is therefore measured PER CACHE,
because nll_null itself moves with the RoPE cache.
THE SCORE IS DERIVED AND GUARDED
score(p) = (nll_absent - nll(p)) / (nll_absent - nll_floor), floor = p1
Marked UNDEFINED (not clipped, not extrapolated) when the denominator is below
MIN_DENOM_NATS -- a small denominator makes the ratio unstable and able to leave
[0,1] legitimately. Raw NLLs are always reported; the score never stands alone.
OFF-BY-ONE
The eval reports NLL at the INPUT position, so the answer token at trial offset
`a` is scored on CSV row `a-1`. We read the manifest's `csv_positions`, never
`answer_offsets`.
Usage
python3 needle_score.py --needle-dir needles --csv-dir needle_csv --T 32768 \
--model 114m_run1 --cache ntk --factor 8 --head-dim 64 --rope-theta 500000 \
--out needle_results.csv
"""
import argparse, csv, json, math, os, sys
MIN_DENOM_NATS = 0.5 # pre-registered (plan A2); revisable only BEFORE C1 launches
# Minimum MATERIAL effect. Statistical significance alone is not retrieval: with many
# blocks and low block-to-block scatter, s.e. can shrink until a 0.01-nat difference
# clears 2*s.e. -- which is a measurable difference, not evidence the model used the
# needle. Retrieval requires clearing BOTH 2*s.e. and this absolute floor.
MIN_EFFECT_NATS = 0.1
P_FRACS = ["p0", "p0.25", "p0.5", "p0.75", "p1"]
FLOOR = "p1" # needle adjacent to the probe = the retrieval floor
def span_nll(csv_path, lo, hi):
"""Count-weighted mean nll over CSV rows [lo, hi). Returns (mean, tokens)."""
s = n = 0.0
seen = 0
with open(csv_path) as f:
for row in csv.DictReader(f):
p = int(row["pos"])
if lo <= p < hi:
c = float(row["count"])
if c > 0:
s += float(row["mean_nll"]) * c
n += c
seen += 1
if seen != hi - lo:
raise SystemExit(f"{csv_path}: found {seen} of the {hi-lo} answer-span rows "
f"[{lo},{hi}) -- wrong T, or a truncated CSV")
if n <= 0:
raise SystemExit(f"{csv_path}: answer span has count=0 everywhere")
return s / n, n
def variant_stats(needle_dir, csv_dir, T, variant, blocks):
"""Per-block span NLLs -> (mean, s.e., n_blocks, tokens_per_block)."""
vals, toks = [], None
for b in range(blocks):
man_p = os.path.join(needle_dir, f"needle_T{T}_{variant}_b{b}.json")
csv_p = os.path.join(csv_dir, f"ppl_needle_T{T}_{variant}_b{b}.csv")
if not os.path.exists(man_p):
raise SystemExit(f"missing manifest {man_p}")
if not os.path.exists(csv_p):
raise SystemExit(f"missing CSV {csv_p} (block {b} of {variant} did not run?)")
man = json.load(open(man_p))
lo, hi = man["csv_positions"] # NOT answer_offsets -- see docstring
v, n = span_nll(csv_p, lo, hi)
vals.append(v)
toks = n
m = sum(vals) / len(vals)
if len(vals) > 1:
var = sum((v - m) ** 2 for v in vals) / (len(vals) - 1)
se = math.sqrt(var / len(vals))
else:
se = float("nan") # 1 block -> no spread; flagged downstream
return m, se, len(vals), toks
def main():
global P_FRACS # --p-fracs may replace the module default (see below)
ap = argparse.ArgumentParser()
ap.add_argument("--needle-dir", required=True, help="where needle_gen.py wrote manifests")
ap.add_argument("--csv-dir", required=True, help="where the sweep wrote ppl_needle_*.csv")
ap.add_argument("--T", type=int, required=True)
ap.add_argument("--blocks", type=int, default=5)
ap.add_argument("--model", required=True, help="e.g. 114m_run1 / 48m_run200 / twinB")
ap.add_argument("--cache", required=True, help="ntk | yarn | pi | searched")
ap.add_argument("--factor", type=int, required=True, help="extension factor s")
ap.add_argument("--head-dim", type=int, default=64)
ap.add_argument("--rope-theta", type=int, default=500000)
ap.add_argument("--out", default="needle_results.csv")
ap.add_argument("--p-fracs", default=None,
help="comma-separated variant names overriding the default "
f"{P_FRACS}. Used to score a REFINED distance grid (extra "
"points bracketing a crossing). The FLOOR variant must be "
"present: it is the denominator of every score.")
args = ap.parse_args()
if args.p_fracs:
P_FRACS = [v.strip() for v in args.p_fracs.split(",") if v.strip()]
if FLOOR not in P_FRACS:
raise SystemExit(f"--p-fracs must include the floor variant '{FLOOR}' -- "
"it is the denominator of every score.")
need = P_FRACS + ["absent", "null"]
st = {v: variant_stats(args.needle_dir, args.csv_dir, args.T, v, args.blocks)
for v in need}
nll_absent, se_absent, nb, ntok = st["absent"][0], st["absent"][1], st["absent"][2], st["absent"][3]
nll_floor, se_floor = st[FLOOR][0], st[FLOOR][1]
nll_null, se_null = st["null"][0], st["null"][1]
denom = nll_absent - nll_floor
denom_ok = denom >= MIN_DENOM_NATS
new = not os.path.exists(args.out)
fields = ["model", "rope_theta", "head_dim", "T", "factor", "cache", "p_frac",
"nll", "se", "nll_null", "se_null", "nll_absent", "nll_floor",
"denom_nats", "score", "significant", "n_trials_tokens", "blocks", "notes"]
with open(args.out, "a", newline="") as fh:
w = csv.DictWriter(fh, fieldnames=fields)
if new:
w.writeheader()
print(f"{'variant':8s} {'nll':>9s} {'se':>7s} {'score':>8s} sig? notes")
for v in P_FRACS:
nll, se, _, _ = st[v]
# Significance on RAW NLL against the per-cache null, requiring BOTH
# statistical separation (2 s.e.) AND a material effect (MIN_EFFECT_NATS).
se_comb = math.sqrt((0 if math.isnan(se) else se) ** 2 +
(0 if math.isnan(se_null) else se_null) ** 2)
margin = max(2 * se_comb, MIN_EFFECT_NATS)
effect = nll_null - nll
notes = []
if denom_ok:
sig = effect > margin
score = (nll_absent - nll) / denom
if 2 * se_comb < MIN_EFFECT_NATS < effect * 10 and effect <= MIN_EFFECT_NATS:
notes.append(f"effect {effect:.3f} nats below the {MIN_EFFECT_NATS}-nat "
f"materiality floor despite small s.e.")
else:
# No dynamic range in the cell: the floor is not meaningfully below the
# ceiling, so NO retrieval claim is supportable either way. Report the
# raw NLLs and withhold both score and verdict rather than emitting a
# significant-looking difference of a few hundredths of a nat.
sig = None
score = None
notes.append(f"UNUSABLE: denom {denom:.3f} < {MIN_DENOM_NATS} nats "
f"(gate A4.1) -- score and verdict withheld")
if args.blocks < 3:
notes.append(f"blocks={args.blocks}: spread has <2 dof")
w.writerow(dict(model=args.model, rope_theta=args.rope_theta, head_dim=args.head_dim,
T=args.T, factor=args.factor, cache=args.cache, p_frac=v,
nll=f"{nll:.6f}", se=("" if math.isnan(se) else f"{se:.6f}"),
nll_null=f"{nll_null:.6f}",
se_null=("" if math.isnan(se_null) else f"{se_null:.6f}"),
nll_absent=f"{nll_absent:.6f}", nll_floor=f"{nll_floor:.6f}",
denom_nats=f"{denom:.6f}",
score=("" if score is None else f"{score:.4f}"),
significant=("" if sig is None else int(bool(sig))),
n_trials_tokens=int(ntok), blocks=args.blocks,
notes="; ".join(notes)))
verdict = "n/a" if sig is None else ("YES" if sig else " no")
print(f"{v:8s} {nll:9.4f} {(' n/a' if math.isnan(se) else f'{se:7.4f}')} "
f"{(' n/a' if score is None else f'{score:8.4f}')} "
f"{verdict:>3s} {'; '.join(notes)}")
# A4.1: only floor < absent is by construction. A distant p_frac sitting at the
# absent level is TOTAL CONTEXT ROT -- the finding, not a failure. Only flag
# values materially OUTSIDE [floor, absent] for inspection.
print()
print(f"floor({FLOOR})={nll_floor:.4f} absent={nll_absent:.4f} null={nll_null:.4f} "
f"denom={denom:.4f} nats")
if not denom_ok:
print(f" GATE A4.1 FAIL: floor is not below absent by >= {MIN_DENOM_NATS} nats -- "
f"the probe has no dynamic range on this model/length.")
tol = 2 * max(se_floor if not math.isnan(se_floor) else 0.0,
se_absent if not math.isnan(se_absent) else 0.0, 1e-9)
for v in P_FRACS:
nll = st[v][0]
if nll < nll_floor - tol:
print(f" FLAG {v}: below the floor by more than 2 s.e. -- distant needle beating "
f"adjacency is suspicious; inspect (not an auto-fail).")
if nll > nll_absent + tol:
print(f" FLAG {v}: above the no-information ceiling by more than 2 s.e. -- the "
f"needle appears to HURT; inspect (not an auto-fail).")
print(f"\nwrote {args.out}")
if __name__ == "__main__":
main()
|