Fix zero-preserving updates: normalize per node and remove magnitude cutoff

src/pyjuice/model/backend/par_update.py

@@ -203,7 +203,7 @@ def cum_pflow_kernel(cum_pflows, mparams, param_flows, nchs, par_start_ids, pflo
         nch = tl.load(nchs + offs_m, mask = mask_m, other = 1)
         pflows += (pseudocount / nch[:,None])
 
-        nflows = tl.sum(tl.where(old_params < 1e-12, 0.0, pflows), axis = 1)
+        nflows = tl.sum(tl.where(old_params == 0.0, 0.0, pflows), axis = 1)
     else:
         nflows = tl.sum(pflows, axis = 1)
 
@@ -249,7 +249,7 @@ def em_par_update_kernel(mparams, param_flows, cum_pflows, nchs, par_start_ids,
     old_param = tl.load(mparams + offs_par, mask = mask_pflow, other = 0)
 
     if keep_zero_params:
-        updated_params = tl.where(old_param < 1e-12, 0.0, (1.0 - step_size) * old_param + step_size * new_param)
+        updated_params = tl.where(old_param == 0.0, 0.0, (1.0 - step_size) * old_param + step_size * new_param)
     else:
         updated_params = (1.0 - step_size) * old_param + step_size * new_param
 
@@ -286,7 +286,7 @@ def sgd_par_update_kernel(mparams, param_grads, par_start_ids, pgrad_start_ids,
     old_param = tl.load(mparams + offs_par, mask = mask_pgrad, other = 0)
 
     if keep_zero_params:
-        updated_params = tl.where(old_param < 1e-12, 0.0, tl.exp(tl.log(old_param) + lr * pgrads))
+        updated_params = tl.where(old_param == 0.0, 0.0, tl.exp(tl.log(old_param) + lr * pgrads))
     else:
         updated_params = tl.exp(tl.log(old_param) + lr * pgrads)
 
@@ -340,7 +340,7 @@ def sgd_par_update(params: torch.Tensor, param_grads: torch.Tensor, par_update_k
     :param lr: learning rate
     :type lr: float
 
-    :param keep_zero_params: whether to freeze zero parameters
+    :param keep_zero_params: whether to freeze exactly zero parameters (no magnitude threshold)
     :type keep_zero_params: bool
     """
 

src/pyjuice/model/tensorcircuit.py

@@ -1479,7 +1479,8 @@ def mini_batch_em(self, step_size: float, pseudocount: float = 0.0, keep_zero_pa
         :param pseudocount: a pseudo count added to the parameter flows
         :type pseudocount: float
 
-        :param keep_zero_params: if set to `True`, do not add pseudocounts to zero parameters
+        :param keep_zero_params: if `True`, freeze exactly zero parameters; positive parameters are updated
+            regardless of their magnitude (no implicit pruning threshold)
         :type keep_zero_params: bool
 
         :param step_size_rescaling: whether to rescale the step size by flows

src/pyjuice/nodes/distributions/categorical.py

@@ -207,21 +207,22 @@ def small_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
 
             if keep_zero_params:
                 param = tl.load(params_ptr + s_pids + cat_id, mask = cat_mask, other = 0)
-                cum_flow += tl.where(param < 1e-12, 0.0, flow + numerate_pseudocount)
+                cum_flow += tl.where(param == 0.0, 0.0, flow + numerate_pseudocount)
             else:
                 cum_flow += flow
 
         # Parameter update
-        cum_flow += pseudocount
+        if not keep_zero_params:
+            cum_flow += pseudocount
         for cat_id in range(max_num_cats):
             cat_mask = mask & (cat_id < num_cats)
 
             param = tl.load(params_ptr + s_pids + cat_id, mask = cat_mask, other = 0)
             flow = tl.load(param_flows_ptr + s_pfids + cat_id, mask = cat_mask, other = 0)
 
             if keep_zero_params:
-                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / (cum_flow - pseudocount)
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / cum_flow
+                new_param = tl.where(param == 0.0, 0.0, new_param)
             else:
                 new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / cum_flow
 
@@ -247,14 +248,15 @@ def large_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
 
             if keep_zero_params:
                 param = tl.load(params_ptr + s_pids[:,None] + cat_ids[None,:], mask = cat_mask, other = 0)
-                cum_flow += tl.sum(tl.where(param < 1e-12, 0.0, flow + numerate_pseudocount[:,None]))
+                cum_flow += tl.sum(tl.where(param == 0.0, 0.0, flow + numerate_pseudocount[:,None]), axis = 1)
             else:
                 cum_flow += tl.sum(flow, axis = 1)
 
             cat_ids += 128
 
         # Parameter update
-        cum_flow += pseudocount
+        if not keep_zero_params:
+            cum_flow += pseudocount
         cat_ids = tl.arange(0, 128)
         for cat_sid in range(0, max_num_cats, 128):
             cat_mask = mask[:,None] & (cat_ids[None,:] < num_cats[:,None])
@@ -263,8 +265,8 @@ def large_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
             flow = tl.load(param_flows_ptr + s_pfids[:,None] + cat_ids[None,:], mask = cat_mask, other = 0)
 
             if keep_zero_params:
-                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / (cum_flow[:,None] - pseudocount)
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / cum_flow[:,None]
+                new_param = tl.where(param == 0.0, 0.0, new_param)
             else:
                 new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / cum_flow[:,None]
             tl.store(params_ptr + s_pids[:,None] + cat_ids[None,:], new_param, mask = cat_mask)

src/pyjuice/nodes/distributions/external_categorical.py

@@ -1640,21 +1640,22 @@ def small_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
 
             if keep_zero_params:
                 param = tl.load(params_ptr + s_pids + cat_id, mask = cat_mask, other = 0)
-                cum_flow += tl.where(param < 1e-12, 0.0, flow + numerate_pseudocount)
+                cum_flow += tl.where(param == 0.0, 0.0, flow + numerate_pseudocount)
             else:
                 cum_flow += flow
 
         # Parameter update
-        cum_flow += pseudocount
+        if not keep_zero_params:
+            cum_flow += pseudocount
         for cat_id in range(max_num_cats):
             cat_mask = mask & (cat_id < num_cats)
 
             param = tl.load(params_ptr + s_pids + cat_id, mask = cat_mask, other = 0)
             flow = tl.load(param_flows_ptr + s_pfids + cat_id, mask = cat_mask, other = 0)
 
             if keep_zero_params:
-                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / (cum_flow - pseudocount)
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / cum_flow
+                new_param = tl.where(param == 0.0, 0.0, new_param)
             else:
                 new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / cum_flow
 
@@ -1680,14 +1681,15 @@ def large_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
 
             if keep_zero_params:
                 param = tl.load(params_ptr + s_pids[:,None] + cat_ids[None,:], mask = cat_mask, other = 0)
-                cum_flow += tl.sum(tl.where(param < 1e-12, 0.0, flow + numerate_pseudocount[:,None]))
+                cum_flow += tl.sum(tl.where(param == 0.0, 0.0, flow + numerate_pseudocount[:,None]), axis = 1)
             else:
                 cum_flow += tl.sum(flow, axis = 1)
 
             cat_ids += 128
 
         # Parameter update
-        cum_flow += pseudocount
+        if not keep_zero_params:
+            cum_flow += pseudocount
         cat_ids = tl.arange(0, 128)
         for cat_sid in range(0, max_num_cats, 128):
             cat_mask = mask[:,None] & (cat_ids[None,:] < num_cats[:,None])
@@ -1696,8 +1698,8 @@ def large_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
             flow = tl.load(param_flows_ptr + s_pfids[:,None] + cat_ids[None,:], mask = cat_mask, other = 0)
 
             if keep_zero_params:
-                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / (cum_flow[:,None] - pseudocount)
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / cum_flow[:,None]
+                new_param = tl.where(param == 0.0, 0.0, new_param)
             else:
                 new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / cum_flow[:,None]
             tl.store(params_ptr + s_pids[:,None] + cat_ids[None,:], new_param, mask = cat_mask)

src/pyjuice/nodes/distributions/softevi_categorical.py

@@ -3320,21 +3320,22 @@ def small_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
 
             if keep_zero_params:
                 param = tl.load(params_ptr + s_pids + cat_id, mask = cat_mask, other = 0)
-                cum_flow += tl.where(param < 1e-12, 0.0, flow + numerate_pseudocount)
+                cum_flow += tl.where(param == 0.0, 0.0, flow + numerate_pseudocount)
             else:
                 cum_flow += flow
 
         # Parameter update
-        cum_flow += pseudocount
+        if not keep_zero_params:
+            cum_flow += pseudocount
         for cat_id in range(max_num_cats):
             cat_mask = mask & (cat_id < num_cats)
 
             param = tl.load(params_ptr + s_pids + cat_id, mask = cat_mask, other = 0)
             flow = tl.load(param_flows_ptr + s_pfids + cat_id, mask = cat_mask, other = 0)
 
             if keep_zero_params:
-                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / (cum_flow - pseudocount)
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / cum_flow
+                new_param = tl.where(param == 0.0, 0.0, new_param)
             else:
                 new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount) / cum_flow
 
@@ -3360,14 +3361,15 @@ def large_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
 
             if keep_zero_params:
                 param = tl.load(params_ptr + s_pids[:,None] + cat_ids[None,:], mask = cat_mask, other = 0)
-                cum_flow += tl.sum(tl.where(param < 1e-12, 0.0, flow + numerate_pseudocount[:,None]))
+                cum_flow += tl.sum(tl.where(param == 0.0, 0.0, flow + numerate_pseudocount[:,None]), axis = 1)
             else:
                 cum_flow += tl.sum(flow, axis = 1)
 
             cat_ids += 128
 
         # Parameter update
-        cum_flow += pseudocount
+        if not keep_zero_params:
+            cum_flow += pseudocount
         cat_ids = tl.arange(0, 128)
         for cat_sid in range(0, max_num_cats, 128):
             cat_mask = mask[:,None] & (cat_ids[None,:] < num_cats[:,None])
@@ -3376,8 +3378,8 @@ def large_ncats_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfid
             flow = tl.load(param_flows_ptr + s_pfids[:,None] + cat_ids[None,:], mask = cat_mask, other = 0)
 
             if keep_zero_params:
-                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / (cum_flow[:,None] - pseudocount)
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / cum_flow[:,None]
+                new_param = tl.where(param == 0.0, 0.0, new_param)
             else:
                 new_param = (1.0 - step_size) * param + step_size * (flow + numerate_pseudocount[:,None]) / cum_flow[:,None]
             tl.store(params_ptr + s_pids[:,None] + cat_ids[None,:], new_param, mask = cat_mask)
@@ -3413,7 +3415,7 @@ def small_ncats_dual_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s
             flow = tl.where(cat_mask, flow, 0.0)
 
             if keep_zero_params:
-                cum_flow += tl.where(param < 1e-12, 0.0, flow)
+                cum_flow += tl.where(param == 0.0, 0.0, flow)
             else:
                 cum_flow += flow
 
@@ -3436,7 +3438,7 @@ def small_ncats_dual_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s
             new_param = tl.maximum(new_param, 1e-30)
 
             if keep_zero_params:
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = tl.where(param == 0.0, 0.0, new_param)
 
             tl.store(params_ptr + s_pids + cat_id, new_param, mask = cat_mask)
 
@@ -3466,7 +3468,7 @@ def large_ncats_dual_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s
             flow = tl.where(cat_mask, flow, 0.0)
 
             if keep_zero_params:
-                cum_flow += tl.sum(tl.where(param < 1e-12, 0.0, flow), axis = 1)
+                cum_flow += tl.sum(tl.where(param == 0.0, 0.0, flow), axis = 1)
             else:
                 cum_flow += tl.sum(flow, axis = 1)
 
@@ -3489,7 +3491,7 @@ def large_ncats_dual_em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s
             new_param = tl.maximum(new_param, 1e-30)
 
             if keep_zero_params:
-                new_param = tl.where(param < 1e-12, 0.0, new_param)
+                new_param = tl.where(param == 0.0, 0.0, new_param)
 
             tl.store(params_ptr + s_pids[:,None] + cat_ids[None,:], new_param, mask = cat_mask)
 

src/pyjuice/nodes/external_params/external_params.py

@@ -617,7 +617,7 @@ def _fused_em_correction(self, ns, params, param_flows, denom_param_flows, step_
         if consts is None:
             consts = cache_c[ckey] = torch.tensor(
                 [float(pseudocount), float(step_size), 1.0 - float(step_size),
-                 1e-38, 1e-30, 1e-12, 0.0], dtype = torch.float32, device = dev)
+                 1e-38, 1e-30, 0.0], dtype = torch.float32, device = dev)
             if len(cache_c) > 8:                  # a schedule sweeping step sizes must not grow it
                 for k in list(cache_c)[:-4]:
                     del cache_c[k]
@@ -772,7 +772,7 @@ def compute_em_correction(self, ns, params, param_flows, denom_param_flows, step
         new_theta = theta * ((1.0 - step_size) + step_size * ratio) / cum[nb][:, None, :]
         new_theta = new_theta.clamp_min(1e-30)                              # momentum-underflow guard
         if keep_zero_params:
-            new_theta = torch.where(theta < 1e-12, torch.zeros_like(new_theta), new_theta)
+            new_theta = torch.where(theta == 0.0, torch.zeros_like(new_theta), new_theta)
         new_theta = new_theta.reshape(-1)
         if out is not None:
             out.copy_(new_theta)

src/pyjuice/nodes/external_params/kernels/dual_em.py

@@ -170,8 +170,8 @@ def _dual_em_update_kernel(mparams, param_flows, denom, cum, out, cids, pids, pf
     new = theta * (one_m_s + step_size * ratio) / c[:,None]
     new = tl.maximum(new, tiny_par)                            # momentum-underflow guard
     if KEEP_ZERO:
-        zero_thr = tl.load(consts + 5)                         # 1e-12, the `keep_zero_params` threshold
-        new = tl.where(theta < zero_thr, tl.load(consts + 6), new)
+        zero = tl.load(consts + 5)                             # freeze exactly zero parameters
+        new = tl.where(theta == zero, zero, new)
 
     if OUT_BY_PAR:
         off = par[None,:] + offs_node[:,None] - out_base

tests/model/external_params/denom_param_flows_test.py

@@ -1954,7 +1954,7 @@ def _ref_dual_m_step(ns, theta, fp, fm, step_size, pseudocount, keep_zero_params
                 out[e, :, m] = th[e, :, m] * ((1.0 - step_size) + step_size * ratios[e]) / cum
     out = torch.clamp(out, min = 1e-30)
     if keep_zero_params:
-        out = torch.where(th < 1e-12, torch.zeros_like(out), out)
+        out = torch.where(th == 0.0, torch.zeros_like(out), out)
     return out.reshape(-1)
 
 
@@ -2126,3 +2126,58 @@ def test_dense_and_rowgroup_denom_modes_agree(bs, nb, gcbs, step, pcount, kz):
     ref = dense_desc.compute_em_correction(*args, denom_sources = srcs[tgt])
     rel2 = ((rowgroup - ref).abs() / ref.abs().clamp_min(1e-20)).max().item()
     assert rel2 < 1e-5, f"the fused modes disagree with the torch M-step (relmax = {rel2})"
+
+
+@cuda_only
+@pytest.mark.parametrize("fused", [False, True])
+@pytest.mark.parametrize("keep", [False, True])
+def test_em_correction_preserves_tiny_positive_parameters(fused, keep, monkeypatch):
+    """The fused and torch M-steps freeze exact zeros, not small positive weights."""
+    dev = torch.device("cuda:0")
+    root, ns = _build(gated = True)
+    pc = juice.compile(root, verbose = False).to(dev)
+    ps, pe = ns._param_range
+    E, cbs, bs = ns.edge_ids.size(1), ns.ch_block_size, ns.block_size
+    theta = pc.params[ps:pe].reshape(E, cbs, bs)
+    theta[:, 0, :] = 0.0
+    theta[:, 1, :] = 1e-16
+    nb = ns.edge_ids[0].to(device = dev, dtype = torch.long)
+    total = torch.zeros(ns.num_node_blocks, bs, device = dev)
+    total.index_add_(0, nb, theta.sum(1))
+    theta /= total[nb][:, None, :]
+    before = theta.clone()
+    zeroed, tiny = before == 0, (before > 0) & (before < 1e-12)
+    assert zeroed.any() and tiny.any()
+    if not fused:
+        monkeypatch.setattr(ExternalSumParams, "_fused_em_correction", lambda *args, **kwargs: None)
+    else:
+        fused_calls = []
+        original = ExternalSumParams._fused_em_correction
+
+        def checked_fused(*args, **kwargs):
+            result = original(*args, **kwargs)
+            assert result is not None, "the fused M-step unexpectedly fell back to torch"
+            fused_calls.append(True)
+            return result
+
+        monkeypatch.setattr(ExternalSumParams, "_fused_em_correction", checked_fused)
+    torch.manual_seed(17)
+    x = torch.randint(0, NUM_CATS, [32, 2], device = dev)
+    ext = {ns: _phi(ns, 32, dev)}
+    pc(x, sum_external_params = ext)
+    pc.backward(x, sum_external_params = ext, logspace_flows = True,
+                flows_memory = 0.0, allow_modify_flows = False)
+    fp = pc.param_flows[slice(*ns._param_flow_range)].clone()
+    fm = _fminus(pc, ns)[slice(*ns._param_flow_range)]
+    expected = _ref_dual_m_step(ns, before, fp, fm, 1.0, 0.1, keep).reshape_as(before)
+    pc.mini_batch_em(step_size = 1.0, pseudocount = 0.1, keep_zero_params = keep)
+    if fused:
+        assert fused_calls, "the fused M-step was not exercised"
+    actual = pc.params[ps:pe].reshape_as(before)
+    torch.testing.assert_close(actual.double(), expected, atol = 2e-6, rtol = 2e-4)
+    torch.testing.assert_close(actual[tiny].double(), expected[tiny], atol = 0, rtol = 2e-4)
+    assert bool((actual[tiny] > 0).all())
+    if keep:
+        assert bool((actual[zeroed] == 0).all())
+    else:
+        assert bool((actual[zeroed] > 0).all())