arXiv preprint ·

M3Former: Fine-Grained Cross-Layer Mixing for Transformers

Jiasong Han, Yunkai Chai, Tong Zhu, Hengyu Zhang, Wei Fang, Jintao Wang, Guoqi Li and Yu Cheng

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Machine Learning

Abstract

Residual connections give each Transformer block access to all earlier layers only through one fixed sum, broadcast to all of the block's inputs. This broadcast bottleneck imposes three constraints: the query, key, value and residual inputs of a block, which we call its roles, read the same vector; all hidden dimensions within a role share one set of weights over earlier layers; and every earlier layer has the same fixed weight for every token. We propose the Multi-Role Multi-Group Mixing Residual (M3), which removes all three: every role, and every group of hidden dimensions within a role, reads its own token-dependent mixture of all earlier layers. M3 starts as an exact residual connection and adds as little as 1.4% parameters and 1.8% FLOPs. From 100M to 1.4B parameters, M3Former outperforms the Transformer and coarse-grained cross-layer connections at every size, matching a Transformer 1.19-1.34x the size of the one it modifies, trained with 1.4-1.8x the compute. Relaxing only some of the constraints recovers only part of this gain. M3Former also improves sparse mixture-of-experts models of up to 4.6B parameters and trains more stably in deeper networks. Code: https://github.com/Kingchet10/m3former

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Machine Learning

Get the few papers worth your time in Machine Learning every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics