Why does ALBERT use einsum in PyTorch implementation while in TF one it does not?

❓ Questions & Help

I wanted to learn internals of ALBERT model from your implementation (which is BTW really clean in comparison to the original one - good job!), but I’ve stumbled upon weird looking part in the AlbertAttention: https://github.com/huggingface/transformers/blob/6af3306a1da0322f58861b1fbb62ce5223d97b8a/src/transformers/modeling_albert.py#L258

Why does PyTorch version use einsum-based notation while calculating hidden state (with manual usage of dense layer’s weights), while the TensorFlow version just reshapes the context_layer and does standard “forward” on dense layer?

https://github.com/huggingface/transformers/blob/6af3306a1da0322f58861b1fbb62ce5223d97b8a/src/transformers/modeling_tf_albert.py#L296

I would really like to know the explanation of this implementation - @LysandreJik cloud you shed some light here?

[LysandreJik]

For no particular reason, but it might not have been the best choice according to this thread on performance.

[ZhuBaohe]

The two implementations are equivalent, but the Pytorch version is cumbersome. I think the code

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

        # Should find a better way to do this
        w = (
            self.dense.weight.t()
            .view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
            .to(context_layer.dtype)
        )
        b = self.dense.bias.to(context_layer.dtype)

        projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b

should be rewritten by

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_shape = context_layer.size()[:-2] + (-1,)
        context_layer = context_layer.view(*new_shape)

        projected_context_layer = self.dense(context_layer)