交换激活值直接查表恢复KV，TRACE让分布式Transformer推理通信计算重叠！
TRACE: Transformer Residual Activation Codebook Encoding for Communication-efficient Distributed Edge Inference
📖 AI 总结
本文提出TRACE框架，用于解决分布式Transformer推理中因设备间交换激活值而导致的通信瓶颈。现有压缩方案虽能降低传输开销，但接收方仍需对收到的数据执行key/value投影才能进行注意力计算，使计算被串行地卡在通信之后。TRACE在每层Transformer的注意力后残差处引入向量量化器，并预计算以码本索引为键的key/value表示；推理时仅传输紧凑的码本索引，接收方通过查表直接恢复key和value，从而省去下游计算，实现通信与计算重叠。作者还利用层次聚类将码本组织为多分辨率结构，支持不重新训练即可动态调整通信比特率，并提供侧重精度与兼顾效率两种运行模式。在ViT-Base（ImageNet-1K、CIFAR-100）和ViViT-B（UCF-101）上的实验显示，TRACE相比最先进的基于VQ的分布式推理框架精度最高提升1.89%，端到端推理速度最高提升1.27倍；相比未压缩的序列并行推理，通信开销最高降低22624倍。
🔑 关键词速览
| Transformer | 一种基于自注意力机制的神经网络架构，广泛应用于自然语言处理和计算机视觉任务。 |
| Vector Quantizer (VQ) | 一种将连续向量映射到离散码本索引的量化技术，用于压缩数据表示。 |
| 序列并行 | 一种分布式推理策略，将输入序列分割到多个设备上并行处理，设备间需要交换激活值。 |
| 键/值投影 | 在Transformer注意力机制中，将输入激活值线性变换为键（Key）和值（Value）向量的计算过程。 |
| 层次聚类 | 一种将数据组织成多级树状结构的方法，用于构建多分辨率码本以支持自适应比特率。 |
📱 每天一份 AI 前沿日报
关注公众号，每天 09:00 推送 · 不错过任何重磅