MatFormer: Nested Transformer for Elastic Inference (2310.07707v1)

Abstract: 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 LLM (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: 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.