Skip to content

Add Spilled Energy uncertainty estimator - #466

Open
neuezeldaa wants to merge 1 commit into
IINemo:mainfrom
neuezeldaa:spilled-energy
Open

neuezeldaa wants to merge 1 commit into
IINemo:mainfrom
neuezeldaa:spilled-energy

Conversation

@neuezeldaa

Copy link
Copy Markdown

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

@ArtemVazh ArtemVazh left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This bypasses WhiteboxModel.tokenize() and therefore does not apply the chat template when instruct=True. Generation and the energy pass can score different prompts.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants