[论文解读] DataMUX: Data Multiplexing for Neural Networks

Vishvak Murahari, Carlos Jimenez-Gomez+2 more|arXiv (Cornell University)|Feb 18, 2022

Topic Modeling被引用 9

一句话总结

DataMUX 提出一种技术，使深度神经网络能够通过单一紧凑的混合表示，同时处理多个输入。该技术利用多路复用层将输入组合，再通过解多路复用层恢复各个独立的预测结果。在 MNLI 和 SST-2 等任务上，该方法实现高达 18 倍的推理吞吐量提升，且精度下降不足 4%，证明了其在 Transformer 架构中的可行性，以及在 MLP 和 CNN 中的有限适用性。

ABSTRACT

In this paper, we introduce data multiplexing (DataMUX), a technique that enables deep neural networks to process multiple inputs simultaneously using a single compact representation. DataMUX demonstrates that neural networks are capable of generating accurate predictions over mixtures of inputs, resulting in increased throughput with minimal extra memory requirements. Our approach uses two key components -- 1) a multiplexing layer that performs a fixed linear transformation to each input before combining them to create a mixed representation of the same size as a single input, which is then processed by the base network, and 2) a demultiplexing layer that converts the base network's output back into independent representations before producing predictions for each input. We show the viability of DataMUX for different architectures (Transformers, and to a lesser extent MLPs and CNNs) across six different tasks spanning sentence classification, named entity recognition and image classification. For instance, DataMUX for Transformers can multiplex up to $20$x/$40$x inputs, achieving $11$x/$18$x increase in throughput with minimal absolute performance drops of $<2\%$ and $<4\%$ respectively on MNLI, a natural language inference task. We also provide a theoretical construction for multiplexing in self-attention networks and analyze the effect of various design elements in DataMUX.

研究动机与目标

• 探索深度神经网络是否能够使用单一紧凑的混合表示，同时处理多个输入，且性能下降不明显。
• 设计一种方法，在最小内存和计算开销下提升推理吞吐量。
• 证明数据多路复用在多种架构（Transformer、MLP、CNN）和任务（NLP、视觉）中的可行性。
• 分析自注意力网络中数据多路复用的理论与实证极限。
• 提供一种端到端训练多路复用模型的框架，且参数增加极少。

提出的方法

• 多路复用层对每个输入应用固定的线性变换（例如，随机正交或低秩变换），然后将它们合并为一个与单个输入大小相同的单一混合表示。
• 基础神经网络（例如，Transformer、MLP、CNN）将混合表示作为单个输入进行处理。
• 解多路复用层将网络输出重构为每个原始输入的独立表示。
• 多路复用和解多路复用层是可微的，支持通过标准反向传播进行联合训练。
• 引入一种特殊的标记检索预训练目标，以提升多路复用模型的训练稳定性。
• 该方法应用于六个任务：句子分类、命名实体识别和图像分类，并对多路复用策略和模型组件进行了消融研究。

实验结果

研究问题

• RQ1 深度神经网络能否通过单一紧凑的混合表示，同时处理多个输入并实现准确预测？
• RQ2 数据多路复用在不同架构中对推理吞吐量和模型精度有何影响？
• RQ3 在自注意力网络中，实现成功数据多路复用的架构和训练条件是什么？
• RQ4 不同的多路复用与解多路复用策略对 MLP、CNN 和 Transformer 的性能有何影响？
• RQ5 数据多路复用在多样化的 NLP 和视觉任务中具有多大程度的泛化能力？

主要发现

• DataMUX 使 Transformer 架构能够同时处理最多 40 个输入，在 MNLI 任务上仅导致 4% 的精度下降，推理吞吐量最高提升 18 倍。
• 在 SST-2 任务上，处理 40 个输入时，多路复用的 Transformer（T-MUX）保持了原始 Transformer 98% 的精度。
• 对于 MLP，采用低秩变换的多路复用在 N=8 时达到约 78% 的准确率，而基线模型为约 95%，表明性能有中等程度下降。
• CNN 可以在 N=4 时实现超过 80% 的性能，但超过此规模后性能急剧下降，尤其是使用正交变换时，会破坏空间局部性。
• 该方法仅引入极小的内存和计算开销，因为混合表示的大小与单个输入相同，无需增加批处理维度，即可实现高吞吐量。
• 理论分析表明，自注意力机制可在 N 个独立子空间中处理输入，支持注意力网络中多路复用的可行性。