谷歌大腦重磅研究：首個具有O(nlogn)時間、O(n)空間復雜度可微分排序算法，速度快出一個數量級
魚羊 十三 發自 凹非寺
量子位 報道 | 公眾號 QbitAI
快排堆排冒泡排。排序，在計算機中是再常見不過的算法。
在機器學習中，排序也經常用于統計數據、信息檢索等領域。
那么問題來了，排序算法在函數角度上是分段線性的，也就是說，在幾個分段的“節點”處是不可微的。這樣，就給反向傳播造成了困難。
現在，谷歌大腦針對這一問題，提出了一種快速可微分排序算法，并且，時間復雜度達到了O(nlogn)，空間復雜度達為O(n)。
速度比現有方法快出一個數量級！
代碼的PyTorch、TensorFlow和JAX版本即將開源。
快速可微分排序算法
現代深度學習架構通常是通過組合參數化功能塊來構建，并使用梯度反向傳播進行端到端的訓練。
這也就激發了像LeCun提出的可微分編程?(differentiable programming)的概念。
雖然在經驗上取得了較大的成功，但是許多操作仍舊存在不可微分的問題，這就限制了可以計算梯度的體系結構集。
諸如此類的操作就包括排序?(sorting)和排名?(ranking)。
從函數角度來看都是分段線性函數，排序的問題在于，它的向量包含許多不可微分的“節點”，而排名的秩要比排序還要麻煩。
首先將排序和排名操作轉換為在排列多面體(permutahedron)上的線性過程，如下圖所示。
△排列多面體說明
在這一過程后，可以發現對于r(θ)，若是θ出現微小“擾動”，就會導致線性程序跳轉到另外一個排序，使得r(θ)不連續。
也就意味著導數要么為null，要么就是“未定義”，這就阻礙了梯度反向傳播。
為了解決上述的問題，就需要對排序和排名運算符，進行有效可計算的近似設計。
谷歌大腦團隊提出的方法，就是通過在線性規劃公式中引入強凸正則化來實現這一目標。
這就讓它們轉換成高效可計算的投影算子(projection operator)，可微分，且服從于形式分析(formal analysis)。
在投影到排列多面體之后，可以根據這些投影來定義軟排序(soft sorting)和軟排名(soft ranking)操作符。
△軟排序和軟排名操作符
在此基礎上，要想完成快速計算和微分，一個關鍵步驟就是將投影簡化為保序優化?(isotonic optimization)。
接下來是將保序優化進行微分，此處采用的是雅可比矩陣(Jacobian)，因為它簡單的塊級結構，使得導數很容易分析。
而后，結合命題3和引理2，可以描述投影到排列多面體上的雅可比矩陣。
需要強調的是，與保序優化的雅可比矩陣不同，投影的雅可比矩陣不是塊對角的，因為我們需要對它的行和列進行轉置。
最終，可以用O(n)時間和空間中的軟算子雅可比矩陣相乘。
實驗結果
研究人員在CIFAR-10和CIFAR-100數據集上進行了實驗。
實驗使用的CNN，包含4個具有2個最大池化層的Conv2D，RelU激活，2個完全連接層；ADAM優化器的步長恒定為10-4，k=1。
與之比較的是O(Tn2)的OT方法，以及O(n2)的All-pairs方法。
△rQ及rE為新算法
結果表明，在CIFAR-10和CIFAR-100上，新算法都達到了與OT方法相當的精度，并且速度明顯更快。
在CIFAR-100上訓練600個epoch，OT耗費的時間為29小時，rQ為21小時，rE為23小時，All-pairs為16小時。在CIFAR-10上結果差不多。
在驗證輸入尺寸對運行時間的影響時，研究人員使用的是64GB RAM的6核Intel Xeon W-2135，以及GeForce GTX 1080Ti。
禁用反向傳播的情況下，進行1個batch的計算，OT和All-pairs分別在n=2000和n=3000的時候出現內存不足。
啟用反向傳播時，OT和All-pairs分別在n=1000和n=2500的時候出現內存不足。
開啟新的可能性
曾就職于谷歌、NASA的機器學習工程師Brad Neuberg認為，從機器學習的角度來說，快速可微分排序、排名算法看上去十分重要。
而谷歌的這一新排序算法，也在reddit和hacker news等平臺上引起了熱烈的討論。
有網友對其帶來的“新可能性”做出了更為詳細的討論：
我想，可微分排序生成的梯度信息量更大，使得梯度下降的速度更快，從而能夠進一步提升訓練速度。
我認為，這意味著某些基于排名的指標，以后可以用可微分的形式來表示。也就是說，神經網絡可以輕松地針對這些結果直接進行優化。
對于谷歌而言，這很顯然會應用于網絡搜索，以及諸如標簽分配之類的東西問題。
也有網友指出，雖然該算法并不是第一個解決了排序不可微問題的方法，但它的效率無疑更高。
傳送門
論文：https://arxiv.org/pdf/2002.08871.pdf
討論：https://news.ycombinator.com/item?id=22393790https://www.reddit.com/r/MachineLearning/comments/f85yp4/r_fast_differentiable_sorting_and_ranking/
— 完 —
- 蘋果芯片主管也要跑路！庫克被曝出現健康問題2025-12-07
- 世界模型和具身大腦最新突破：90%生成數據，VLA性能暴漲300%｜開源2025-12-02
- 谷歌新架構突破Transformer超長上下文瓶頸！Hinton靈魂拷問：后悔Open嗎？2025-12-05
- 90后華人副教授突破30年數學猜想！結論與生成式AI直接相關2025-11-26
相關閱讀
o3并非獨門秘技，谷歌已發背后關鍵機制，方法更簡單、成本更低
DeepSeek-Coder通過重復采集5個樣本，性能優于GPT-4o，而成本卻僅為后者的三分之一。
