Add row norm function#253
Conversation
Add row_norm_fn, a composable row-wise normalization update usable as scaled_orthogonalize_fn of OrthogonalizedOptimizer. Follow-up to NVIDIA-NeMo#188. Signed-off-by: Abdul Basit Tonmoy <abdulbasittonmoy@gmail.com>
Greptile SummaryThis PR adds
Confidence Score: 5/5Safe to merge — standalone additive change with no modifications to existing optimizers The implementation is a pure-function module with correct math for both center_rows paths, proper dtype round-trip, and exact zero output for zero-gradient inputs. No existing code is modified. No files require special attention. Important Files Changed
Reviews (2): Last reviewed commit: "Address review comments" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| r"""Row-wise normalization of the update. | ||
|
|
||
| Scales each row of ``G`` to unit norm: | ||
|
|
||
| .. math:: | ||
| u_i = G_i \, / \max(\lVert G_i \rVert_2, \epsilon) | ||
|
|
||
| Usable as ``scaled_orthogonalize_fn`` of | ||
| :class:`~emerging_optimizers.orthogonalized_optimizers.OrthogonalizedOptimizer`. | ||
|
|
||
| Args: | ||
| grad: The (momentum) tensor to normalize. | ||
| center_rows: If True, subtract the per-column mean (the average over the row axis, ``dim=0``) | ||
| before and after the update, so each column is zero-mean. | ||
| eps: Floor on the row norms. | ||
| extra_scale_factor: Extra multiplier on the update. | ||
|
|
||
| Returns: | ||
| The row-normalized update, same shape and dtype as ``grad``. | ||
| """ |
There was a problem hiding this comment.
center_rows=True
The docstring opens with "Scales each row of G to unit norm" and shows the math formula u_i = G_i / max(||G_i||_2, ε), implying unit-norm output. This holds only when center_rows=False. With center_rows=True, the post-centering step (update - update.mean(dim=0)) projects the update back onto the zero-column-mean subspace, which generally breaks the unit-row-norm property. A user who reads the docstring and passes center_rows=True expecting unit-norm rows will get incorrectly scaled rows. Notably, test_rows_close_to_unit_norm only exercises the default (center_rows=False) case, which confirms the guarantee doesn't extend to the centered variant. The docstring should clarify this (e.g., "when center_rows=False").
Signed-off-by: Abdul Basit Tonmoy <abdulbasittonmoy@gmail.com>
1 participant
What does this PR do?
Adds
row_norm_fn, a composable update normalization from #188: each row of the momentum is scaled to unit norm, with optionalcenter_rowsmatching the polargrad functions. Usable asscaled_orthogonalize_fnofOrthogonalizedOptimizer; no new optimizer class.Placed in
orthogonalized_optimizersfor composability. As noted in #188 this is not strictly an orthogonalization, so happy to move it if you prefer another home.Tests
Unit row norms, per-row scale invariance, row-permutation equivariance, right-orthogonal equivariance, centered-output projection, zero-gradient edge case, and composition through
OrthogonalizedOptimizer.step(). Verified locally on--device=cpuand--device=cuda, random seed and--seed=42; pre-commit clean.