Add animation of lm head - #870
Open
klei22 wants to merge 5 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a new training-time visualization for tracking how per-token lm_head vector magnitudes evolve, via TensorBoard histogram logging plus an exported interactive HTML report for inspection outside TensorBoard.
Changes:
- Adds TensorBoard histogram logging for per-token
lm_headweight-vector L2 magnitudes. - Captures per-eval snapshots and exports an interactive Plotly-based HTML report.
- Adds a runnable demo script and documentation for viewing the new visualizations.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| train.py | Implements histogram logging, snapshot capture, and HTML export integrated into the metrics/training flow. |
| train_args.py | Adds CLI flags to enable/parameterize lm_head histogram logging and HTML export. |
| demos/README.md | Documents how to run the new demo and where to view the outputs. |
| demos/lm_head_vocab_histogram_demo.sh | Provides a short end-to-end example run enabling the new TensorBoard/HTML visualizations. |
Comments suppressed due to low confidence (3)
train.py:1518
- Same issue as above: prefer
key in self.model.transformer+self.model.transformer[key]instead of.get(...)when reading from the model’snn.ModuleDicttransformer container (seemodel.py:116).
lm_head = self.model.transformer.get(f"lm_head_{dataset_idx}", lm_head)
train.py:1441
- If
--lm_head_vocab_hist_html_pathis set to a filename without a directory (e.g.report.html),os.path.dirname(out_path)is empty andos.makedirs('')will raise. Create the directory only when the dirname is non-empty.
os.makedirs(os.path.dirname(out_path), exist_ok=True)
train.py:1442
- The snapshots JSON is inlined directly into a
<script>block. If any token string contains</script>(or similar), it can break out of the script context when viewing the report (HTML/JS injection risk). At minimum, escape</sequences before embedding.
payload = json.dumps(self.lm_head_hist_snapshots)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1390
to
+1396
| if ( | ||
| not self.args.tensorboard_log | ||
| or self.writer is None | ||
| or not self.args.log_lm_head_vocab_hist | ||
| or self.iter_num % self.args.log_lm_head_vocab_hist_interval != 0 | ||
| ): | ||
| return |
Comment on lines
+1398
to
+1405
| lm_head = getattr(self.model, "lm_head", None) | ||
| if self.args.training_mode == "multicontext" and hasattr(self.model, "transformer"): | ||
| try: | ||
| dataset_idx = self.args.multicontext_datasets.index(target_dataset) | ||
| except ValueError: | ||
| dataset_idx = 0 | ||
| lm_head = self.model.transformer.get(f"lm_head_{dataset_idx}", lm_head) | ||
|
|
Comment on lines
+1524
to
+1527
| for i, m in enumerate(magnitudes): | ||
| token_raw, token_display = self._get_vocab_label_parts(i) | ||
| vocab_data.append({"id": i, "magnitude": float(m), "token_raw": token_raw, "token_display": token_display}) | ||
| self.lm_head_hist_snapshots.append({ |
| logging_group.add_argument('--log_areq', default=True, action=argparse.BooleanOptionalAction, help='Log aReQ representation metric during validation') | ||
| logging_group.add_argument('--log_lm_head_vocab_hist', default=False, action=argparse.BooleanOptionalAction, help='Log TensorBoard histogram of per-token lm_head vector magnitudes over training') | ||
| logging_group.add_argument('--log_lm_head_vocab_hist_interval', default=100, type=int, help='Training-step interval for logging lm_head vocab magnitude histogram') | ||
| logging_group.add_argument('--export_lm_head_vocab_hist_html', default=False, action=argparse.BooleanOptionalAction, help='Export an interactive HTML report of final lm_head vocab-vector magnitudes') |
| Use the Histograms or Distributions tab and scrub over steps to animate the | ||
| change in lm_head vocab-vector magnitude distribution during training. | ||
|
|
||
| An interactive final-snapshot HTML is also written to: |
gkielian
approved these changes
Jul 28, 2026
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.
This pull request adds a new feature for visualizing the evolution of per-token
lm_headvector magnitudes during training, with both TensorBoard histogram logging and an interactive HTML report. It introduces new logging options, a demo script, and supporting code to enable, capture, and export these visualizations.New visualization and logging features:
--log_lm_head_vocab_histand related CLI options totrain_args.pyto enable logging of per-tokenlm_headvector magnitudes as TensorBoard histograms, control logging interval, and export an interactive HTML report._log_lm_head_vocab_magnitude_histogram,_capture_lm_head_hist_snapshot, and_export_lm_head_vocab_histogram_htmlmethods intrain.pyto log histogram data, capture snapshots for HTML, and generate an interactive Plotly-based HTML report. Methods are integrated into the training and metrics logging flow. [1] [2] [3] [4] [5]demos/lm_head_vocab_histogram_demo.sh, and updateddemos/README.mdwith instructions to run a sample job and view the new visualizations in TensorBoard and HTML. [1] [2]These changes make it easy to observe and analyze how the output embedding matrix (
lm_head) evolves during training, which can be useful for debugging and research on model representations.