基於 Top-k 記憶體內 ADC 的 Transformer 低能耗、低延遲推理：Topkima-Former

核心概念

本文提出了一種軟硬體協同設計方案 Topkima-Former，通過電路、架構和演算法層面的創新，有效降低了 Transformer 模型推理過程中的能耗和延遲。

摘要

Topkima-Former：一種低能耗、低延遲的 Transformer 推理方案

簡介

本文介紹了一種名為 Topkima-Former 的創新設計，旨在加速 Transformer 模型的推理過程，並降低其能耗和延遲。Transformer 模型在自然語言處理和電腦視覺領域展現出色的效能，但其複雜的結構和大量的非線性運算，例如 softmax，也帶來了巨大的計算成本和延遲。為了解決這些問題，Topkima-Former 從電路、架構和演算法三個層面進行了優化。

電路層面優化：Topkima 宏

Topkima-Former 的核心創新在於其電路設計，特別是 Topkima 宏。該宏結合了 top-k 激勵選擇和記憶體內 ADC (IMA)，實現了低能耗、低延遲的 softmax 運算，且無需排序延遲。

• Top-k 激勵選擇： Topkima 只將 k 個最大的激勵值傳遞給 softmax 計算模組，從而顯著降低了 softmax 的計算成本。
• 記憶體內 ADC (IMA)： IMA 技術將 ADC 整合到每個 IMC 宏的行中，允許位元單元在乘法累加 (MAC) 和 ADC 操作之間共用，從而提高了能源效率。
• 遞減斜坡電壓： Topkima 採用遞減斜坡電壓取代了傳統 IMA 中的遞增斜坡電壓，使得較大的 MAC 電壓能夠更早地被檢測到，從而實現提前停止數據轉換，進一步降低延遲和能耗。

演算法層面優化：TFCBP 訓練方法

為了減輕 top-k 方法中只保留 k 個最大值所帶來的精度損失，Topkima-Former 引入了一種改進的訓練方法：top-k 前向-完整反向傳播 (TFCBP)。

• 前向傳播階段： 僅使用 top-k 激勵值來計算 softmax 機率。
• 反向傳播階段： 所有激勵值都參與梯度計算。

這種訓練方法有效地減輕了 top-k 近似帶來的精度損失，實驗結果顯示，在 CIFAR-10、CIFAR-100 和 SQuAD 數據集上，當 k=5 時，ViT、distilBERT 和 BERT-base 模型的精度僅下降了 0.4% 到 1.2%。

架構層面優化：無縮放設計

Transformer 模型中的注意力機制需要進行縮放操作，這在硬體實現中會帶來額外的設計複雜度。Topkima-Former 提出了一種改進的無縮放設計，通過調整 WQ 權重來實現縮放功能，而無需額外的硬體開銷。

系統性能評估

通過對 BERT-base 模型在 SQUAD 數據集上的評估，Topkima-Former 展現出優異的性能：

• 速度提升： 相比於傳統的 softmax 宏和基於數位 top-k 的 softmax 宏，Topkima-SM 宏分別實現了約 15 倍和 8 倍的速度提升。
• 能效提升： Topkima-Former 的能效比現有的基於 IMC 的加速器高出 1.3 到 35 倍。

總結

Topkima-Former 是一種針對 Transformer 模型推理過程進行優化的軟硬體協同設計方案，通過電路、架構和演算法層面的創新，有效降低了模型推理過程中的能耗和延遲，同時保持了較高的精度。

統計資料

softmax 運算在 Transformer 模型推理中佔據高達 40% 的延遲。
使用 Top-k 方法，當 k=30 時，可以實現 2.6 倍的加速，但會損失一些精度。
Topkima-Former 在 CIFAR-10 數據集上使用 Top-5 方法，精度僅下降 0.4%。
Topkima-Former 在 CIFAR-100 和 SQuAD 數據集上使用 Top-5 方法，精度下降不超過 1.2%。
Topkima-SM 宏比傳統 softmax 宏速度快 15 倍，比基於數位 top-k 的 softmax 宏速度快 8 倍。
Topkima-Former 比現有的基於 IMC 的加速器能效高出 1.3 到 35 倍。
Topkima-Former 在 200 MHz 頻率下可以達到 6.70 TOPS 的吞吐量和 16.84 TOPS/W 的能效。

引述

