MemoryFormer：透過移除全連接層來最小化 Transformer 的計算量

核心概念

MemoryFormer 是一種新穎的 Transformer 架構，它透過使用記憶體層和局部敏感雜湊算法取代全連接層，在不犧牲效能的情況下，顯著減少了 Transformer 模型的計算複雜度。

摘要

MemoryFormer: 透過移除全連接層來最小化 Transformer 的計算量

研究目標

本研究旨在解決大型語言模型 (LLM) 計算複雜度高的問題，提出了一種名為 MemoryFormer 的新型 Transformer 架構，透過移除全連接層來減少計算量，同時保持模型效能。

方法

MemoryFormer 使用記憶體層 (Memory Layer) 和局部敏感雜湊 (LSH) 算法來取代標準 Transformer 模型中的全連接層。記憶體層包含一組儲存大量離散向量的雜湊表。它使用 LSH 算法根據輸入詞嵌入動態地從雜湊表中檢索相關的向量子集。然後，將檢索到的向量以不同的權重聚合在一起，形成記憶體層的輸出。這種雜湊和聚合操作的結果提供了對全連接層中矩陣乘法結果的估計。

關鍵發現

• MemoryFormer 在多個自然語言處理基準測試中，與基準 Transformer 模型相比，計算複雜度顯著降低，同時保持了可比的效能。
• 與其他高效 Transformer 方法（如 Linformer、Cosformer 和 Performer）相比，MemoryFormer 在效能和效率方面均表現出色。
• 透過調整 Memory Layer 的超參數，可以在記憶體使用和效能之間取得平衡。

主要結論

MemoryFormer 為減少 Transformer 模型的計算複雜度提供了一種新的解決方案。透過移除全連接層，MemoryFormer 可以顯著降低計算量，同時保持模型效能。這項工作為下一代平行運算平台的硬體設計提供了指導意義，例如更大的匯流排寬度和更高的快取命中率。

意義

MemoryFormer 的提出對於促進大型語言模型在資源受限設備上的部署和應用具有重要意義。它為開發更高效、更輕量級的 Transformer 模型提供了新的思路。

局限性和未來研究方向

• MemoryFormer 的效能仍有提升空間，未來可以探索更先進的雜湊算法和記憶體層設計。
• MemoryFormer 的記憶體使用量仍然較大，未來可以研究如何進一步壓縮模型大小。
• MemoryFormer 目前主要應用於自然語言處理領域，未來可以探索其在其他領域（如電腦視覺）的應用。

統計資料

當序列長度為 2048，隱藏層大小為 2048 時，MemoryFormer 塊所需的浮點運算量僅為基準 Transformer 塊的 19%。
MemoryFormer-tiny、-small 和 -base 模型的雜湊表數量分別為 64、96 和 128。
考慮到雜湊表梯度的稀疏性，MemoryFormer 的學習率設定為基準 Pythia 模型學習率的 3 倍。
在 Memory Block 中，第一個 Memory Layer 的輸出維度設定為 (τ + 2) * K，其中 τ 是每個子向量的位元寬度，K 是雜湊表的數量。
实验结果表明，将 τ 设置为 8 是 MemoryFormer 的一个较好的选择。

引述

"相比於執行矩陣乘法，從記憶體中檢索數據塊是一個更經濟的操作，只需要很少的計算量。"
"這項工作不僅提出了一種不同於現有方法的新 FLOPs 降低策略，還為下一代平行運算平台的硬體設計（例如更大的匯流排寬度和更高的快取命中率）提供了指導意義。"

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

MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers

by Ning Ding, Y... 於 arxiv.org 11-21-2024

https://arxiv.org/pdf/2411.12992.pdf

深入探究

MemoryFormer 如何與其他模型壓縮技術（如剪枝和量化）相結合，以進一步減少模型大小和計算成本？

MemoryFormer 可以與其他模型壓縮技術（如剪枝和量化）相結合，進一步減少模型大小和計算成本：
1. 與剪枝技術結合:

