EFFICIENT GENERATIVE AI BY TUCKER NETS

We propose a trainable token mixing transformation based on Tucker maps — which are of (quasi)linear complexity in the input. On experiments on ViT and GPTs we show that replacing attention layers by our Tucker maps caused $\sim 50$% reduction in the number of FLOPs used in training — while almost never causing any deterioration of performance. We verify that the advantages of Tucker Nets on transformers, hold for both Adam(W) and Muon being used for training. Our Tucker layers are designed as drop-in replacements of the attention layers in any transformer architecture and thus pave a new way to efficient generative AI.