[论文解读] MatFormer: Nested Transformer for Elastic Inference

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

Opportunistic and Delay-Tolerant NetworksComputer Science被引用 3

一句话总结

MatFormer 引入了一种嵌套的 Transformer 架构，通过在每个前馈网络块内联合训练多个不同大小的子网络，实现了弹性推理。这使得无需额外训练即可自由提取数百个准确且行为一致的子模型，在语言和视觉任务中实现了与独立训练模型相当的性能，包括在推测解码和检索任务中的最先进结果。

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/

研究动机与目标

• 解决在移动设备和大规模计算集群等多样化推理环境中部署基础模型时面临的高训练成本和模型粒度有限的问题。
• 克服不同大小独立训练模型之间在行为一致性与推理优化兼容性方面的不足。
• 通过从单一训练好的通用模型中提取数百个准确的子模型，实现对延迟、成本与准确率权衡的细粒度控制。
• 在不重新训练的前提下，保持在不同模型尺寸、模态（语言与视觉）和任务上的高性能。
• 通过在子模型之间保持度量空间结构和模型一致性，支持推测解码与自适应检索。

提出的方法

• 在 Transformer 的前馈网络（FFN）模块中引入嵌套结构，使较小的子模型嵌入到较大的子模型中，形成层级结构 T₁ ⊂ T₂ ⊂ … ⊂ T₉。
• 通过在单次训练过程中组合所有嵌套子模型的独立损失，联合训练所有子模型，确保参数共享与一致性。
• 应用套娃表示学习原理，按重要性从高到低组织神经元与注意力头，实现分层提取。
• 支持 Mix'n'Match 提取：跨层任意组合子模块可生成 gˡ 个不同的模型，指数级提升模型粒度。
• 聚焦于占非嵌入参数与延迟超过 60% 的 FFN 模块，以最大化推理效率提升。
• 在视觉编码器（MatViT）中保持度量空间结构，使其可有效用于自适应检索系统。

实验结果

研究问题

• RQ1 是否可以训练一个单一的通用 Transformer 模型，使其在无需额外训练的情况下支持数百个准确且行为一致的、不同大小的子模型？
• RQ2 与独立训练的模型相比，MatFormer 提取的子模型在下游自然语言处理和视觉任务中的表现如何？
• RQ3 嵌套架构在视觉编码器中在多大程度上保持了度量空间结构，以支持检索应用？
• RQ4 使用提取的子模型进行推测解码是否能降低推理延迟，同时保持准确性？
• RQ5 嵌套子模型的联合训练在多大程度上影响了模型在不同粒度上的泛化能力？

主要发现

• 一个 2.6B 参数的 MatLM 解码器模型可提取从 1.5B 到 2.6B 参数的子模型，其验证损失和零样本下游性能均与独立训练的对应模型相当。
• 从 MatViT 提取的子模型保持了度量空间结构，使其在自适应大规模检索中表现有效。
• 使用从 MatFormer 提取的准确且一致的子模型进行推测解码，可降低推理延迟而不损失准确性。
• 在 TriviaQA 基准测试中，一个 1.55B 的 MatFormer 子模型达到了 26.15% 的 EM 分数，与一个 1.65B 的独立训练模型性能相当。
• 在 14 个下游任务中，提取的 1.55B 子模型的平均 EM 得分为 48.39%，接近全模型 2.6B 的 49.18%。
• 提取的 1.55B 子模型在开发集上的对数困惑度为 2.663，与全模型的 2.65 非常接近，表明其具有强大的泛化能力。