[Paper Review] MatFormer: Nested Transformer for Elastic Inference

Devvrit Devvrit, Sneha Kudugunta+9 more|arXiv (Cornell University)|Oct 11, 2023

Opportunistic and Delay-Tolerant NetworksComputer Science3 citations

TL;DR

MatFormer introduces a nested Transformer architecture that enables elastic inference by jointly training a single universal model with multiple subnetworks of varying sizes within each feed-forward network block. This allows for the free extraction of hundreds of accurate, behaviorally consistent submodels without additional training, achieving performance comparable to independently trained models across language and vision tasks, including state-of-the-art results in speculative decoding and retrieval.

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 language model (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/

Motivation & Objective

• To address the high training cost and limited model granularity in deploying foundation models across diverse inference environments such as mobile devices and large-scale clusters.
• To overcome the lack of behavioral consistency and inference optimization compatibility between independently trained models of different sizes.
• To enable fine-grained control over latency, cost, and accuracy trade-offs by extracting hundreds of accurate submodels from a single trained universal model.
• To maintain high performance across diverse model sizes, modalities (language and vision), and tasks without retraining.
• To support speculative decoding and adaptive retrieval by preserving metric-space structure and model consistency across submodels.

Proposed method

• Introduces a nested structure in the Transformer's feed-forward network (FFN) block, where smaller submodels are embedded within larger ones, forming a hierarchy T₁ ⊂ T₂ ⊂ … ⊂ T₉.
• Jointly trains all nested submodels by combining their individual losses during a single training run, ensuring parameter sharing and consistency.
• Applies the matryoshka representation learning principle to organize neurons and attention heads from most to least significant, enabling hierarchical extraction.
• Enables Mix'n'Match extraction: arbitrary combinations of subblocks across layers yield gˡ distinct models, exponentially increasing model granularity.
• Focuses on the FFN block, which accounts for over 60% of non-embedding parameters and latency, to maximize inference efficiency gains.
• Preserves metric-space structure in vision encoders (MatViT), enabling effective use in adaptive retrieval systems.

Experimental results

Research questions

• RQ1 Can a single universal Transformer model be trained to support hundreds of accurate, behaviorally consistent submodels of varying sizes without additional training?
• RQ2 How well do extracted submodels from MatFormer perform compared to independently trained models on downstream NLP and vision tasks?
• RQ3 To what extent does the nested architecture preserve metric-space structure in vision encoders for retrieval applications?
• RQ4 Can speculative decoding with extracted submodels reduce inference latency while maintaining accuracy?
• RQ5 How does the joint training of nested submodels affect the model's ability to generalize across different granularities?

Key findings

• A 2.6B-parameter MatLM decoder model enables extraction of submodels from 1.5B to 2.6B parameters, each achieving validation loss and one-shot downstream performance comparable to independently trained counterparts.
• Extracted submodels from MatViT preserve metric-space structure, enabling effective adaptive large-scale retrieval.
• Speculative decoding using accurate and consistent submodels extracted from MatFormer reduces inference latency without sacrificing accuracy.
• On the TriviaQA benchmark, a 1.55B MatFormer submodel achieved 26.15% EM, matching the performance of a 1.65B independently trained model.
• The average EM score across 14 downstream tasks for extracted 1.55B submodels reached 48.39%, approaching the 49.18% of the full 2.6B model.
• The development set log perplexity of extracted 1.55B submodels was 2.663, very close to the full 2.6B model's 2.65, indicating strong generalization.