Add two-stage TrajMixer mixing

test_traj_mixer.py

@@ -21,14 +21,98 @@ class TrajMixerTest(unittest.TestCase):
 
         x = torch.randn(2, 7, 120)
         self.assertEqual(mixer(x).shape, x.shape)
-        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 15_840)
+        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 33_164)
 
         expected = torch.eye(12).expand(10, 12, 12)
         torch.testing.assert_close(mixer.group_align.detach(), expected)
+        self.assertEqual(mixer.intra_hidden, 48)
+        self.assertEqual(
+            tuple(mixer.intra_gate_proj.shape),
+            (10, 12, 48),
+        )
+        self.assertEqual(
+            tuple(mixer.intra_value_proj.shape),
+            (10, 12, 48),
+        )
+        self.assertEqual(
+            tuple(mixer.intra_output_proj.shape),
+            (10, 48, 12),
+        )
         self.assertEqual(mixer.hidden_group, 40)
         self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
         self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
         self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
+        self.assertEqual(tuple(mixer.intra_norm.normalized_shape), (12,))
+        self.assertEqual(tuple(mixer.cross_norm.normalized_shape), (10,))
+
+    def test_zero_output_projections_make_both_stages_identity(self) -> None:
+        torch.manual_seed(0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
+        with torch.no_grad():
+            mixer.intra_output_proj.zero_()
+            mixer.output_proj.zero_()
+        x = torch.randn(2, 5, 120)
+        torch.testing.assert_close(mixer(x), x)
+
+    def test_intra_stage_is_independent_across_groups(self) -> None:
+        torch.manual_seed(0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
+        mixer.eval()
+        with torch.no_grad():
+            mixer.output_proj.zero_()
+
+        grouped = torch.randn(2, 4, 10, 12)
+        changed = grouped.clone()
+        changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :])
+
+        original_out = mixer(grouped.reshape(2, 4, 120)).reshape(
+            2, 4, 10, 12
+        )
+        changed_out = mixer(changed.reshape(2, 4, 120)).reshape(
+            2, 4, 10, 12
+        )
+        unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9])
+        torch.testing.assert_close(
+            original_out.index_select(2, unchanged_groups),
+            changed_out.index_select(2, unchanged_groups),
+        )
+
+    def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
+        mixer = TrajMixer(6, n_head=3, dropout=0.0)
+        mixer.eval()
+        with torch.no_grad():
+            mixer.intra_output_proj.zero_()
+            mixer.gate_proj.zero_()
+            mixer.value_proj.zero_()
+            mixer.output_proj.zero_()
+
+            # For coordinate 0 only, read group 0 through hidden unit 0 and
+            # write the resulting gated value into group 1.
+            mixer.gate_proj[0, 0, 0] = 1.0
+            mixer.value_proj[0, 0, 0] = 1.0
+            mixer.output_proj[0, 0, 1] = 1.0
+
+        grouped = torch.tensor(
+            [[[
+                [-1.0, 4.0],
+                [0.0, 5.0],
+                [1.0, 6.0],
+            ]]]
+        )
+        changed = grouped.clone()
+        changed[0, 0, 0, 0] = 2.0
+
+        original_out = mixer(grouped.reshape(1, 1, 6)).reshape(1, 1, 3, 2)
+        changed_out = mixer(changed.reshape(1, 1, 6)).reshape(1, 1, 3, 2)
+
+        self.assertNotEqual(
+            original_out[0, 0, 1, 0].item(),
+            changed_out[0, 0, 1, 0].item(),
+        )
+        torch.testing.assert_close(
+            original_out[..., 1],
+            changed_out[..., 1],
+        )
 
     def test_mixer_does_not_mix_sequence_positions(self) -> None:
         torch.manual_seed(0)
@@ -58,11 +142,12 @@ class TrajMixerTest(unittest.TestCase):
             self.assertIsNotNone(parameter.grad, name)
             self.assertTrue(torch.isfinite(parameter.grad).all(), name)
 
-    def test_gpt_block_defaults_to_traj_mixer_and_standard_layer_norm(self) -> None:
+    def test_gpt_block_delegates_both_mixer_residuals_to_traj_mixer(self) -> None:
         block = GPTBlock(n_embd=120, n_head=10)
         self.assertIsInstance(block.mlp, TrajMixer)
-        self.assertIsInstance(block.ln2, torch.nn.LayerNorm)
-        self.assertEqual(tuple(block.ln2.normalized_shape), (120,))
+        self.assertFalse(hasattr(block, "ln2"))
+        self.assertIsInstance(block.mlp.intra_norm, torch.nn.LayerNorm)
+        self.assertIsInstance(block.mlp.cross_norm, torch.nn.LayerNorm)
 
         x = torch.randn(2, 6, 120)
         self.assertEqual(block(x).shape, x.shape)
@@ -75,6 +160,10 @@ class TrajMixerTest(unittest.TestCase):
             validate_traj_mixer_config({})
         with self.assertRaisesRegex(ValueError, "only accepts models trained"):
             validate_traj_mixer_config({"model_architecture": "delphi_swiglu"})
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v2"}
+            )
 
     def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None:
         block = GPTBlock(n_embd=120, n_head=10)
@@ -84,7 +173,7 @@ class TrajMixerTest(unittest.TestCase):
         }
         validate_traj_mixer_state_dict(state_dict)
 
-        state_dict.pop("blocks.0.mlp.group_align")
+        state_dict.pop("blocks.0.mlp.intra_gate_proj")
         with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"):
             validate_traj_mixer_state_dict(state_dict)
 
@@ -97,8 +186,8 @@ class TrajMixerTest(unittest.TestCase):
         self.assertEqual(
             get_model_parameter_counts(mixer),
             {
-                "model_parameter_count": 15_840,
-                "trainable_parameter_count": 15_840,
+                "model_parameter_count": 33_164,
+                "trainable_parameter_count": 33_164,
             },
         )