factorized embedding vectors: \(x^1_{k, l} + x^2_{l, k}\), where as the config.max_embedding_size dimension
\(j\) is factorized into \(k \text{ and } l\). This design ensures that each position embedding vector
\(x_j\) is unique.
Using the above example again, axial position encoding with \(d^1 = 2^5, d^2 = 2^5, n_s^1 = 2^9, n_s^2 = 2^{10}\)
can drastically reduced the number of parameters to \(2^{14} + 2^{15} \approx 49000\) parameters.
In practice, the parameter config.axial_pos_embds_dim is set to a tuple \((d^1, d^2)\) which sum has to be
equal to config.hidden_size and config.axial_pos_shape is set to a tuple \((n_s^1, n_s^2)\) which
product has to be equal to config.max_embedding_size, which during training has to be equal to the sequence
length of the input_ids.
LSH Self Attention
In Locality sensitive hashing (LSH) self attention the key and query projection weights are tied. Therefore, the key
