ME-ViT: A Single-Load Memory-Efficient FPGA Accelerator for Vision Transformers

原文链接: arXiv:2402.09709 PDF

作者: Kyle Marino, Pengmiao Zhang, Viktor K. Prasanna (University of Southern California)

摘要

Vision Transformers (ViTs) 已成为目标分类任务的最先进解决方案，但其巨大的计算需求和参数量使其难以实现实时推理。现有 ViT 硬件加速器面临频繁的外部内存访问问题，限制了可达到的吞吐量。本文提出 ME-ViT，一种新颖的内存高效 FPGA 加速器，通过单次加载策略最小化内存流量：模型参数仅加载一次，中间结果存储在片上，所有操作在单一处理单元中实现。

核心成果:

  • 内存带宽提升 9.22× (32阵列) 和 17.89× (16阵列)
  • 每 DSP 吞吐量提升 2.16×
  • 能效提升最高达 4.00× (相比 GPU)
  • 在 Xilinx Alveo U200 上可实现 5 个 ME-PE 实例，吞吐量提升 5.10×

1. 问题定义

1.1 Vision Transformer 的内存瓶颈

“Computing capabilities of hardware such as GPUs and TPUs have outpaced memory bandwidth improvements, resulting in a poor Compute-to-Communication (C2C) ratio and limited model performance”

Vision Transformers 虽然在图像分类任务上取得了卓越性能，但其实现面临严峻挑战：

  1. 巨大的参数量 - ViT-Base 模型包含超过 86M 参数
  2. 高内存带宽需求 - 自注意力机制需要频繁的矩阵乘法
  3. 内存墙限制 - GPU/TPU 的计算能力远超内存带宽提升速度

1.2 现有方法的局限

现有 ViT 加速方法包括：

  • 模型压缩 (pruning, distillation)
  • 权重量化 (INT8, INT4)
  • 算法优化 (稀疏注意力)

然而，这些方法没有直接解决核心问题：内存带宽瓶颈。块矩阵乘法 (BMM) 方法虽然将大矩阵分块，但仍需要频繁地将块写回和重新加载到片外内存。

1.3 FPGA 的机会与挑战

FPGA 提供了高计算并行性和可定制架构，但同样面临内存带宽限制。如图 1 所示，ViT 和 DeiT 模型在没有专门优化的情况下严重受限于内存带宽。

2. 方法框架

2.1 单次加载策略 (Single-Load Policy)

ME-ViT 的核心设计原则是单次加载策略：

  1. 模型参数仅加载一次 - 所有权重在推理开始时加载到片上 BRAM
  2. 中间结果片上存储 - 避免写回片外内存
  3. 单一处理单元执行所有操作 - 线性投影、自注意力、MLP 在同一架构上执行

这与传统的块矩阵乘法方法不同，后者需要在不同块之间重新加载相同的权重。

2.2 内存高效处理单元 (ME-PE)

为实现单次加载策略，作者设计了 ME-PE (Memory-Efficient Processing Element)：

多用途缓冲区设计

“Multi-purpose buffers are designed to achieve the goals of the single-load policy while addressing its requirement of large BRAM allocations”

关键创新：

  • 缓冲区复用 - 同一缓冲区在不同计算阶段存储不同数据
  • 战略排序 - 精心安排计算顺序以最小化资源使用
  • BRAM 高效打包 - 为不同计算阶段重新打包和复用 BRAM

集成 Softmax 和 LayerNorm

ME-PE 将 Softmax 和 LayerNorm 函数集成到处理单元中：

  • 最小化矩阵乘法之间的停顿
  • 减少与主机 CPU 的通信
  • 避免中间结果写回内存

2.3 支持的 ViT 操作

ME-PE 执行 ViT 推理的三个关键操作：

  1. 线性投影 (Linear Projection, LP) - 将输入映射到 Q, K, V
  2. 多头自注意力 (Multi-headed Self-Attention, MSA) - 核心注意力计算
  3. 多层感知机 (Multi-Layer Perceptron, MLP) - 前馈网络

通过为 MSA 操作设计最小化架构，并将其复用于 MLP 计算，无需额外的 BRAM 资源。

3. 实验结果

3.1 实验设置

硬件平台: Xilinx Alveo U200 FPGA 评估模型:

  • ViT-Base (ViT-B)
  • DeiT-Base, DeiT-Small, DeiT-Tiny (DeiT-B, DeiT-S, DeiT-T)

