[AINode] Replace DiT-derived Sundial conditioning head

iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py

@@ -70,10 +70,6 @@ def sample(self, z, num_samples=1):
         return x
 
 
-def modulate(x, shift, scale):
-    return x * (1 + scale) + shift
-
-
 class TimestepEmbedder(nn.Module):
     """
     Embeds scalar timesteps into vector representations.
@@ -119,11 +115,8 @@ def forward(self, t):
         return t_emb
 
 
-class ResBlock(nn.Module):
-    """
-    A residual block that can optionally change the number of channels.
-    :param channels: the number of input channels.
-    """
+class ConditionalResidualBlock(nn.Module):
+    """A residual MLP controlled by a per-sample conditioning vector."""
 
     def __init__(self, channels):
         super().__init__()
@@ -140,17 +133,21 @@ def __init__(self, channels):
             nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
         )
 
-    def forward(self, x, y):
-        shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
-        h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
-        h = self.mlp(h)
-        return x + gate_mlp * h
+    def reset_conditioning_parameters(self):
+        nn.init.zeros_(self.adaLN_modulation[-1].weight)
+        nn.init.zeros_(self.adaLN_modulation[-1].bias)
 
+    def forward(self, features, condition):
+        offset, gain_delta, update_scale = torch.tensor_split(
+            self.adaLN_modulation(condition), 3, dim=-1
+        )
+        conditioned = torch.addcmul(offset, self.in_ln(features), gain_delta.add(1))
+        update = self.mlp(conditioned)
+        return torch.addcmul(features, update, update_scale)
 
-class FinalLayer(nn.Module):
-    """
-    The final layer adopted from DiT.
-    """
+
+class ConditionalOutputProjection(nn.Module):
+    """Map condition-normalized features to the requested output width."""
 
     def __init__(self, model_channels, out_channels):
         super().__init__()
@@ -162,11 +159,20 @@ def __init__(self, model_channels, out_channels):
             nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True)
         )
 
-    def forward(self, x, c):
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
+    def reset_conditioning_parameters(self):
+        nn.init.zeros_(self.adaLN_modulation[-1].weight)
+        nn.init.zeros_(self.adaLN_modulation[-1].bias)
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, features, condition):
+        offset, gain_delta = torch.tensor_split(
+            self.adaLN_modulation(condition), 2, dim=-1
+        )
+        conditioned = torch.addcmul(
+            offset, self.norm_final(features), gain_delta.add(1)
+        )
+        return self.linear(conditioned)
 
 
 class SimpleMLPAdaLN(nn.Module):
@@ -199,16 +205,10 @@ def __init__(
 
         self.input_proj = nn.Linear(in_channels, model_channels)
 
-        res_blocks = []
-        for i in range(num_res_blocks):
-            res_blocks.append(
-                ResBlock(
-                    model_channels,
-                )
-            )
-
-        self.res_blocks = nn.ModuleList(res_blocks)
-        self.final_layer = FinalLayer(model_channels, out_channels)
+        self.res_blocks = nn.ModuleList(
+            ConditionalResidualBlock(model_channels) for _ in range(num_res_blocks)
+        )
+        self.final_layer = ConditionalOutputProjection(model_channels, out_channels)
 
         self.initialize_weights()
 
@@ -225,16 +225,10 @@ def _basic_init(module):
         nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
         nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
 
-        # Zero-out adaLN modulation layers
         for block in self.res_blocks:
-            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-
-        # Zero-out output layers
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
-        nn.init.constant_(self.final_layer.linear.weight, 0)
-        nn.init.constant_(self.final_layer.linear.bias, 0)
+            block.reset_conditioning_parameters()
+
+        self.final_layer.reset_conditioning_parameters()
 
     def forward(self, x, t, c):
         """

iotdb-core/ainode/tests/test_sundial_flow_loss.py

