re-normalize probability distribution when clip_logits_epsilon is used.

[entrpn]

Description

Fixes b/534432708

This PR renormalizes the probability distribution when clamping using clip_logits_epsilon for calculating kl divergence. Also enables using clip_logits_epsilon when run_hf_model=True.

Tests

Used llama3.1-8b with clip_logits_epsilon=1e-8.

--- Prompt: I love to ---
INFO:absl:
--- MaxText model top 10 tokens ---
[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer TokenizersBackend. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1518       | see                  | 14.3817    |
| 1373       | read                 | 13.2008    |
| 6865       | hear                 | 12.9473    |
| 5944       | travel               | 12.9390    |
| 4394       | cook                 | 12.8764    |
| 3139       | grow                 | 12.8584    |
| 8343       | eat                  | 12.8224    |
| 1005       | use                  | 12.6565    |
| 3821       | watch                | 12.4830    |
| 3137       | talk                 | 12.3774    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1518       | see                  | 14.3889    |
| 1373       | read                 | 13.2144    |
| 5944       | travel               | 12.9658    |
| 6865       | hear                 | 12.9652    |
| 4394       | cook                 | 12.8976    |
| 3139       | grow                 | 12.8762    |
| 8343       | eat                  | 12.8403    |
| 1005       | use                  | 12.6674    |
| 3821       | watch                | 12.4977    |
| 3137       | talk                 | 12.3810    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 80.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 3.56e-04
INFO:absl:Per-token KL Divergences: 
['5.8001e-05', '4.1124e-04', '3.5516e-04', '5.9813e-04']
INFO:absl:
Max KL divergence for a single token in the set: 5.98e-04
INFO:absl:
--- Prompt: Today is a ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1938       | day                  | 13.4890    |
| 1695       | good                 | 13.2722    |
| 1633       | very                 | 13.2558    |
| 2466       | big                  | 13.1216    |
| 2294       | great                | 13.0514    |
| 3361       | special              | 12.8912    |
| 502        | new                  | 11.9839    |
| 6366       | beautiful            | 11.5664    |
| 12703      | sad                  | 11.3384    |
| 13560      | holiday              | 11.3374    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1938       | day                  | 13.4852    |
| 1695       | good                 | 13.2647    |
| 1633       | very                 | 13.2567    |
| 2466       | big                  | 13.1190    |
| 2294       | great                | 13.0474    |
| 3361       | special              | 12.8965    |
| 502        | new                  | 11.9764    |
| 6366       | beautiful            | 11.5652    |
| 13560      | holiday              | 11.3234    |
| 12703      | sad                  | 11.3214    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 80.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 4.35e-04
INFO:absl:Per-token KL Divergences: 
['5.8001e-05', '5.0937e-04', '5.6535e-04', '6.0815e-04']
INFO:absl:
Max KL divergence for a single token in the set: 6.08e-04
INFO:absl:
--- Prompt: What is the ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 6811       | difference           | 11.0520    |
| 1888       | best                 | 10.7357    |
| 1455       | most                 | 10.6928    |
| 1925       | main                 | 10.0193    |
| 7580       | purpose              | 9.9071     |
| 3560       | role                 | 9.8912     |
| 6864       | capital              | 9.5586     |
| 7438       | meaning              | 9.5398     |
| 12939      | importance           | 9.4890     |
| 5536       | impact               | 9.4724     |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 6811       | difference           | 11.0629    |
| 1888       | best                 | 10.7349    |
| 1455       | most                 | 10.6874    |
| 1925       | main                 | 10.0299    |
| 7580       | purpose              | 9.9002     |
| 3560       | role                 | 9.8755     |
| 6864       | capital              | 9.5623     |
| 7438       | meaning              | 9.5414     |
| 12939      | importance           | 9.4840     |
| 5536       | impact               | 9.4479     |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 100.0                |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 3.17e-04
INFO:absl:Per-token KL Divergences: 
['5.8001e-05', '5.2395e-04', '2.8289e-04', '4.0487e-04']
INFO:absl:
Max KL divergence for a single token in the set: 5.24e-04

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[aireenmei]

Thanks for the fix!