On your R2 rotation Matrix

[kimjoohyungsd]

In multi-head attention, the rotation matrix must be block-diagonal with one block per head. Why does SpinQuant train a single shared R2 matrix of shape [head_dim, head_dim] and replicate it across all heads (creating a block-diagonal with identical blocks), rather than training independent rotation matrices for each head?

[fxmarty-amd]

How would you go about fusing into https://github.com/huggingface/transformers/blob/d1808f2c36c02faad537f9737a76165e49b041f9/src/transformers/models/llama/modeling_llama.py#L245 then?

edit: you might be able to train num_key_value_heads rotations [head_dim, head_dim] I think (edit2: confirm it works. You can even train a cross-head online rotation, like in quarot paper!)

[fxmarty-amd]

@kimjoohyungsd I attempted tuning R2 per-kv head. I could not achieve significant improvement this way, as compared to a shared orthogonal matrix across all heads. Curious to know if you have attempted it as well?

[YilanWang]

@fxmarty-amd yeah per head training r2 matrix has no improvements, actually I find train r2 has little improvements