text: Title: Should init scale matrix as diagonal form?
  Body: Hi, Phil:
  
  I noticed the LayerScale part in the `CaiT`, in the original paper the scale matrix is a diagonal form `(b,d,d)`, but in this implement, it just initialized in a form of vector(maybe can broadcast afterwards, but would it be better just initialize as a diagonal form?)
  https://github.com/lucidrains/vit-pytorch/blob/3f754956fbfb1f97ae4f1e244a7ecb16eab79296/vit_pytorch/cait.py#L41
  
  Best,