Repository navigation
Freeze exact zeros in every EM update kernel (#64 on current main) - #72
Merged
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
Conflicts: - model/tensorcircuit.py: main moved the class's methods out; main's file. This branch's `keep_zero_params` docstring goes with `mini_batch_em` to model/learning.py. - nodes/distributions/categorical.py: main already sums the normalizer per node (31cb81c); this branch's line, which also has the exact-zero test. The kernels main has added since this branch's base still have the magnitude cutoff: the next commit.
…reeze exact zeros too, with no cutoff -- `em_blocks.py` (sum layers with node blocks of 16 or more), `params_master.py` (bf16 sum layers with smaller blocks), `input_em.py` (Categorical up to 256 categories) and `dual_em_master.py` (gated bf16 layers whose master copy is larger than the L2) still took parameters below 1e-12 for zero. MEASURED on the merge alone: all 585 tiny weights of an HMM's sum layer zeroed, at fp32 and bf16, which this branch's test of the sum update could not see (block size 1). And `dual_em_master.py` reads the constants table by position, which this branch shortens: `theta < consts[5]` (now 0) never held, so exact zeros came back as the momentum floor, 1e-30 (16 of 16, the L2 condition forced; HMM-6144 reaches it), and `consts[6]` read past the table's end. `input_em.py` also drops the pseudocount's add-then-subtract, as this branch does in the template. Tests: the references of the per-form and input-layer tests take `== 0`; tiny weights added to them and to the master copy's correction test, whose relative tolerance cannot see 1e-30 against 0, so it checks them outright; the zero-preserving tests assert exact zeros. Each ported site put back to the cutoff (11), the input kernel's denominator, and the shortened table's bug: 13 mutants, 13 killed; the same tests pass on that lane unmutated. `keep_zero_params = False`, the default: the update from one saved state, main against this tree, HMM, HCLT, PD and a gated layer at fp32 and bf16, two steps each, plus the three input distributions at 64 to 512 categories: 25 of 28 cases bit-identical. The other three are gated: they differ between repeats of one tree too, by up to 3e-7, and main against this tree by the same, with some pairs identical. The tests of these kernels: 348 passed, with `--slow`. The full suite has not run on this commit: it runs once the pending PRs are merged.
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.
Supersedes #64 by Zoe Shao, on top of current
main. Zoe's three commits are unchanged and keep their authorship. This branch adds a merge ofmainand one commit that ports the change to the update kernels #71 added.What changes
keep_zero_params=Truefreezes only exact zeros. It used to freeze every parameter below1e-12as well, which silently deleted live weights. The docs already described the flag as preserving exact zeros.ExternProductCategoricalandSoftEvidenceCategoricalstill summed the retained flows over both axes, giving several nodes one denominator.Categoricalwas fixed earlier in31cb81ca.The flag is off by default, and nothing in
srcturns it on.Why #64 could not merge as it was
It merged with two small conflicts, but #71 added parallel kernels that #64 could not know about. MEASURED on a plain merge:
em_blocks.pyparams_master.pyinput_em.pyCategoricalup to 256 categoriesdual_em_master.py1e-30, with the L2 condition forcedIn the last row, #64 shortens a constants table that this kernel reads by position, so it never froze zeros and read past the table's end. The port makes all four use the exact-zero test.
Verification
--slow. The per-form, input-layer and master-copy tests now include tiny weights.keep_zero_params=False: the update from one saved state, againstmain, on HMM, HCLT, PD and a gated layer at fp32 and bf16, plus the three input distributions at 64 to 512 categories. 25 of 28 cases are bit-identical. The other 3 are gated: they differ between repeats of one tree by up to 3e-7, againstmainby the same, and some pairs match exactly.Not run: the full suite. It will run once the pending PRs are merged.