Repository navigation
Fix zero-preserving updates: normalize per node and remove magnitude cutoff - #64
Merged
liuanji merged 3 commits intoNov 4, 2026
Merged
Conversation
…izer over both axes `Categorical.large_ncats_em_fn` (selected when `num_cats > 256`) accumulates the per-node normalizer over a `[BLOCK_SIZE, 128]` tile, so the reduction must be over the category axis. The `keep_zero_params` branch called `tl.sum(...)` with no `axis`, which reduces over *both* axes and broadcasts one scalar back to `cum_flow`, giving every node in the block the same denominator: the summed flow of all `em_block_size = 8` nodes in it. The `else` branch next to it already used `axis = 1`. Each input distribution is then divided by roughly 8x too much and loses its normalization on every EM step, silently, with no error raised. One `mini_batch_em(step_size = 1.0, keep_zero_params = True)` on an HMM with `num_emits = 512` leaves the emission rows summing to 0.125 instead of 1. The same line appears in `ExternProductCategorical` and `SoftEvidenceCategorical`; the second `large_ncats_em_fn` in `softevi_categorical.py` already had `axis = 1` and is unchanged. Existing coverage missed this because `tests/optim/zero_param_preserving_test.py` uses `num_emits = 256`, which takes the `small_ncats_em_fn` path. The added test is parametrized over both sides of that threshold and fails only at 512 before this change. `keep_zero_params` defaults to `False`, so only callers that opt in -- typically to hold pruned parameters at zero -- were affected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017RQQtEYKLLjVFh9AgJbYMB
liuanji
added a commit
that referenced
this pull request
Nov 4, 2026
Freeze exact zeros in every EM update kernel (#64 on current main)
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.
keep_zero_params=Truecurrently freezes both exact zeros and positive weights below1e-12. For example, a live weight of1e-16is deleted even when its current flow would increase it. This can silently remove valid probability paths. The optimizer documentation already describes the flag as preserving exactly zero parameters.There is also a normalization bug in the large-category EM kernels: with more than 256 categories, the retained-flow sum reduces over both the node and category axes, giving multiple nodes a shared denominator.
This PR:
axis=1) so each node has its own normalizer.parameter == 0is frozen. Applies this consistently to categorical EM, ordinary sum EM/SGD, soft-evidence dual-flow EM, and external-sum fused/fallback updates.I propose keeping magnitude pruning separate from this flag, without introducing a replacement cutoff. Feedback is welcome if the original threshold was intentional. This does not guarantee that every positive weight stays positive: ordinary EM with zero counts and no smoothing, or floating-point underflow, can still produce zero. The
keep_zero_params=Falsepaths and existing division/underflow guards are unchanged.Validation against upstream
3b44e329, on an NVIDIA RTX PRO 6000 Blackwell GPU with Python 3.12.12, PyTorch 2.9.1+cu130 and Triton 3.5.1:Focused tests cover both flag values, repeated updates, zero and tiny positive weights, partial/full step sizes, zero/nonzero pseudocounts, category counts 256/257/512/50,262, partial tiles, and pseudocount-only rows. External-sum tests explicitly exercise the fused and fallback paths. Compatibility checks cover HMM training, categorical/external/soft-evidence updates, and external-sum normalization, orchestration, and denominator layouts.
PYTHONPATH=src python -m pytest \ tests/optim/categorical_zero_normalization_test.py \ tests/optim/exact_zero_sum_update_test.py \ tests/optim/zero_param_preserving_test.py \ tests/model/external_params/denom_param_flows_test.py::test_em_correction_preserves_tiny_positive_parameters \ -m "not slow" -qOptional CUTLASS/Ninja dependencies were unavailable, so optional CUDA extensions used their fallbacks. The modified update kernels are Triton kernels. The full suite and documentation build were not run.