PaLM-flax VS x-transformers

Google, Deepmind, and OpenAI normally provide a section in their research papers for replicating the pre-training and fine-tuning architectures of the models. For example, a replication of the pre-training architecture outlined in the LaMDA research paper in PyTorch (https://github.com/conceptofmind/LaMDA-pytorch/blob/main/lamda\_pytorch/lamda\_pytorch.py) or another implementation of Google's SOTA Pathways Language Model in JAX/FLAX (https://github.com/conceptofmind/PaLM-flax).

it is all documented here, in writing and in code https://github.com/lucidrains/x-transformers

you will want to use rotary embeddings, if you do not need length extrapolation

it is just a single transformer encoder, so just use https://github.com/lucidrains/x-transformers with ff_glu set to True

However, I've come across the concept of Key, Value Caching in Transformer-Decoders recently (e.g. Figure 3 here), wherein because each output (and hence each input, since the model is autoregressive) only depends on previous outputs (inputs), we don't need to re-compute Key and Value vectors for all t < t_i at timestep i of the sequence. My intuition leads me to believe, then, that (unconditioned) inference for a decoder-only model uses an effective sequence length of 1 (the most recently produced token is the only real input that requires computation on), making Attention a linear-complexity operation. This thinking seems to be validated by this github issue, and this paper (2nd paragraph of Introduction).

The Project - Model: The primary architecture consists of a CNN with a transformer encoder and decoder. At first, I used my implementation of self-attention. Still, due to it not converging, I switched to using x-transformer implementation by lucidrains - as it includes improvements from many papers. The objective is simple; the CNN encoder converts images to a high-level representation; feeds them to the transformer encoder for information flow. Finally, a transformer decoder tries to decode the text character-by-character using autoregressive loss. After two weeks of trying around different things, the training did not converge within the first hour - as this is the usual mark I use to validate if a model is learning or not.

X-Transformers: A fully-featured transformer with experimental features\ (25 comments)

One thing I’ve looked at is the fact that there’s no obvious reason to distinguish between W_K and W_Q in the formulation of a transformer as far as I can tell. However if you build a transformer where you merge the two matrices, it doesn’t learn as well. It still learns, but not as well. You can try out the code here. The training loss can be seen here, though we aborted the run because of how poorly it was doing.