Skip to content

Fix EOS confidence suppression at nonzero temperature - #139

Open
primorLee wants to merge 1 commit into
ML-GSAI:mainfrom
primorLee:fix/eos-remasking-confidence
Open

Fix EOS confidence suppression at nonzero temperature#139
primorLee wants to merge 1 commit into
ML-GSAI:mainfrom
primorLee:fix/eos-remasking-confidence

Conversation

@primorLee

Copy link
Copy Markdown

Summary

  • suppress both EOS and EoT in the logits used to compute remasking confidence
  • keep token sampling unchanged, so EOS/EoT can still be proposed normally
  • add an order-sensitive regression test for the nonzero-temperature path

Problem

confidence_eos_eot_inf currently writes the EOS suppression to logits_with_noise after x0 has already been sampled. At nonzero temperature, logits_with_noise is a separate tensor, while low-confidence remasking computes probabilities from logits. EOS therefore retains its original confidence and can be transferred too early.

The zero-temperature path hides the defect because add_gumbel_noise returns logits itself, so the two names alias the same tensor.

Validation

  • regression test fails on main: expected the ordinary token to transfer first, but EOS is fixed first
  • python -m pytest tests/test_generate.py -q (1 passed)
  • Python 3.12 compile check on Linux
  • git diff --check

Fixes #137

@johnrunjiali-svg

Copy link
Copy Markdown

Thanks for the fix!

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.

confidence_eos_eot_inf does not correctly suppress EOS confidence

2 participants