本文介紹CVPR2023的中稿論文：Temporal Attention Unit: Towards Efficient Spatiotemporal Predictive Learning。這篇論文介紹了一種用于高效時空預測的時間注意力單元（Temporal Attention Unit，TAU）。該方法改進了現有框架，對時間和空間上的依賴關系分別學習，提出了時間維度上的可并行化時序注意力單元，實現了高效的視頻預測。

代碼開源在我們的時空預測學習框架OpenSTL中。OpenSTL是一個全面的時空預測學習基準，涵蓋了廣泛的方法和不同的任務，從合成的移動物體軌跡到現實世界的場景，如人類運動、駕駛場景、交通流和天氣預報。歡迎大家關注！

引言

時空預測學習是一種通過學習歷史幀來預測未來幀的自監督學習范式，可以利用海量的無標注視頻數據學習豐富的視覺信息，在氣象預測、交通流量預測、人體姿勢變化估計等領域有著廣泛的應用場景。時空預測學習需要考慮視頻中的空間相關性和時間演變規律，這是一項具有挑戰性的任務。傳統的方法主要基于循環神經網絡來建模時間依賴關系，但是RNN有著計算效率低、難以捕捉長期依賴、容易出現梯度消失或爆炸等缺點。因此，如何設計一個高效、準確、穩定的時空預測學習模型，是一個亟待解決的問題。為了解決這個問題，我們首先研究現有的方法,并提出時空預測學習的通用框架，如下圖所示。

TAU

如下圖所示，TAU模型不使用循環神經網絡，而是使用注意力機制來并行化地處理時間演變。TAU模型將時空注意力分解為兩個部分：幀內靜態注意力和幀間動態注意力。幀內靜態注意力使用小核心深度卷積和擴張卷積來實現大感受野，從而捕捉幀內的長距離依賴關系。幀間動態注意力使用通道間注意力的方式來學習不同幀之間的通道權重，從而捕捉幀間的變化趨勢。

TAU模塊將時間注意力分為兩部分：幀內靜態注意力和幀間動態注意力。前者通過獲得的大感受野捕捉幀內的長期依賴關系；而后者以擠壓和激發的方式學習通道的注意力權重，以捕捉時間線上的時序演變。最后的注意力是動態注意力和靜態注意力結合的產物。受ViTs和大核卷積的啟發，研究者使用了深度卷積（DW Conv）、深度擴張卷積（DW-D Conv）和1x1通道卷積來建模大核卷積。

此外，我們還提出了一種新穎的差分散度正則化方法，用于優化時空預測學習的損失函數。該方法同時考慮了幀內誤差和幀間變化量。通過將預測幀和真實幀之間的差分轉換為概率分布，并計算它們之間的KL散度，來強制模型學習到視頻中固有的變化規律。差分散度正則化（differential divergence regularization）是預測幀與其對應的真實幀之間的Kullback-Leibler（KL）散度。具體而言，它是預測幀差分和真實幀差分之間的KL散度。

τ 代表溫度參數，經驗性地將其設置為0.1以增強概率分布的差異。直觀來說，均方誤差損失（MSE）僅考慮幀內誤差，而差分散度正則化克服了這一缺點，迫使模型學習連續幀之間的差異并意識到固有的變化，以改善模型的預測。

因此目標損失函數：

實驗

Moving MNIST

下圖是在Moving MNIST上測試的兩個實例，對于隨機運動的數字，預測與目標的絕對差異很細微，說明TAU能很好地處理時空預測：

相對于SOTA的循環模型，TAU的性能增益是較大的，在三個度量指標下，TAU的表現都超越了其他方法：

TaxiBJ

在真實復雜環境的交通流量數據集上，TAU具有良好的表現：

不同數據的泛化

為了檢驗模型的泛化能力，我們先在KITTI原始數據上進行訓練，接著使用Caltech dataset進行評估，評估時輸入前十幀預測下一幀。

靈活長度的預測

我們的模型可以通過模仿RNN，將預測的幀作為輸入并遞歸產生預測來處理靈活長度的預測。對于KTH數據集，人體運動預測任務的難點不僅在于預測幀的靈活長度，還在于涉及人類意識隨機性的復雜動力學，這增加了任務的困難程度。TAU可以從給定的10幀中預測接下來的20或40幀，也有出色的表現。

運行效率

此外，我們的模型不僅可以提高視頻生成質量，還可以提高計算效率和訓練速度。如下圖所示，收斂速度極快，50輪訓練即可達到MSE 35.0的水準。在相同實驗環境下，TAU模型在基準數據集上每個周期只需要2.5分鐘，而此前的SOTA方法需要7到30分鐘不等。

總結

本文提出了一個通用的時空預測學習框架，使用基于靜態和動態結合的時間注意力模塊替代循環單元，還引入了差分散度正則化方法來解決僅考慮幀內誤差的MSE損失的問題，為高效的時空預測學習提供了新的范式。

編輯：黃飛

?