From 80579d6cc254d337a23e71404ae7ecab1849d1e5 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Mon, 9 Feb 2026 14:40:31 -0600 Subject: Fix cascading gate: exempt layer 0 from disconnection check Layer 0 has no incoming edges structurally (no prior layers), but receives the embedding as input. The cascading gate was killing its outgoing edges (hard: g=0, soft: g=0.5), causing nll_hard to be ~2x worse than baseline. Fix: set g=1 for layer 0 nodes. Co-Authored-By: Claude Opus 4.6 --- tests/test_predictor.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) (limited to 'tests') diff --git a/tests/test_predictor.py b/tests/test_predictor.py index 5a092d4..104c204 100644 --- a/tests/test_predictor.py +++ b/tests/test_predictor.py @@ -159,15 +159,22 @@ class TestCascadingGate: assert (A_gated >= 0).all() and (A_gated <= 1).all() def test_hard_mode_kills_disconnected(self): - """Nodes with no incoming edges should have all outgoing edges zeroed.""" + """Non-layer-0 nodes with no incoming edges should have outgoing edges zeroed. + + Layer 0 is exempt (receives embedding, not truly disconnected). + """ A = torch.zeros(1, 256, 256) - # Only set edges from node 0 to node 16 (layer 0 → layer 1) - A[0, 0, 16] = 1.0 + # Node 32 (layer 2, head 0) has no incoming edges but has outgoing to node 48 + A[0, 32, 48] = 1.0 A_gated = cascading_gate(A, k=5.0, hard=True) - # Node 0 has no incoming edges → its outgoing should be zeroed - assert A_gated[0, 0, 16] == 0.0, "Node 0 has no incoming but wasn't gated to 0" - # Node 16 has incoming from node 0 (but node 0 was gated to 0) - # In one-pass mode, inc uses ORIGINAL A, so node 16 has inc > 0 + # Node 32 has no incoming → outgoing should be zeroed + assert A_gated[0, 32, 48] == 0.0, "Node 32 has no incoming but wasn't gated to 0" + + # Layer 0 should be exempt: node 0 has no incoming but keeps outgoing + A2 = torch.zeros(1, 256, 256) + A2[0, 0, 16] = 1.0 + A2_gated = cascading_gate(A2, k=5.0, hard=True) + assert A2_gated[0, 0, 16] == 1.0, "Layer 0 should be exempt from cascading gate" def test_hard_mode_preserves_connected(self): """Nodes with incoming edges keep their outgoing edges.""" -- cgit v1.2.3