From 48bb900f4515490a84786c3ce8f5c5aca4011ea1 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Tue, 23 Jun 2026 21:31:15 +0800
Subject: [PATCH] revert F1: keep NaN-at-collapse in the answer-token renorm
 (it is the signal)

The prior commit "fixed" p_a/pmass into a softmax to avoid NaN at coherence collapse.
That was wrong for this codebase: renormalizing within allowed tokens discards the mass,
so a distribution built from ~zero mass and one from real mass look equally comparable
after renorm -- they are not (the mean of 10 != the mean of 130). p/pmass -> NaN at
collapse poisons that item's factor, which is the honest "do not compare" signal, not a
bug. softmax was a silent fallback fabricating a comparable-looking number from incoherent
output. Restored, with a comment so it does not get re-"fixed". (Over-trusted the external
reviewer here against the repo's own no-defensive-programming rule.)

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
---
 src/tinymfv/read.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/tinymfv/read.py b/src/tinymfv/read.py
index 8f6445e..269d16b 100644
--- a/src/tinymfv/read.py
+++ b/src/tinymfv/read.py
@@ -72,12 +72,14 @@ def read_items(model, tok, instr: Instrument, items: list[InstrItem], answer_ids
         enc = tok(texts, return_tensors="pt", padding=True, add_special_tokens=False).to(device)
         logits = model(**enc).logits[:, -1, :].float()        # [B, V] next-token
         logp = F.log_softmax(logits, dim=-1)
-        logp_a = logp[:, gid]                                 # [B, A] logprob on each answer token
-        pmass = logp_a.exp().sum(dim=-1)                      # [B] coherence check: mass on allowed tokens
-        # softmax over the allowed logprobs == p_a / pmass when pmass > 0, but NaN-safe: at coherence
-        # collapse pmass underflows to 0 and the divide would poison the whole profile with NaN; the
-        # softmax still returns a valid within-allowed distribution and pmass separately flags the drop.
-        p_norm = F.softmax(logp_a, dim=-1)                    # [B, A] within allowed
+        p_a = logp[:, gid].exp()                              # [B, A] prob on each answer token
+        pmass = p_a.sum(dim=-1)                               # [B] coherence check: mass on allowed tokens
+        # Renormalize within allowed. INTENTIONALLY NOT NaN-guarded: at full coherence collapse
+        # pmass -> 0 so p_norm -> NaN and poisons that item's factor. That is the honest signal, a
+        # distribution renormalized from ~zero mass is NOT comparable to one from real mass (the mean
+        # of 10 != the mean of 130), so it must not be silently turned into a comparable-looking
+        # number. NaN marks "do not compare". Do not "fix" this with a softmax/eps fallback.
+        p_norm = p_a / pmass[:, None]                         # [B, A] within allowed (NaN at collapse, by design)
         for j, it in enumerate(chunk):
             out.append({
                 "id": it.id, "frame": it.frame,