

Axial Positional Encodings
Axial Positional Encodings were first implemented in Google's trax library
and developed by the authors of this model's paper. In models that are treating very long input sequences, the
conventional position id encodings store an embeddings vector of size \(d\) being the config.hidden_size for
every position \(i, \ldots, n_s\), with \(n_s\) being config.max_embedding_size. This means that having
a sequence length of \(n_s = 2^{19} \approx 0.5M\) and a config.hidden_size of \(d = 2^{10} \approx 1000\)
would result in a position encoding matrix:
$$X_{i,j}, \text{ with } i \in \left[1,\ldots, d\right] \text{ and } j \in \left[1,\ldots, n_s\right]$$
which alone has over 500M parameters to store. Axial positional encodings factorize \(X_{i,j}\) into two matrices:
$$X^{1}_{i,j}, \text{ with } i \in \left[1,\ldots, d^1\right] \text{ and } j \in \left[1,\ldots, n_s^1\right]$$
and
$$X^{2}_{i,j}, \text{ with } i \in \left[1,\ldots, d^2\right] \text{ and } j \in \left[1,\ldots, n_s^2\right]$$
with:
$$d = d^1 + d^2 \text{ and } n_s = n_s^1 \times n_s^2 .$$
Therefore the following holds:
$$X_{i,j} = \begin{cases}
X^{1}{i, k}, & \text{if }\ i < d^1 \text{ with } k = j \mod n_s^1 \
X^{2}{i - d^1, l}, & \text{if } i \ge d^1 \text{ with } l = \lfloor\frac{j}{n_s^1}\rfloor
\end{cases}$$
Intuitively, this means that a position embedding vector \(x_j \in \mathbb{R}^{d}\) is now the composition of two
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^9, d^2 = 2^9, n_s^1 = 2^9, n_s^2 = 2^{10}\)
can drastically reduced the number of parameters from 500 000 000 to \(2^{18} + 2^{19} \approx 780 000\) parameters, this means 85% less memory usage.
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
query embedding vectors are also tied. LSH self attention uses the locality sensitive hashing mechanism proposed in
Practical and Optimal LSH for Angular Distance to assign each of the tied key
query embedding vectors to one of config.num_buckets possible buckets. The premise is that the more "similar"
key query embedding vectors (in terms of cosine similarity) are to each other, the more likely they are assigned to
the same bucket.
The accuracy of the LSH mechanism can be improved by increasing config.num_hashes or directly the argument
num_hashes of the forward function so that the output of the LSH self attention better approximates the output
of the "normal" full self attention. The buckets are then sorted and chunked into query key embedding vector chunks
each of length config.lsh_chunk_length. For each chunk, the query embedding vectors attend to its key vectors
(which are tied to themselves) and to the key embedding vectors of config.lsh_num_chunks_before previous
neighboring chunks and config.lsh_num_chunks_after following neighboring chunks.
For more information, see the original Paper or this great blog post.
Note that config.num_buckets can also be factorized into a list \((n_{\text{buckets}}^1,
n_{\text{buckets}}^2)\). This way instead of assigning the query key embedding vectors to one of \((1,\ldots,
n_{\text{buckets}})\) they are assigned to one of \((1-1,\ldots, n_{\text{buckets}}^1-1, \ldots,
1-n_{\text{buckets}}^2, \ldots, n_{\text{buckets}}^1-n_{\text{buckets}}^2)\). This is crucial for very long sequences to
save memory.
When training a model from scratch, it is recommended to leave config.num_buckets=None, so that depending on the
sequence length a good value for num_buckets is calculated on the fly. This value will then automatically be
saved in the config and should be reused for inference.
Using LSH self attention, the memory and time complexity of the query-key matmul operation can be reduced from
\(\mathcal{O}(n_s \times n_s)\) to \(\mathcal{O}(n_s \times \log(n_s))\), which usually represents the memory
and time bottleneck in a transformer model, with \(n_s\) being the sequence length.
Local Self Attention
Local self attention is essentially a "normal" self attention layer with key, query and value projections, but is
chunked so that in each chunk of length config.local_chunk_length the query embedding vectors only attends to
the key embedding vectors in its chunk and to the key embedding vectors of config.local_num_chunks_before
previous neighboring chunks and config.local_num_chunks_after following neighboring chunks.
Using Local self attention, the memory and time complexity of the query-key matmul operation can be reduced from
\(\mathcal{O}(n_s \times n_s)\) to \(\mathcal{O}(n_s \times \log(n_s))\), which usually represents the memory
and time bottleneck in a transformer model, with \(n_s\) being the sequence length.
Training
During training, we must ensure that the sequence length is set to a value that can be divided by the least common
multiple of config.lsh_chunk_length and config.local_chunk_length and that the parameters of the Axial
Positional Encodings are correctly set as described above. Reformer is very memory efficient so that the model can
easily be trained on sequences as long as 64000 tokens.
For training, the [ReformerModelWithLMHead] should be used as follows:
python
input_ids = tokenizer.encode("This is a sentence from the training data", return_tensors="pt")
loss = model(input_ids, labels=input_ids)[0]
Resources