MultiheadAttentionBlock: Unusual FFN

[ldv1]

Hi,

In the MultiheadAttentionBlock class, the position-wise feed-forward network consists of a single linear transformation with a ReLU activation:

out = out + self.lin(out).relu()

whereas in almost any other implementation it consists of two linear transformations with a ReLU activation in between (based on the paper Attention Is All You Need):

out = out + self.lin2(self.lin1(out).relu())

(and nowadays, GELU is standard)

What is the rational behind this unusual implementation?

[rusty1s]

This is mostly taken from the official implementation, see https://github.com/davidbuterez/gnn-neural-readouts/blob/main/code/models/set_transformer_modules.py#L8.

[ldv1]

I see.
This seems strange to me but this readout is entirely based on the Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks by Lee et al. whose original implementation also uses a single linear layer followed by ReLU.
Thanks for the prompt answer.

[rusty1s]

I am happy to change this if you see a strong need to do this.

[ldv1]

Related to this topic:

In the implementation you mentioned, we found for the PMA

class PMA(nn.Module):
    def __init__(...)
    ...
    def forward(self, X):
        return self.mab(self.S.repeat(X.size(0), 1, 1), X)

but in PyG, we have

class PoolingByMultiheadAttention(torch.nn.Module):
    def __init__(...)
    ...
    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        x = self.lin(x).relu()
        return self.mab(self.seed.expand(x.size(0), -1, -1), x, y_mask=mask)

Why do we have x = self.lin(x).relu() that is also not specified in the Set Transformer paper ? (unless I oversaw it) ?
Thanks for the help.

[rusty1s]

This is taken from Eq (11) in https://arxiv.org/pdf/1810.00825.pdf. It depends on whether you want to implement the feature transformation within or without the module. The paper does integrate it into PMA while the code does not.

[ldv1]

You're right! Many thanks!