阵列配置:

  • 32×32 脉动阵列
  • 16×16 脉动阵列

3.2 性能对比

指标 ME-ViT (32) ME-ViT (16) SOTA FPGA 提升幅度
内存带宽效率 优化 优化 瓶颈 9.22× / 17.89×
每 DSP 吞吐量 峰值 峰值 受限 2.16×
能效 (vs GPU) 4.00× - 1.00× 4.00×
能效 (vs FPGA) 1.03× - 1.00× 1.03×

3.3 Roofline 模型分析

如图 1 所示，ME-ViT 优化了内存带宽，使其在 GOPS (每秒十亿次操作) 上达到接近峰值的性能，而现有最先进实现受限于内存带宽。

3.4 多实例扩展

在 Xilinx Alveo U200 上：

  • 可实现 5 个 ME-PE 实例
  • 吞吐量提升 5.10× (相比 SOTA FPGA)
  • 能效提升 5.85× (相比 GPU) 和 1.51× (相比 FPGA)

4. 技术亮点

4.1 核心创新

  1. 单次加载策略 - 从根本上消除重复内存访问
  2. 多用途缓冲区 - 通过复用实现大容量片上存储
  3. 统一处理单元 - 单一架构支持所有 ViT 操作
  4. 函数集成 - Softmax/LayerNorm 嵌入 PE 减少通信

4.2 设计权衡

设计选择 优势 代价
大 BRAM 分配 避免片外访问 限制 PE 实例数量
缓冲区复用 减少资源使用 复杂的数据流控制
单一 PE 架构 简化设计 操作间需重新配置

4.3 适用场景

最适合:

  • 边缘 FPGA 设备（内存带宽受限）
  • 实时图像分类应用
  • 功耗敏感的部署环境

限制:

  • 需要足够的片上 BRAM
  • 模型大小受限于 FPGA 资源
  • 不适用于超大模型（如 ViT-Huge）

5. 对 AI 硬件设计的启示

5.1 内存优先设计

ME-ViT 证明：对于 Transformer 类模型，内存带宽优化比计算优化更重要。未来的 AI 加速器应该：

  • 优先考虑数据局部性
  • 最大化片上存储利用率
  • 减少片外内存访问次数

5.2 专用 vs 通用架构

“Custom FPGA architectures can meet the computational and memory demands of ViT inference more effectively than the general-purpose architecture of a GPU”

对于特定工作负载（如 ViT 推理），专用架构可以显著优于通用 GPU。这支持了领域专用加速器 (DSA) 的发展趋势。

5.3 边缘 AI 的机会

ME-ViT 的能效优势（4× 相比 GPU）使其特别适合边缘部署：

  • 低功耗设备上的实时推理
  • 无需云端连接的本地处理
  • 隐私敏感应用（数据不上传）

6. 局限与未来方向

6.1 当前局限

  1. 模型大小限制 - 受 FPGA BRAM 容量限制
  2. 仅支持推理 - 不支持训练
  3. 特定模型优化 - 主要针对 DeiT 和 ViT-Base

6.2 未来研究方向

  1. 扩展到大模型 - 支持 ViT-Large, ViT-Huge
  2. 训练支持 - 探索梯度计算和权重更新
  3. 其他 Transformer 变体 - Swin, DETR 等
  4. 稀疏性利用 - 结合稀疏注意力机制
  5. 混合精度 - 支持 INT8/INT4 量化

7. 总结

ME-ViT 通过单次加载策略和内存高效处理单元设计，成功解决了 Vision Transformer 在 FPGA 上的内存带宽瓶颈问题。其核心贡献在于：

  • 9.22× 内存带宽效率提升
  • 2.16× 每 DSP 吞吐量提升
  • 4.00× 能效提升 (相比 GPU)

这项工作为边缘设备上的高效 ViT 推理提供了可行方案，并强调了内存优化在 Transformer 加速器设计中的关键作用。

参考文献

  1. K. Marino, P. Zhang, and V. K. Prasanna, “ME-ViT: A Single-Load Memory-Efficient FPGA Accelerator for Vision Transformers,” arXiv:2402.09709, 2024.
  2. A. Vaswani et al., “Attention is All You Need,” NeurIPS, 2017.
  3. A. Dosovitskiy et al., “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale,” ICLR, 2021.
  4. T. Xiao et al., “Early Convolutions Help Transformers See Better,” NeurIPS, 2021.

本文使用 paper_to_blog 工具自动生成