Remove the need for `einsum` in Albert's attention computation
This change makes it easier to optimize for exporting libraries such as ONNX and/or TensorRT. 
We did some with Lysandre, but I can rerun some more checks before we merge, just to be sure we are not breaking anything 👍🏻 