Question about Zero_module

[Chuong98CC]

Hi, thanks for open source the code. It is a great help to understand how to implement Flow matching in practice.

I have a hard time to understand the purpose of zero_module function when implement UNET model for Image example. Particularly,

def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module

In attention block:

        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))

or ResNet block layer

        self.out_layers = nn.Sequential(
            normalization(self.out_channels),
            nn.SiLU(),
            nn.Dropout(p=dropout),
            zero_module(
                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
            ),
        )

In these blocks, the final layer in the block are all zero_module. Can you point me to the resource to explain why we need to do that? thank you