From bcf5187277076eb0869cf4c99af4496433c32730 Mon Sep 17 00:00:00 2001 From: Yongzao <532741407@qq.com> Date: Thu, 10 Sep 2026 17:41:02 +0800 Subject: [PATCH] [AINode] Replace DiT-derived Sundial conditioning head --- .../ainode/core/model/sundial/flow_loss.py | 78 +++++------ .../ainode/tests/test_sundial_flow_loss.py | 123 ++++++++++++++++++ 2 files changed, 159 insertions(+), 42 deletions(-) create mode 100644 iotdb-core/ainode/tests/test_sundial_flow_loss.py diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py index b3fe95dbe2d2..76a49431edf1 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py +++ b/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): """ diff --git a/iotdb-core/ainode/tests/test_sundial_flow_loss.py b/iotdb-core/ainode/tests/test_sundial_flow_loss.py new file mode 100644 index 000000000000..cb81f5f7a675 --- /dev/null +++ b/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()