Add lm head metrics - #869
Open
klei22 wants to merge 3 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds richer instrumentation around lm_head magnitude statistics (raw and post-normalization) to support analysis of output-layer geometry during training—particularly for modular arithmetic grokking experiments—and wires those metrics through CSV/TensorBoard and the experiment monitoring tooling. It also introduces a new modular-addition grokking demo and a sweep YAML for running standardized experiments.
Changes:
- Added computation + logging/export of
lm_headrow-vector magnitude summary stats (raw + post-norm) intrain.py, plus new CLI flags to enable them. - Updated experiment monitoring / metrics parsing schemas to recognize the new metrics keys.
- Added a modular arithmetic grokking demo script and an exploration YAML for running sweeps.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| variations/norm_variations.py | Extends CappedHyperSphereNorm radius initialization to support hsnorm_radius. |
| train.py | Implements lm_head magnitude stat collection, TensorBoard logging, CSV export, and best-loss file export. |
| train_args.py | Adds CLI flags to enable raw and post-norm lm_head magnitude stat logging. |
| run_exploration_monitor.py | Extends the monitor’s metric schema to include the new lm_head magnitude stats. |
| optimization_and_search/run_experiments.py | Extends metric keys and parsing/backward compatibility to include the new stats. |
| explorations/modular_arithmetic_grokking.yaml | Adds a sweep config for modular addition grokking runs with lm_head stat logging enabled. |
| demos/modular_arithmetic_grokking_demo.py | Adds an end-to-end data prep/train/eval demo for modular addition grokking. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+211
to
+217
| # Determine radius initialization value | ||
| radius_init = None | ||
| if config.hsnorm_radius is not None: | ||
| self.radius = config.hsnorm_radius | ||
| else: | ||
| self.radius = math.sqrt(config.n_embd) | ||
| print(self.radius) |
Comment on lines
+1632
to
+1648
| vectors = [] | ||
| with torch.no_grad(): | ||
| for lm_head in self._iter_lm_head_modules(): | ||
| weight = lm_head.weight.detach() | ||
| if apply_lm_head_norm: | ||
| weight = self.raw_model.apply_lm_head_norm(weight) | ||
| vectors.append(weight.float()) | ||
| if not vectors: | ||
| return self._empty_lm_head_magnitude_stats() | ||
| magnitudes = torch.linalg.vector_norm(torch.cat(vectors, dim=0), ord=2, dim=1) | ||
| return { | ||
| "max": magnitudes.max().item(), | ||
| "avg": magnitudes.mean().item(), | ||
| "median": magnitudes.median().item(), | ||
| "std": magnitudes.std(unbiased=False).item(), | ||
| "min": magnitudes.min().item(), | ||
| } |
Comment on lines
+32
to
+50
| - run_name_override: ["mod_add_p113_train30_wd1"] | ||
| tensorboard_run_name: ["mod_add_p113_train30_wd1"] | ||
| norm_variant_lm_head: ["cappedhyperspherenorm"] | ||
|
|
||
| - run_name_override: ["mod_add_p113_train30_wd1"] | ||
| tensorboard_run_name: ["mod_add_p113_train30_wd1"] | ||
| weight_decay: [1.0, 5.0, 10.0, 50.0, 100.0] | ||
| norm_variant_lm_head: ["cappedhyperspherenorm"] | ||
|
|
||
| - run_name_override: ["mod_add_p113_train30_wd1"] | ||
| tensorboard_run_name: ["mod_add_p113_train30_wd1"] | ||
| optimizer: ["muon"] | ||
| weight_decay: [0] | ||
|
|
||
| - run_name_override: ["mod_add_p113_train30_wd1"] | ||
| tensorboard_run_name: ["mod_add_p113_train30_wd1"] | ||
| norm_variant_lm_head: ["cappedhyperspherenorm"] | ||
| optimizer: ["muon"] | ||
| weight_decay: [0] |
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 introduces a comprehensive set of features and logging improvements centered around analyzing and reporting the magnitude statistics of the
lm_headoutput layer(s) in transformer models, with a focus on supporting modular arithmetic grokking experiments. It adds a new demo for modular addition grokking, expands metric logging and CSV export to include detailed statistics for thelm_headweights (both raw and post-normalization), and updates experiment monitoring and argument parsing to support these new features.Key changes include:
New Modular Arithmetic Grokking Demo
demos/modular_arithmetic_grokking_demo.py, an end-to-end script to prepare data, train, and evaluate a small transformer on modular addition, following best practices from the literature.explorations/modular_arithmetic_grokking.yamlfor running parameter sweeps on modular addition grokking tasks, with recommended settings and options for weight decay and normalization variants.Logging and Metric Reporting Enhancements
--log_lm_head_vector_stats,--log_post_norm_lm_head_vector_stats) to control logging oflm_headvector magnitude statistics, both raw and post-norm, to TensorBoard and CSV.train.pyto compute, update, log, and export max/avg/median/std/min statistics forlm_headrow vectors, supporting both raw and normalized weights. These are now included in CSV logs and TensorBoard summaries. [1] [2] [3] [4] [5] [6] [7]Experiment Monitoring and Metrics Schema Updates
run_exploration_monitor.pyandoptimization_and_search/run_experiments.pyto include the newlm_headandpost_norm_lm_headmagnitude statistics in their metric schemas, ensuring these are tracked and displayed in experiment dashboards. [1] [2] [3] [4]These changes provide better tools for diagnosing and understanding model behavior in modular arithmetic grokking tasks, and make it easier to monitor and compare experiments involving normalization strategies and output layer geometry.