Add Spilled Energy uncertainty estimator - #466
neuezeldaa wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
I think the implementation can be simplified by removing the additional EnergyCalculator forward pass. WhiteboxModel._ScoresProcessor replaces out.scores with normalized log-probabilities, but Hugging Face generation can also return unprocessed logits through output_logits=True. These remain available as out.logits.
GreedyProbsCalculator could therefore request output_logits=True and extract the sampled-token logits and log-partitions directly from the generation pass. This would use the exact logits produced during generation and avoid the additional compute and prompt-tokenization mismatch.
Please rerun the experiments with instruct: true, the simple_instruct TriviaQA dataset, and AlignScore as the generation-quality metric. The existing simple_instruct evaluation config should be used as the base, with only the estimator set extended for Spilled Energy.
The current continuation setup with instruct: false does not exercise the instruct path. The primary comparison should therefore use the standard simple-instruct protocol; the continuation experiment may remain as a separate ablation.
| tokenizer = model.tokenizer | ||
| ids = [] | ||
| for t in texts: | ||
| enc = tokenizer(t, add_special_tokens=True, return_attention_mask=False) |
There was a problem hiding this comment.
This bypasses WhiteboxModel.tokenize() and therefore does not apply the chat template when instruct=True. Generation and the energy pass can score different prompts.
Spilled Energy as a training-free, logits-only sequence-level estimator.
stat_calculators/energy.py — EnergyCalculator, producing energy_token_logits (raw logit of each sampled token), energy_lse (per-step log-partition, one entry longer) and energy_trailing_terminators.
estimators/spilled_energy.py — SpilledEnergy, with four variants (logit, marginal, spilled, scaled_spilled), three poolings (min/max/mean), and optional exclusion of trailing terminator tokens from the pooling window.
Four one-line registrations plus one entry in the default calculator list.
Two evaluation configs under examples/configs/, following the existing polygraph_eval_* layout.
14 unit tests