feat: add margin, max_prob_raw, entropy, and top3_mass confidence methods - #4
Open
Jakob Drachmann Havtorn (JakobHavtorn) wants to merge 4 commits into
Conversation
- margin: p_max - p_second (top-2 probability gap) - max_prob_raw: raw softmax probability without V-normalization Both improve class separation when V is large (e.g. 8193 for TDT v3)
- entropy: normalized Shannon entropy using log_softmax, mapped to [0,1] - top3_mass: sum of top-3 softmax probabilities - Fix margin lambda to use [..., 0] instead of [:, :, 0] for 2D tensor compat - Override entropy dispatch to use our custom implementation instead of NeMo's broken built-in variants (entropy_gibbs_lin etc. produce inf/values >1)
…er.py Co-authored-by: black <black@psf.dev>
Signed-off-by: Jakob D. Havtorn <jdh@corti.ai>
Copilot started reviewing on behalf of
Jakob Drachmann Havtorn (JakobHavtorn)
September 7, 2026 12:50
View session
There was a problem hiding this comment.
🟡 Changes recommended
As implemented, new measures conflict with existing test invariants (uniform distribution expected to map to 0 for all measures) and the entropy config semantics appear to be a backward-incompatible API change that needs clarification/adjustment.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds additional confidence metrics to NeMo ASR’s per-frame confidence machinery by extending get_confidence_measure_bank() and updating method-name selection in ConfidenceMethodMixin._init_confidence_method().
Changes:
- Added new confidence measures:
margin,max_prob_raw,top3_mass, and a new Shannon-styleentropyentry in the measure bank. - Expanded supported method names and updated
_init_confidence_method()branching to recognize the new options. - Updated
ConfidenceMethodConfigdocumentation for some (but not all) newly supported method names.
File summaries
| File | Description |
|---|---|
| nemo/collections/asr/parts/utils/asr_confidence_utils.py | Adds new confidence measure implementations and updates method-name routing to enable selecting them via config. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 6
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+241
to
+244
| # margin: p_max - p_second (top-2 gap, no V-normalization) | ||
| confidence_measure_bank["margin"] = lambda x, v, t: ( | ||
| x.topk(2, dim=-1)[0][..., 0].exp() - x.topk(2, dim=-1)[0][..., 1].exp() | ||
| ) |
Comment on lines
+245
to
+248
| # max_prob_raw: raw softmax probability (no V-normalization) | ||
| confidence_measure_bank["max_prob_raw"] = lambda x, v, t: ( | ||
| x.max(dim=-1)[0].exp() if t == 1.0 else (x.max(dim=-1)[0] * t).exp() | ||
| ) |
Comment on lines
+269
to
+270
| # top3_mass: sum of top-3 softmax probabilities, in [0, 1] | ||
| confidence_measure_bank["top3_mass"] = lambda x, v, t: (torch.softmax(x * t, dim=-1).topk(3, dim=-1)[0].sum(-1)) |
Comment on lines
+271
to
+278
| # entropy: 1 - H(p)/log(V), normalized Shannon entropy confidence in [0, 1] | ||
| confidence_measure_bank["entropy"] = lambda x, v, t: ( | ||
| 1.0 | ||
| + (torch.nn.functional.log_softmax(x * t, dim=-1).exp() * torch.nn.functional.log_softmax(x * t, dim=-1)).sum( | ||
| -1 | ||
| ) | ||
| / math.log(v) | ||
| ) |
Comment on lines
341
to
345
| elif confidence_method_cfg.name == "entropy": | ||
| measure_name = '_'.join( | ||
| [confidence_method_cfg.name, confidence_method_cfg.entropy_type, confidence_method_cfg.entropy_norm] | ||
| ) | ||
| measure_name = "entropy" | ||
| elif confidence_method_cfg.name == "entropy_renyi_exp": | ||
| measure_name = "entropy_renyi_exp" | ||
| else: |
Comment on lines
54
to
61
| Args: | ||
| name: The method name (str). | ||
| Supported values: | ||
| - 'max_prob' for using the maximum token probability as a confidence. | ||
| - 'max_prob_raw' for the raw softmax probability (no V-normalization). | ||
| - 'margin' for the difference between top-2 probabilities (p_max - p_second). | ||
| - 'entropy' for using a normalized entropy of a log-likelihood vector. | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds 4 new confidence measures to NeMo's
confidence_measure_bankinasr_confidence_utils.py, stacked on top of #3.New confidence methods
marginp_max - p_secondmax_prob_rawexp(x_max)entropy1 - H(p) / log(V)top3_massΣ_{i≤3} p_iChanges
All changes are in
nemo/collections/asr/parts/utils/asr_confidence_utils.py:confidence_measure_bank_init_confidence_methodto accept the new method namesNotes
marginandentropyandtop3_masscannot be post-hoc temperature-retuned (they depend on the full token distribution)max_probandmax_prob_rawsupport post-hoc alpha recomputationentropy_gibbs_lin, etc.) produceinfand values >1 — ourentropyis reimplemented from scratch withlog_softmaxStacked on