How to train a causal TAAE model?

[yanzhou3366]

In released code ,the model seems to be non-causal, how can we train a causal version?

[Aceticia]

I second this. @julian-parker Is it possible to release the weight or at least the model code so we can reproduce the causal results?

[Aceticia]

From the other issue #6 , setting sliding window to something like [63, 0] can make the attention part causal. I have hacked the following patching code to convert non-causal convolutions into causal convs:

class CausalConv1d(nn.Conv1d):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        dilation=1,
        groups=1,
        bias=True,
        padding_mode="zeros",
        **kw,
    ):
        super().__init__(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=0,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            **kw,
        )
        k_eff = (self.kernel_size[0] - 1) * self.dilation[0] + 1
        self.left_pad = k_eff - 1
        self.k_eff = k_eff

    @classmethod
    def from_conv(cls, conv: nn.Conv1d) -> "CausalConv1d":
        obj = cls(
            conv.in_channels,
            conv.out_channels,
            conv.kernel_size,
            stride=conv.stride,
            dilation=conv.dilation,
            groups=conv.groups,
            bias=conv.bias is not None,
            padding_mode=conv.padding_mode,
        )
        obj.weight.data.copy_(conv.weight.data)
        if conv.bias is not None:
            obj.bias.data.copy_(conv.bias.data)
        return obj

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.pad(x, (self.left_pad, 0))
        y = super().forward(x)

        # Anything longer than the original input length is due to padding
        L_in = x.shape[-1] - self.left_pad
        L_exp = (L_in - self.k_eff) // self.stride[0] + 1
        extra = y.shape[-1] - L_exp
        if extra:
            y = y[..., :-extra]
        return y


def make_model_causal(module: nn.Module) -> nn.Module:
    for name, child in list(module.named_children()):
        if isinstance(child, nn.Conv1d):
            setattr(module, name, CausalConv1d.from_conv(child))
        else:
            make_model_causal(child)
    return module

Testing this out... Hopefully this is close enough to the official causal version to produce similar results?

[julian-parker]

@Aceticia, this approach seems reasonable. There's also causal Conv1d implementations inside adp.py in stable-audio-tools. The 'official' version was indeed just this + the causal sliding window.

Weights release is unfortunately unlikely for the causal version from the paper. It should not need much compute to reproduce though. I trained it in around a day by finetuning the released base checkpoint.