Meta HSTU：統一的生成式推薦
一、背景
在之前的文章中我們介紹了字節的 HLLM 方案（???字節 HLLM：大模型與推薦系統結合的新范式??），其中也提到了 Meta 的HSTU，本文中我們簡單介紹一些這篇論文。
對應的論文為：[2402.17152] Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations
對應的代碼庫為：GitHub - facebookresearch/generative-recommenders
二、摘要
大規模推薦系統具有以下特征：依賴于高基數、異質特征，并且需每天處理數百億用戶行為。盡管在海量數據（含數千種特征）上進行訓練，但多數工業界的深度學習推薦模型（DLRM）無法在計算規模上取得突破。受到 Transformer 在語言與視覺領域的啟發，作者重新審視了推薦系統的基本設計選擇，將推薦問題重新表述為生成式建模下的生成式任務（生成式推薦，Generative Recommender，GR），并提出了一種新型架構 —— HSTU（Hierarchical Sequential Transduction Units），專為處理高基數、非恒定的流式推薦數據場景。
在合成和公開數據集上的實驗表明，HSTU 相比基線模型在 NDCG 指標上最多提升了 65.8%，且在處理 8192 長的序列時，比基于 FlashAttention2 的 Transformer 快 5.3x - 15.2x。基于 HSTU 的生成式推薦系統（包含 1.5T 參數量）在在線 A/B 測試中指標提升 12.4%，并已部署于擁有數十億用戶的互聯網平臺的多項服務中。更重要的是，生成式推薦的模型呈現出跨三個數量級的冪律增長特性，直到與 GPT-3/LLaMa-2 的規模相當。
三、引言
推薦系統作為在線內容平臺與電子商務領域的核心技術，每日為數十億用戶提供個性化體驗，發揮著關鍵作用。當前最先進的推薦方法已基于 DLRM 發展了約十年。DLRM 的特點在于其利用了異質特征，包括數值特征（如計數器和比率）、嵌入向量以及分類特征（如創作者 ID、用戶 ID等）。由于每分鐘都有新內容和產品不斷加入，特征空間呈現出極高的基數，通常達到數十億級別。為充分利用數萬種此類特征，DLRM 采用多種神經網絡來組合特征、轉換中間表示并生成最終輸出。
盡管采用了廣泛的人工設計特征并基于海量數據進行訓練，大多數 DLRM 在實際工業應用中仍面臨計算效率低下的問題。這一局限性值得關注，且至今尚未得到有效解決。
受 Transformer 在語言與視覺領域取得成功的啟發，作者重新審視了現代推薦系統中的基礎設計選擇。作者觀察到，在處理億級用戶規模的推薦系統時，需克服三大挑戰。
- 首先，推薦系統中的特征缺乏顯式結構。盡管在小規模場景下已探索了序列化表達，但在工業級 DLRM 中，異質特征，如高基數 ID、交叉特征、計數器、比率等，扮演著關鍵角色。
- 其次，推薦系統采用的詞表量達數十億，且持續變化。與語言模型中使用的 10 萬級靜態詞表相比，億級動態詞表帶來了訓練難題，并因需以 Target-aware 方式評估數萬候選 Item，導致推理成本高昂。
- 最后，計算成本成為實現大規模序列模型的主要瓶頸。GPT-3 在數千塊 GPU 上歷經 1-2 個月，處理了總計 300B 個 Token。這一規模看似驚人，但與用戶行為規模相比則相形見絀。最大的互聯網平臺每日服務數十億活躍用戶，這些用戶每日與數十億帖子、圖片和視頻互動。用戶序列長度可達 10 的五次方。因此，推薦系統每日需處理的 Token 量，遠超語言模型在 1-2 個月內處理的數量級。
本文中，作者將用戶行為視為生成建模中的新型模態。作者的核心洞見包括：
- 在工業級推薦系統中，核心的排序與檢索任務在適當的新特征空間下可被轉化為生成建模問題；
- 這一范式使作者能夠系統性地利用特征、訓練及推理中的冗余，從而提升效率。
四、推薦系統作為生成式任務：從 DLRM 到 GR
4.1 在 DLRM 中統一異質特征空間
現代 DLRM 通常采用大量類別型（“稀疏”）和數值型（“密集”）特征進行訓練。在 GR 中，作者將這些特征整合并編碼為單一的時間序列，如下圖 Figure 2 所示。
類別型（“稀疏”）特征：其示例包括：用戶喜歡的物品、用戶關注的特定類別（如“戶外”）中的創作者、用戶語言、用戶加入的社區、請求發起的城市等。作者對這些特征進行序列化處理，具體步驟如下：首先，選擇最長的時序數據，通常通過合并表示用戶互動 Item 的特征來構建主時間序列。其余特征通常為隨時間變化緩慢的時間序列，如人口統計信息或用戶關注的創作者。通過保留每個連續片段的最早 Item 來壓縮這些時間序列，然后將結果合并到主時間序列中。鑒于這些時間序列變化極為緩慢，此方法并未顯著增加整體序列長度。
數值型（“密集”）特征：其示例包括：加權衰減計數器、比率等。例如，某一特征可能表示用戶過去在特定主題匹配項上的點擊率（CTR）。相較于類別特征，這些特征變化更為頻繁，幾乎每發生一次（User，Item）交互都可能發生變化。因此，從計算和存儲角度來看，完全序列化這些特征是不可行的。然而，一個關鍵發現是，進行這些聚合操作所基于的類別特征（如 Item 主題、位置）已在 GR 中被序列化并編碼。因此，可以使用足夠表達能力的生成式架構并結合目標感知（[1706.06978] Deep Interest Network for Click-Through Rate Prediction）公式的情況下，移除 GR 中的數值特征。通過增加序列長度和 GR 中的計算量，能夠有效地捕捉數值特征。
4.2 將排序和檢索重新表述為生成式任務
給定一個按時間順序排列的 n 個 Token 序列 x?, x?, ..., xn-1（其中 xi ∈ X），以及這些 Token 被觀察到的時間點 t?, t?, ..., tn-1，生成式任務將該輸入序列映射為輸出 Token y?, y?, ..., yn-1：
- 其中 yi ∈ X ∪ {?}），yi = ? 表示 yi未定義。
- 用 Φi ∈ Xc（Xc ? X）表示系統提供給用戶的 Content（例如，圖像或視頻）。由于新內容不斷生成，Xc 和 X 不是固定不變的。
- 用 ai 表示用戶對 Φi 做出某種 Action（例如，點贊、跳過、視頻完成并分享），其中 ai ∈ X。
- 用 nc 表示用戶已交互的 Content 總數。
在因果自回歸設置下，標準的排序和檢索任務可以定義為生成式任務（見如下圖 Table 1）。可以得出以下觀察結果：
- 檢索。在推薦系統的檢索環節，針對 Φi+1 ∈ Xc 學習一個分布 p(Φi+1 | ui)，其中 ui 表示用戶在第 i 個 Token 處的表征。典型的目標是通過選擇 arg maxΦ∈Xc p(Φ | ui) 來最大化某種獎勵。這與標準的自回歸設置有兩點不同。
a.首先，對于 xi 和 yi 的監督信號不一定就是 Φi+1，因為用戶可能對 Φi+1作出負面反應。
b.其次，當 xi+1 代表與非交互相關的分類特征（如人口統計信息）時，yi 是未定義的。
- 排序。在 GR 中，排序面臨獨特的挑戰，因為工業推薦系統通常需要“Target-aware”的公式化表示。在這種情境下，目標Φi+1與歷史特征的“交互”需要盡早發生，這在標準自回歸設置中是不可行的，因為“交互”通常發生較晚（例如，通過 Encoder 輸出后的 Softmax 函數）。作者通過在如下圖 Table 1 中交錯排列 Item 和 Action 來解決這一問題，使得排序可以表述為 p(ai+1 |Φ0 ,a0 ,Φ1 ,a1 ,...,Φi+1)（在分類特征之前）。在實踐中，作者應用一個小型神經網絡將 Φi+1 處的輸出轉換為多任務預測。重要的是，這使得能夠在一次處理中對所有nc 個交互應用目標感知的交叉注意力機制。?
如下圖 Table 10 所示為以前的序列推薦工作與 GR 的對比：
如下圖 Figure 8 所示：
- 左圖為傳統的序列推薦模型：輸入可以沒有 Item Action 信息 ai，或者將 ai 融合到 Item Content 信息 Φi 中，而輸出都是 Item Content 信息 Φi。
- 右圖為本文的 GR 模型：Action ai 和 Content Φi 同時作為輸入和輸出。其中
a.下一個 Action ai Token 預測任務可以理解為排序（Ranking）任務。
b.下一個 Content Φi 預測可以理解為檢索（Retrieval）任務。
4.3 生成式訓練
工業推薦系統通常采用流式訓練模式，其中每個樣本按其可用順序依次處理。在此模式下，基于 SelfAttention 機制的生成式架構的總計算需求與
成比例，其中 ni 為用戶 i 的序列長度（Token 數量），d 為 Embedding 維度。括號中的第一項源自 SelfAttention 機制，第二項則來自 MLP 層，其隱藏層大小為 O(dff) = O(d)。設 N = maxini，整體時間復雜度可簡化為 O(N3d + N2d2)，對于推薦系統場景來說成本過高。
為了以可擴展方式在長序列上訓練，作者轉向生成式訓練，從而將計算復雜度降低 O(N) 倍，如下圖頂部所示。通過這種方式，Encoder 成本可以在多個目標間分攤。更具體地，當以速率 su(ni) 采樣第 i 個用戶時，總訓練成本可以按
縮放，通過將 su(ni) 設為 1/ni，可降至 O(N2d + Nd2)。在工業規模系統中實現此采樣的一種方式是在用戶請求或會話結束時生成訓練樣本，從而使得
。
五、用于 GR 的高性能 SelfAttention Encoder
為了適應工業級推薦系統中詞表大且不固定的特性，并實現 GR 的規模化應用，作者提出了一種新的 Encoder 設計：HSTU，HSTU 由一系列通過殘差連接相互關聯的相同層堆疊而成。每一層包含 3 個子層：
- 公式 1：Pointwise 投影。
- 公式 2：空間聚合。
- 公式 3：Pointwise 變換。?
其中，fi(X) 表示 MLP；為降低計算復雜度并優化 Query Q(X)、Key K(X)、Value V(X) 及門控權重 U(X) 的 Batch 處理效率，f1 和 f2 均采用單一線性層，即 fi(X) = Wi(X) + bi, 并且會使用 Kernel 融合。?1 和 ?2 代表非線性函數，此處選用 SiLU；Norm 為 LayerNorm，rabp,t 表示相對注意力偏置，其整合了位置（p）與時序（t）信息。完整的符號信息如下圖 Table 9 所示：
如下圖 Figure 3 所示，DLRM 實際可理解為 3 個主要階段：
- 特征提取：檢索分類特征的池化 Embedding 表征，代表作包含 Pointwise Attention 和 Target-aware Attention 池化。
- 特征交互：DLRM 中最關鍵的部分。HSTU 通過Norm (A(X)V(X)) ⊙ U (X) 實現 Attention 池化特征與其他特征的“交互”來取代 DLRM 的特征交互。
- 表征轉換：DLRM 中通常通過 MoE（Mixture of Expert）來實現，以處理多樣化、異質的群體。其關鍵思想是：提供不同用戶的特化子網絡以實現條件計算。HSTU 中的 Pointwise 實際上可以執行 MoE 中的門控操作。?
5.1 Pointwise 聚合 Attention
HSTU 采用新的 Pointwise 聚合（歸一化）Attention 機制（Softmax Attention 在整個序列上計算歸一化因子）。這一設計基于兩個因素：
- 與 Target 相關的先驗數據點的數量作為強特征，能夠指示用戶偏好的強度，而Softmax 歸一化后很難捕獲這一特征。這一點至關重要，因為需要同時預測用戶參與的強度（如，在某一個 Item 上花費的時間）以及 Item 的相對順序（如，預測最大 AUC 的排序）。
- 盡管 Softmax 激活在結構上對噪聲具有魯棒性，但在流式場景中處理非恒定詞表時，其適應性較差。
如上述公式（2）和（3）所示，在 Pointwise 池化后需要進行 LayerNorm 以穩定訓練過程。如下圖 Table 2 所示，作者合成了非恒定詞表的流式數據，在此場景下，Softmax Attention 與 Pointwise Attention 的差距達到 44.7%。
5.2 增加稀疏性
在推薦系統中，用戶歷史序列的長度通常呈 Skewed 分布，導致輸入序列非常稀疏，特別是在長序列場景尤為明顯，這種稀疏性可被利用以顯著提升 Encoder 的效率。為此，作者開發了一種高效的 GPU Attention Kernel，以融合連續的 GEMM 操作，實質上是將 Attention 計算轉換為可變大小的 Grouped GEMM 計算。因此，HSTU 中的 SelfAttention 機制變為 Memory Bound，并且在 Memory 訪存方面按照
的比例擴展，其中 ni 表示樣本 i 的序列長度，dqk 為 Attention 維度，R 為寄存器大小。這種方法本身可以帶來 2x-5x 的吞吐提升。
作者進一步通過隨機長度（Stochastic Length, SL）增加用戶歷史序列的稀疏性。推薦系統中用戶歷史序列的一個關鍵特征是：用戶行為在時間維具有重復性。這提供了一個在不損害模型質量的前提下人為增加稀疏性的機會，從而顯著降低按 
比例擴展的成本。可以將第 j 個用戶的歷史表示為序列
，其中 nc,j 是用戶交互的 Content 數量。設
，
表示從原始序列
構建的長度為 L 的子序列。SL 選擇輸入序列的方式如下所示：
該方法將 Attention 相關復雜度降低至 
，需要說明的是，由于訓練成本往往遠大于推理成本，因此將稀疏 Attention 機制應用于訓練成本收益更顯著。
如下圖 Table 3 展示了不同序列長度和 α 值下，針對用戶 30 天歷史的稀疏度。
- 2.0 表示未應用稀疏注意力機制的基礎稀疏情況。
- 藍色下劃線表示損失可忽略不計的情況。
- 較低的 α 值適用于更長的序列。?
5.3 最大限度地減少激活內存使用
在推薦系統中，大 Batch 對于訓練吞吐量和模型質量至關重要。因此，激活內存的占用成為主要瓶頸，而在 LLM 場景，通常 Batch 不是特別大，參數占的內存更大。
相較于 Transformer 結構，HSTU 采用了更加簡化及完全融合的設計，可以顯著降低激活內存的使用。
- 首先，HSTU 將 Attention 機制外的線性層數量從 6 個減少到 2 個。
- 其次，HSTU 采用了 Kernel 融合，將多個 OP 融合為一個，將每層的激活內存使用量減少到 BF16 格式的 2d+2d+4hdqk +4hdv +2hdv = 14d。
作為對比，Transformer 在 Attention 后還有 FFN + Dropout（中間激活 3hdv），隨后還有 LayerNorm 等，總的激活為 2d + 4dff +2d+1d = 5d+4dff，通常 dff=4d，hdv >= d，考慮其他部分，總的激活內存大概為 33d。因此，HSTU 實際可以支持超過 2x 的深度網絡。
此外，由于詞表也會占據大量內存。以 10B 詞表，512 的 Embedding 維度及 Adam 優化器為例，在 FP32 精度下存儲 Embedding 向量及優化器狀態需要 10B*512*3*4 Byte=60TB 內存。其中 3 表示原始 Embedding 及一階矩和二階矩。為此，作者采用 Rowwise AdamW 優化器，并將優化器狀態放在 DRAM 中，從而將 HBM 顯存開銷從 12 byte 降低到 2 byte。
5.4 通過成本分攤實現推理規模擴展
作者提到的最后一個挑戰是推薦系統在 Serving 階段需要處理大量候選 Item。
- 在檢索方面，Encoder 成本完全可以分攤，并且存在量化、哈希或束搜索、層次檢索等手段。
- 在排序方面，則要面臨可能多達數萬個候選 Item。
為此，作者提出了一種名為 M-FALCON（Micro Batched-Fast Attention Leveraging Cacheable OperationNs）的算法，用于處理輸入序列長度為 n 的 m 個候選 Item 的推理任務。在單次 Forward 中，M-FALCON 通過調整 Attention Mask 和相對位置偏置，同時處理 bm 個候選 Item。使得針對這 bm 個候選 Item 執行的 Attention 操作完全相同。這一改進將 Cross Attention 的計算成本從 O(bmn2d) 降低到 O((n+bm)2d)=O(n2d) ，前提是 bm 是遠小于 n 的常數。可以選擇將全部 m 個候選 Item 劃分為 [m/bm] 個大小為 bm 的 Micro Batch，以利用 Encoder 級別的 KV Cache。總體而言，M-FALCON 使得模型復雜度能夠與 DLRM 模型排序階段的候選 Item 數量線性擴展。
六、實驗
6.1 驗證 HSTU Encoder 的歸納假設
6.1.1 傳統序列場景
如下圖 Table 4 所示，作者在廣泛采用的 2 個推薦系統數據集（MovieLens 和 Amazon Reviews）上評估了 HSTU 性能，并將 SASRec 作為基線（值得注意的是，此處采用的評估方法與工業級應用場景存在顯著差異）。“SASRec (2023)” 表示論文中報告的最佳 SASRec 配置；“HSTU” 表示使用了與 SASRec 相同的配置（如層數、Head 數等）。“HSTU-large”則表示更大規模的 HSTU Encoder（層數增加 4 倍，Head 數增加 2 倍）。結果表明：
- 在相同配置下，HSTU 憑借專為推薦進行的優化設計，顯著超越了基線；
- 當規模擴大時，HSTU 的性能進一步提升。?
6.1.2 工業級流式場景
如下圖 Table 5 所示，作者進一步在工業級流式場景下對比了 HSTU 及其變種相比標準 Transformer 模型的優勢。總體而言，在小規模實驗環境下，HSTU 以 1.5x - 2x 的更快速度實現了更高的質量，并減少了 50% 的 HBM 顯存占用。
6.2 Encoder 效率
隨機長度：如下圖 Figure 4 與 Figure 5a 展示了 SL 對模型指標的影響。當 α=1.6 時，長度為 4096 的序列大部分時間可以縮減至 776 的長度，即移除了超過 80% 的 Token。即便稀疏率提升至 64%-84%，主要任務獲得的 NE 下降不超過 0.002（0.2%）。表明適當的 α 值，SL 不會對模型質量產生負面影響，并能實現高稀疏度以降低訓練成本。
Encoder 效率：如下圖 Figure 5 對比了 HSTU 與 Transformer Encoder 在訓練和推理環境下的效率。對于 Transformer，作者采用了 FlashAttention2。序列長度范圍從 1,024 到 8,192，并在訓練過程中應用 SL。在評估中，對 HSTU 和 Transformer 采用相同配置（d=512，h=8，dqk=64），并消融相對注意力偏置。作者在 H100 GPU 上以 BF16 比較了 Encoder 的表現。總體而言，HSTU 在訓練和推理中的效率分別比 Transformer 高出 15.2x 和 5.6x。
6.3 GR 與 DLRM 在工業級流式場景中的對比
最后，作者在工業級流式場景中，將 GR 的端到端性能與最先進的 DLRM 基線進行了對比。如下圖 Table 6 和 Table 7 所示，GR 不僅在離線測試中顯著優于 DLRM，在 A/B 測試中也帶來了 12.4% 的提升。
如下圖 Figure 6 所示，作者進一步對比了 GR 和生產環境 DLRM 的效率，盡管 GR 模型計算復雜度高出 285x，但由于 HSTU 和新型 M-FALCN 算法，在 1024/16384 個候選 Item 排序時，分別獲得了 1.50x/2.99x 的 QPS。
6.4 推薦系統的縮放法則
眾所周知，在大規模工業場景中，DLRM 在特定的計算和參數規模下，其性能會趨于飽和。為了深入理解這一現象，作者對比了 GR 與 DLRM 的可擴展性。
鑒于特征交叉層對 DLRM 性能至關重要，作者采用了多種方法來擴展 DLRM 基線模型在 Ranking 場景的表現。具體實驗包括 Transformer 架構、DHEN、以及生產環境中應用的 DCN 變體。對于檢索基線模型，由于已采用殘差結構，作者通過擴大隱藏層大小、Embedding 維度及層數來提升性能。
對于 HSTU ，則通過調整 HSTU 的超參數來擴展模型，這些超參數包括殘差層數量、序列長度、Embedding 維度以及 Attention Head 數目。此外，作者還對檢索任務中的數目進行了調整。
最終，所有主要指標，包括用于檢索的 Hit Rate@100 和 Hit Rate 500，以及用于排序的 NE，在給定適當超參數的情況下，均經驗性地遵循所用計算量的冪律關系。如下圖 Figure 1 所示，作者在 3 個數量級的范圍內觀察到這一現象，直到能夠測試的最大規模（序列長度 8192，Embedding 維度 1024，層數 24），此時所使用的總計算量接近 GPT-3 和 LLaMA-2 所使用的總計算量，作者首次證明了 LLM 的縮放法則同樣可能適用于大規模推薦系統。
七、參考鏈接