@@ -0,0 +1,123 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import unittest
+
+import torch
+
+from iotdb.ainode.core.model.sundial.flow_loss import (
+    ConditionalOutputProjection,
+    ConditionalResidualBlock,
+    FlowLoss,
+    SimpleMLPAdaLN,
+)
+
+
+class SundialFlowLossTest(unittest.TestCase):
+    def create_network(self):
+        return SimpleMLPAdaLN(
+            in_channels=4,
+            model_channels=8,
+            out_channels=4,
+            z_channels=6,
+            num_res_blocks=2,
+        )
+
+    def test_checkpoint_parameter_contract(self):
+        network = self.create_network()
+        expected_shapes = {
+            "time_embed.mlp.0.weight": (8, 256),
+            "time_embed.mlp.0.bias": (8,),
+            "time_embed.mlp.2.weight": (8, 8),
+            "time_embed.mlp.2.bias": (8,),
+            "cond_embed.weight": (8, 6),
+            "cond_embed.bias": (8,),
+            "input_proj.weight": (8, 4),
+            "input_proj.bias": (8,),
+            "res_blocks.0.in_ln.weight": (8,),
+            "res_blocks.0.in_ln.bias": (8,),
+            "res_blocks.0.mlp.0.weight": (8, 8),
+            "res_blocks.0.mlp.0.bias": (8,),
+            "res_blocks.0.mlp.2.weight": (8, 8),
+            "res_blocks.0.mlp.2.bias": (8,),
+            "res_blocks.0.adaLN_modulation.1.weight": (24, 8),
+            "res_blocks.0.adaLN_modulation.1.bias": (24,),
+            "res_blocks.1.in_ln.weight": (8,),
+            "res_blocks.1.in_ln.bias": (8,),
+            "res_blocks.1.mlp.0.weight": (8, 8),
+            "res_blocks.1.mlp.0.bias": (8,),
+            "res_blocks.1.mlp.2.weight": (8, 8),
+            "res_blocks.1.mlp.2.bias": (8,),
+            "res_blocks.1.adaLN_modulation.1.weight": (24, 8),
+            "res_blocks.1.adaLN_modulation.1.bias": (24,),
+            "final_layer.linear.weight": (4, 8),
+            "final_layer.linear.bias": (4,),
+            "final_layer.adaLN_modulation.1.weight": (16, 8),
+            "final_layer.adaLN_modulation.1.bias": (16,),
+        }
+
+        actual_shapes = {
+            key: tuple(value.shape) for key, value in network.state_dict().items()
+        }
+        self.assertEqual(expected_shapes, actual_shapes)
+
+        restored = self.create_network()
+        result = restored.load_state_dict(network.state_dict(), strict=True)
+        self.assertEqual([], result.missing_keys)
+        self.assertEqual([], result.unexpected_keys)
+
+    def test_conditioned_layers_propagate_gradients(self):
+        features = torch.randn(3, 8, requires_grad=True)
+        condition = torch.randn(3, 8, requires_grad=True)
+        block = ConditionalResidualBlock(8)
+        projection = ConditionalOutputProjection(8, 4)
+
+        output = projection(block(features, condition), condition)
+
+        self.assertEqual((3, 4), tuple(output.shape))
+        output.square().mean().backward()
+        self.assertIsNotNone(features.grad)
+        self.assertIsNotNone(condition.grad)
+
+    def test_network_starts_with_zero_output(self):
+        network = self.create_network()
+
+        output = network(
+            torch.randn(3, 4),
+            torch.tensor([0.0, 500.0, 999.0]),
+            torch.randn(3, 6),
+        )
+
+        torch.testing.assert_close(output, torch.zeros_like(output))
+
+    def test_sample_shape(self):
+        flow_loss = FlowLoss(
+            target_channels=4,
+            z_channels=6,
+            depth=2,
+            width=8,
+            num_sampling_steps=2,
+        )
+
+        samples = flow_loss.sample(torch.randn(2, 6), num_samples=3)
+
+        self.assertEqual((2, 3, 4), tuple(samples.shape))
+
+
+if __name__ == "__main__":
+    unittest.main()