Question about key_padding_mask in TSTPlus

[michaelyma12]

Hello,

Thanks for the great work on the library.

Noticed something that confused me in the source for TSTPlus.

I get how the key_madding_mask is fed through the layers in _TSTBackbone, which outputs a tensor of shape:
(batch_size, n_features, time_dimension).

Accordingly, all the rows along the time dimension that were padded become 0s.

The part that confuses me is how the subsequent head module produced by self.create_head(...) simply feeds the flattened tensor into a LinBnDrop layer.

Specifically referencing this block in tsai.models.TSTPlus:

    def create_head(self, nf, c_out, seq_len, flatten=True, concat_pool=False, act="gelu", fc_dropout=0., bn=False, y_range=None):
        layers = [get_activation_fn(act)]
        if flatten:
            nf *= seq_len
            layers += [Flatten()]
        else:
            if concat_pool: nf *= 2
            layers = [GACP1d(1) if concat_pool else GAP1d(1)]
        layers += [LinBnDrop(nf, c_out, bn=bn, p=fc_dropout)]
        if y_range: layers += [SigmoidRange(*y_range)]
        return nn.Sequential(*layers)

Wouldn't the 0 vectors also be fed in? Doesn't that hinder the model in some way? Is this simply the best workable solution?

Any insight would be appreciated 🙇