Transformer models are deployed in a wide range of settings, from
multi-accelerator clusters to standalone mobile phones. The diverse inference
constraints in these scenarios necessitate practitioners to train foundation
models such as PaLM 2, Llama, & ViTs as a series of models of varying sizes.
Due to significant training costs, only a select few model sizes are trained
and supported, limiting more fine-grained control over relevant tradeoffs,
including latency, cost, and accuracy. This work introduces MatFormer, a nested
Transformer architecture designed to offer elasticity in a variety of
deployment constraints. Each Feed Forward Network (FFN) block of a MatFormer
model is jointly optimized with a few nested smaller FFN blocks. This training
procedure allows for the Mix'n'Match of model granularities across layers --
i.e., a trained universal MatFormer model enables extraction of hundreds of
accurate smaller models, which were never explicitly optimized. We empirically
demonstrate MatFormer's effectiveness across different model classes (decoders
& encoders), modalities (language & vision), and scales (up to 2.6B
parameters). We find that a 2.6B decoder-only MatFormer language model (MatLM)
allows us to extract smaller models spanning from 1.5B to 2.6B, each exhibiting
comparable validation loss and one-shot downstream evaluations to their
independently trained counterparts. Furthermore, we observe that smaller
encoders extracted from a universal MatFormer-based ViT (MatViT) encoder
preserve the metric-space structure for adaptive large-scale retrieval.
Finally, we showcase that speculative decoding with the accurate and consistent
submodels extracted from MatFormer can further reduce inference latency.

MatFormer 介绍了一种嵌套的 Transformer 架构，允许精细控制模型大小和推理延迟，同时能提取准确的子模型以降低推理时间。

MatFormer: 嵌套变压器用于弹性推理

MatFormer: Nested Transformer for Elastic Inference

We consider representation learning on periodic graphs encoding crystal
materials. Different from regular graphs, periodic graphs consist of a minimum
unit cell repeating itself on a regular lattice in 3D space. How to effectively
encode these periodic structures poses unique challenges not present in regular
graph representation learning. In addition to being E(3) invariant, periodic
graph representations need to be periodic invariant. That is, the learned
representations should be invariant to shifts of cell boundaries as they are
artificially imposed. Furthermore, the periodic repeating patterns need to be
captured explicitly as lattices of different sizes and orientations may
correspond to different materials. In this work, we propose a transformer
architecture, known as Matformer, for periodic graph representation learning.
Our Matformer is designed to be invariant to periodicity and can capture
repeating patterns explicitly. In particular, Matformer encodes periodic
patterns by efficient use of geometric distances between the same atoms in
neighboring cells. Experimental results on multiple common benchmark datasets
show that our Matformer outperforms baseline methods consistently. In addition,
our results demonstrate the importance of periodic invariance and explicit
repeating pattern encoding for crystal representation learning.

本研究在周期性图上考虑了表示学习问题以编码晶体材料。针对周期性图的独特挑战，我们提出了一种名为 Matformer 的转换器体系结构，用于周期性图的表示学习，它具有周期性不变性和显式重复模式编码的特点。我们的实验结果表明，相对于基线方法，Matformer 在多个公共基准数据集上始终表现更好。此外，我们的结果还表明了周期性不变性和显式重复模式编码对于晶体表示学习的重要性。