結構化剪枝: 可以剪枝掉對模型貢獻度低的 Memory Layer 或 Hash Table，減少模型參數量和計算量。例如，可以根據每個 Hash Table 的平均激活值或梯度大小進行排序，剪枝掉排名靠後的 Hash Table。
非結構化剪枝: 可以剪枝掉 Hash Table 中不重要的向量，進一步壓縮模型。例如，可以根據每個向量的 L2 范數或對模型性能的影響進行排序，剪枝掉排名靠後的向量。
2. 與量化技術結合:

將 Hash Table 中的向量量化為低比特表示: 例如，使用 8 位整數或更低比特表示來存儲向量，可以顯著減少模型大小和內存佔用。
對 Memory Layer 中的計算進行量化: 例如，使用量化感知訓練（Quantization-Aware Training）技術，在訓練過程中模擬量化操作帶來的影響，可以減少量化誤差，提高模型性能。
3. 結合剪枝和量化:

可以先進行剪枝，再對剪枝後的模型進行量化，以最大程度地減少模型大小和計算成本。
需要注意的是，剪枝和量化技術可能會導致模型性能下降。因此，需要在模型壓縮和性能之間進行權衡，選擇合適的壓縮策略。

MemoryFormer 中使用的局部敏感雜湊算法是否會引入新的偏差，如何評估和減輕這種偏差？

是的，MemoryFormer 中使用的局部敏感雜湊算法會引入新的偏差，主要體現在以下兩個方面：

哈希碰撞: 不同的輸入向量可能被映射到同一個哈希桶，導致信息丢失，影響模型的表達能力。
量化誤差: Memory Layer 中使用有限的哈希桶來近似無限的向量空間，會引入量化誤差，影響模型的精度。

評估和減輕這種偏差的方法包括：

評估指標:

可以使用多個下游任務的性能指標來評估 MemoryFormer 的偏差，例如困惑度（Perplexity）、準確率（Accuracy）等。
可以比較 MemoryFormer 與標準 Transformer 在相同任務上的性能差異，量化偏差的影響。

減輕偏差:

增加哈希桶數量: 可以通過增加哈希桶的數量來降低哈希碰撞的概率，提高模型的表達能力。
使用更精確的哈希函數: 可以使用更精確的哈希函數，例如多個哈希函數的組合，來降低哈希碰撞的概率。
使用更精細的量化策略: 可以使用更精細的量化策略，例如非均勻量化，來降低量化誤差。
數據增強: 可以通過數據增強來增加訓練數據的多樣性，提高模型的泛化能力，從而減輕偏差。

如果將 MemoryFormer 的概念應用於其他需要大量矩陣運算的領域，例如圖神經網路或推薦系統，會產生什麼樣的影響？

將 MemoryFormer 的概念應用於其他需要大量矩陣運算的領域，例如圖神經網路或推薦系統，具有以下潛在影響：
1. 圖神經網路:

降低計算複雜度: 圖神經網路中的訊息傳遞步驟通常涉及節點特徵與鄰居節點特徵的聚合，這可以看作是一種矩陣乘法運算。 MemoryFormer 可以用於近似這些矩陣乘法，從而降低圖神經網路的計算複雜度，尤其是在處理大規模圖數據時。
提升效率: Memory Layer 可以高效地處理稀疏數據，這對於通常具有稀疏連接性的圖數據非常有利。
新的研究方向: MemoryFormer 可以啟發圖神經網路領域的新研究方向，例如探索更適合圖數據特性的哈希函數和 Memory Layer 設計。
2. 推薦系統:

加速推薦過程: 推薦系統中通常需要計算用戶與物品之間的相似度，這也涉及大量的矩陣運算。 MemoryFormer 可以用於近似這些矩陣運算，從而加速推薦過程，尤其是在面對海量用戶和物品數據時。
處理冷啟動問題: Memory Layer 可以通過哈希函數將新用戶或新物品映射到已有的向量空間，從而緩解推薦系統中的冷啟動問題。
個性化推薦: MemoryFormer 可以通過學習用戶和物品的哈希表示，捕捉用戶的個性化偏好，從而提高推薦的準確性。
總體而言，將 MemoryFormer 的概念應用於其他需要大量矩陣運算的領域具有很大的潛力，可以有效降低計算成本、提升效率，並帶來新的研究方向。 然而，也需要考慮不同領域數據特性的差異，對 MemoryFormer 進行相應的調整和優化，才能取得更好的效果。