"Transformer can achieve superior performance in various applications involving neural language processing (NLP) and computer vision (CV) due to the attention mechanism."
"However, the softmax operation consumes up to 40% inference time, severely impeding the efficient implementation of transformer [7]."
"Topkima-Former achieves 1.8×−84× faster and 1.3×−35× more EE than prior IMC-based accelerators."
"Compared to a conventional softmax macro and a digital top-k (Dtopk) softmax macro, our proposed topkima softmax macro (topkima-SM) achieves about 15× and 8× faster speed."

從以下內容提煉的關鍵洞見

Topkima-Former: Low-energy, Low-Latency Inference for Transformers using top-k In-memory ADC

by Shuai Dong, ... 於 arxiv.org 11-21-2024

https://arxiv.org/pdf/2411.13050.pdf

深入探究

Topkima-Former 主要針對 Transformer 模型進行優化，那麼它是否可以應用於其他類型的深度學習模型？

Topkima-Former 的核心設計理念是利用 top-k 選擇和 In-Memory ADC 來加速 softmax 計算，從而降低 Transformer 模型的能耗和延遲。 雖然 Topkima-Former 主要針對 Transformer 模型設計，但其設計理念可以應用於其他包含 softmax 操作且對延遲和能耗敏感的深度學習模型。
例如：

循環神經網絡 (RNN)：RNN 中的 softmax 層可以用 Topkima-Former 加速，特別是在處理長序列數據時，Topkima-Former 可以顯著降低計算成本。
卷積神經網絡 (CNN)：雖然 CNN 中 softmax 的使用不像 Transformer 那麼普遍，但一些 CNN 架構，例如用於圖像分類的模型，仍然會使用 softmax。 Topkima-Former 可以應用於這些模型，尤其是在資源受限的邊緣設備上。
然而，需要注意的是，Topkima-Former 的性能提升與模型結構和應用場景密切相關。 在應用於其他模型時，需要根據具體情況進行評估和調整。

Topkima-Former 雖然降低了能耗和延遲，但犧牲了一定的精度，在對精度要求極高的應用場景中，如何平衡性能和精度？

的確，Topkima-Former 在降低能耗和延遲的同時，會損失一定的精度。 對於精度要求極高的應用場景，可以考慮以下方法來平衡性能和精度：

增加 k 值: Topkima-Former 中的 k 值決定了選擇多少個最大的激活值進行 softmax 計算。 增加 k 值可以提高精度，但同時也會增加計算成本。 可以根據具體應用需求，通過調整 k 值來找到精度和性能之間的平衡點。
改進 Top-k 選擇算法: Topkima-Former 使用簡單的排序算法來選擇 top-k 激活值。 可以考慮使用更精確的 top-k 選擇算法，例如基於哈希或樹的算法，以減少精度損失。
混合精度計算: 可以將 Topkima-Former 與混合精度計算技術相結合。 例如，對於精度要求高的層使用高精度計算，而對於精度要求相對較低的層使用 Topkima-Former 進行加速。
模型微調: 在使用 Topkima-Former 進行模型壓縮後，可以使用原始高精度數據對模型進行微調。 這有助於在保持較低能耗和延遲的同時，進一步提高模型精度。
總之，在對精度要求極高的應用場景中，需要綜合考慮精度、能耗和延遲等因素，並根據具體情況選擇合適的優化策略。

如果未來出現新的記憶體技術或神經網路架構，Topkima-Former 的設計理念是否仍然適用？如何進行相應的調整和改進？

Topkima-Former 的設計理念是基於現有的 SRAM/RRAM 技術和 Transformer 架構。 然而，其核心思想，即利用 top-k 選擇和 In-Memory 計算來加速 softmax 操作，具有較強的通用性和可擴展性，可以適應未來新的記憶體技術和神經網路架構。
面對新的技術和架構，可以進行以下調整和改進：

適配新的記憶體技術: 如果出現新的記憶體技術，例如相變記憶體 (PCM) 或電阻式 RAM (ReRAM)，可以根據其特性調整 Topkima-Former 的電路設計和讀寫方案，以充分利用新技術的優勢。
適配新的神經網路架構: 如果出現新的神經網路架構，需要分析其計算模式和數據流特點，調整 Topkima-Former 的數據映射和計算流程，以適應新的架構。 例如，可以根據新的架構設計專用的 top-k 選擇和 softmax 計算單元。
結合新的算法: 可以將 Topkima-Former 與新的算法相結合，例如神經架構搜索 (NAS) 或剪枝算法，以進一步提高模型效率和性能。
總之，Topkima-Former 的設計理念具有良好的可擴展性和適應性。 通過不斷調整和改進，可以將其應用於更廣泛的應用場景，並持續發揮其在降低能耗和延遲方面的優勢。