加速 TPUs 上的影片擴散模型時空注意力機制

影片擴散模型（Video diffusion models）生成影片通常很慢，原因是要跑大量去噪步驟（denoising steps），而且每個步驟的計算成本都很高。序列長度一增加，自注意力機制（self-attention）就可能成為單一去噪步驟延遲的主要來源之一。以目前流行的開源影片生成模型來說，光是 81 幀 720p 影片，序列長度就落在 50,000 到 400,000 之間。如果把解析度從 720p（HD）提高到 1440p（2K），序列長度可能增加到四倍。全注意力機制（full attention）的計算量隨序列長度呈二次方成長，它在每層延遲中的佔比可能從 55.5% 升到 88.2%。既然注意力機制是單一步驟內延遲的最大來源，任何激進的推論優化都必須先從注意力機制下手。

在影片擴散模型中，注意力機制的結構性很強：許多查詢-鍵（query-key）互動帶有的注意力品質（attention mass）很少，所以大量成對計算（pairwise computation）通常可以跳過。稠密注意力（Dense attention）不管重要性高低，會計算所有互動；稀疏注意力（Sparse attention）則用遮罩（mask）只保留最重要的互動，其餘捨棄。稀疏影片生成（Sparse VideoGen, SVG）是這個方向的奠基性工作，它找出並利用了一種注意力模式：不同的注意力頭（heads）通常可分為空間頭（spatial heads）或時間頭（temporal heads）。在空間頭中，一個影像區塊（patch）主要關注相同或鄰近幀中的其他區塊。在時間頭中，一個影像區塊則關注橫跨多個幀的小範圍空間區域。

根據這個結構，SVG 在推論時動態分析注意力頭，並把每個頭分派給空間遮罩或時間遮罩。做法是先採樣少量查詢，計算它們在稠密、空間和時間注意力下的輸出，再選出與稠密基準輸出偏差最小的稀疏遮罩。這讓 SVG 在給定的稀疏度下能維持較高的品質。如圖 3 所示，兩種遮罩都保留對第一幀（F00）的全注意力，作為錨定全局場景外觀的注意力匯流排（attention sink），同時也包含用於幀間互動的局部帶。

為了進一步優化，第二個迭代版本（B2：完整/邊界影像區塊專業化）替完整影像區塊提供完全不加遮罩的快速路徑，只在邊界影像區塊上執行精確座標遮罩。這兩條路徑共同產生相同的注意力輸出，其部分結果會用對應的 softmax 歸一化統計數據合併。過程中，保留的查詢-鍵對維持不變。這項優化讓延遲從 96.37 毫秒降到 54.12 毫秒，比樸素的稀疏遍歷（B1）減少 44%，也比稠密 Splash 快 31%。

較小的影像區塊能更貼近遮罩邊界，因此需要在區塊內加遮罩的影像區塊比例會降低。影像區塊大小也會改變核心如何分組計算和資料搬移。不過，把邊界工作降到最低，並不一定能把延遲降到最低。圖 6 比較了大影像區塊與小影像區塊的邊界，其中深色線條描出 SVG 遮罩邊界（包括左側垂直的第一幀匯流排）。較小的影像區塊更貼近邊界，黃色的邊界影像區塊變少，但代價是計算影像區塊的效率變差。

實驗中，查詢影像區塊大小（BQ）會調整，鍵/值影像區塊大小（BKV）則固定為 1536。所有配置都使用獨立的完整路徑和邊界路徑。BQ = 1024 時，只有 16.95% 的執行影像區塊需要內部遮罩，但延遲是 66.88 毫秒。BQ 增加到 3328 時，這個比例升到 27.55%，延遲則降到 54.12 毫秒。

激進的調度方案（aggressive schedule）把去噪時間從 153.50 秒縮短到 119.86 秒，加速 1.28 倍，延遲降低 22%。這些結果來自多項做法的結合，包括分別調整的稠密與稀疏配置，以及它們的分散式佈局。由於去噪還包含稀疏注意力以外的工作，這些增益比單獨核心的加速小。隨著序列長度增加，自注意力在 Transformer 區塊延遲中的佔比越來越大。稀疏注意力的端到端加速比會隨解析度明顯擴大。評估不同解析度下的激進調度方案，可以看出這些節省在較長序列上會疊加放大。

表 4 呈現在八個 TPU v6e 晶片上，Aggressive SVG 稀疏調度方案下，端到端去噪延遲和品質如何隨解析度變化。在 1080p（171K tokens）下，激進調度方案把去噪時間從 683 秒降到 457 秒。在 1440p（302K tokens，為 720p 序列長度的四倍）下，去噪時間從 2,471 秒降到 1,461 秒。圖 10 並排比較 720p 81 幀影片在稠密注意力（左）與 Aggressive SVG 稀疏調度方案（右，1.28 倍加速，24.80 dB PSNR）下的生成結果。