MatFormer: Nested Transformer for Elastic Inference

Abstract: Foundation models are applied in a broad spectrum of settings with different inference constraints, from massive multi-accelerator clusters to resource-constrained standalone mobile devices. However, the substantial costs associated with training these models often limit the number of unique model sizes that can be offered. Consequently, practitioners are compelled to select a model that may not be optimally aligned with their specific latency and cost requirements. We present MatFormer, a novel Transformer architecture designed to provide elastic inference across diverse deployment constraints. MatFormer achieves this by incorporating a nested Feed Forward Network (FFN) block structure within a standard Transformer model. During training, we optimize the parameters of multiple nested FFN blocks with varying sizes, enabling the extraction of hundreds of accurate smaller models without incurring additional computational costs. We empirically validate the efficacy of MatFormer across different model classes (decoders and encoders) and modalities (language and vision), demonstrating its potential for real-world deployment. We show that a 850M decoder-only MatFormer LLM (MatLM) allows us to extract multiple smaller models spanning from 582M to 850M parameters, each exhibiting better validation loss and one-shot downstream evaluations than 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 lead to significant reduction in inference latency. Project website: https://devvrit.github.io/matformer/

• The paper introduces MatFormer, a nested transformer architecture that extracts accurate submodels without additional retraining.
• It demonstrates competitive accuracy and superior scalability across language and vision tasks by balancing computation and performance.
• Experimental results highlight improved resource efficiency and adaptive performance for elastic inference in diverse deployment scenarios.

MatFormer: Nested Transformer for Elastic Inference

The paper "MatFormer: Nested Transformer for Elastic Inference" proposes a novel architecture in the domain of transformer-based models to address the critical challenge of adaptability and elasticity in diverse deployment environments. Traditional transformer models, such as those used in LLMs or vision transformers (ViTs), require a predefined model size for each deployment scenario, thus necessitating a series of independently trained models. This approach comes with significant training overheads and limited flexibility, especially when fine-grained control over trade-offs between latency, cost, and accuracy is required.

Key Contributions

1. Introduction of MatFormer:

MatFormer is introduced as a nested transformer architecture facilitating elastic inference. Each feed-forward network (FFN) block in a MatFormer incorporates a few nested smaller FFNs, enabling the extraction of hundreds of accurate submodels without additional retraining. This inherently nested structure offers unprecedented flexibility, allowing practitioners to tailor the model granularity dynamically based on deployment constraints.

2. Empirical Validation Across Modalities:

The authors empirically validate MatFormer across multiple model classes (decoders and encoders), modalities (language and vision), and scales (up to 2.6 billion parameters). For LLMs, MatFormer-based LLMs (MatLMs) are benchmarked against traditional independently trained baseline models. For vision models, MatFormer-based Vision Transformers (MatViTs) are tested on tasks such as image classification and retrieval. The results demonstrate that MatFormer not only matches the accuracy of the baseline models but also exhibits better scalability and flexibility.

3. Speculative Decoding and Elastic Encoders:

The paper showcases how MatFormer submodels can be utilized for faster autoregressive generation through speculative decoding, leveraging the consistent behavior of the smaller submodels with the largest model. Additionally, MatFormer-based encoders are shown to enable elastic query encoding for adaptive dense retrieval, reducing compute overhead significantly while maintaining high accuracy.

Experimental Findings

LLMs (MatLMs):

For MatLMs, spanning scales from 78M to 2.6B parameters, the authors report that the models trained with MatFormer architecture generalize well and provide competitive performance compared to their baseline counterparts. Specifically:

• The validation loss and downstream evaluation scores of MatLM submodels are comparable to those of independently trained models.
• MatFormer’s Mix’n’Match capability allows extracting numerous models along the accuracy-compute curve, providing a fine-grained balance without additional training costs.
• Consistency metrics reveal that submodels extracted from MatFormer are significantly more consistent, enhancing their utility in speculative decoding.

Vision Transformers (MatViTs):

For MatViTs, the experiments conducted on ImageNet-1K reveal:

• MatViT models often outperform the corresponding baseline ViT models.
• The ability to adaptively use Mix’n’Match models enhances elastic inference, leading to better utilization of available computational resources while preserving accuracy.
• For large-scale adaptive image retrieval, MatViTs demonstrate the capability to preserve metric-space consistency, allowing real-time adaptive query encoding.

Implications

Practical Implications:

MatFormer architecture addresses the pressing need for adaptable, efficient models capable of catering to diverse deployment scenarios, from mobile devices with limited computational power to large-scale multi-accelerator clusters. By providing a single universal model that can dynamically adjust its computational requirements, MatFormer reduces the necessity to train and maintain multiple model versions, significantly optimizing resource usage.

Theoretical Implications:

The nested structure of MatFormer challenges the conventional independent training paradigm, proposing a shift towards joint optimization of model granularities. This could pave the way for future research into more generalized and universally adaptable model architectures, potentially influencing how both foundational and specialized models are designed and trained.

Future Directions

Several future research directions stem from this work:

• Hyperparameter optimization and initialization strategies: Fine-tuning the training procedure to address the limitations identified, such as improvement in embedding and token-level operations.
• Real-time adaptation algorithms: Developing efficient algorithms to dynamically select the best-performing model configuration from the nested submodels according to real-time constraints.
• Extension to other architectures: Exploring the adaptability of the nested structure in other neural network architectures beyond transformers.

In conclusion, MatFormer represents a significant advancement in the design of adaptable AI models, with practical benefits in deployment flexibility and resource efficiency. Its empirical success across multiple tasks and modalities suggests it as a promising direction for future research and application in AI deployment frameworks.