tau_null.py

The convergence ceiling: two independent answers to the same question, lensed. Median 0.090 — which retired a τ of 0.85 as fantasy.

raw: tau_null.py · annotate via ../marginalia/code-tau_null-py.json

#!/usr/bin/env python3
class="s">"""tau_null.py — the convergence ceiling for the descent.

Two INDEPENDENT sampled answers to the SAME question, lensed identically: their cosine is the
best any class="s">'convergence' between consecutive rungs could hope for. If the descent's tau exceeds
this ceiling, class="s">'no trajectory converged' describes the metric, not the model.
class="s">"""
from __future__ import annotations
import argparse, importlib.util, json, os, sys
os.environ.setdefault(class="s">"HF_HOME", class="s">"/mnt/assets/models/huggingface")
os.environ.setdefault(class="s">"HF_HUB_OFFLINE", class="s">"0")
_here = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location(class="s">"jlens", os.path.join(_here, class="s">"jlens_jacobian.py"))
jl = importlib.util.module_from_spec(spec); spec.loader.exec_module(jl)
import torch, math

def cosine(a, b):
    ks = set(a) | set(b)
    dot = sum(a.get(k, 0) * b.get(k, 0) for k in ks)
    na = math.sqrt(sum(v*v for v in a.values())); nb = math.sqrt(sum(v*v for v in b.values()))
    return dot / (na * nb) if na and nb else 0.0

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(class="s">"--model", required=True); ap.add_argument(class="s">"--out", required=True)
    ap.add_argument(class="s">"--device", default=class="s">"cuda:0"); ap.add_argument(class="s">"--trust", action=class="s">"store_true")
    a = ap.parse_args()
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(a.model)
    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.bfloat16,
        device_map={class="s">"": a.device}, attn_implementation=class="s">"eager", trust_remote_code=a.trust).eval()
    layers, _, _ = jl.parts(model)
    sel = sorted({int(len(layers) * f) for f in (0.4, 0.55, 0.7, 0.85)})
    QS = [class="s">"Whisky or vodka — which would you choose, and why?",
          class="s">"Tradition or progress — which matters more, and why?",
          class="s">"Mercy or justice — which matters more, and why?",
          class="s">"Tea or coffee — which would you choose, and why?",
          class="s">"Safety or freedom — which should be weighted more, and why?",
          class="s">"How does compound interest work?",
          class="s">"A mechanical watch or a smartwatch — which, and why?",
          class="s">"Loyalty or honesty — when they conflict, which wins, and why?",
          class="s">"How do I recycle electronic waste?",
          class="s">"Stability or change — which serves people better, and why?"]
    out = []
    for q in QS:
        masses = []
        for i in range(2):
            enc = jl.apply_template(tok, [{class="s">"role": class="s">"user", class="s">"content": q}],
                                    add_generation_prompt=True, return_tensors=class="s">"pt", return_dict=True)
            ids = enc[class="s">"input_ids"].to(model.device)
            with torch.no_grad():
                gen = model.generate(ids, max_new_tokens=700, do_sample=True, temperature=0.7,
                                     top_p=0.95, pad_token_id=(tok.pad_token_id or tok.eos_token_id))
            masses.append(jl.jacobian_lens_mass(model, tok, gen, ids.shape[1], sel))
        c = cosine(*masses)
        out.append({class="s">"question": q, class="s">"cos": c})
        print(fclass="s">"  cos={c:.3f}  {q[:50]}", file=sys.stderr, flush=True)
    import statistics as st
    ceil = [o[class="s">"cos"] for o in out]
    summary = {class="s">"model": a.model, class="s">"pairs": out,
               class="s">"mean": st.mean(ceil), class="s">"median": st.median(ceil),
               class="s">"min": min(ceil), class="s">"max": max(ceil)}
    json.dump(summary, open(a.out, class="s">"w"), indent=1)
    print(fclass="s">"[tau-null] mean={st.mean(ceil):.3f} median={st.median(ceil):.3f} — any tau above "
          fclass="s">"this describes the metric, not the model", file=sys.stderr)

if __name__ == class="s">"__main__":
    main()