Google Research 範例#
Google Research 透過 Flax 所做的研究精選。
注意#
快速注意力 (FAVOR+) 和使用 Performers 重新思考注意力#
GitHub 上的程式碼
研究論文
使用 Performers 重新思考注意力（Choromanski 等，2020 年）
介紹“Performer，這是一種變壓器架構，能夠估計具有可證明準確性的正則（softmax）全等級注意力變壓器，但只使用線性（相對於二次）空間和時間複雜度，而不依賴於任何先驗，例如稀疏性或低等級。為了估計 softmax 注意力核心，Performer 使用一種創新的透過正交隨機特徵的快速注意力方法 (FAVOR+)，這種方法對於可調整的核方法可能有獨立的優點。FAVOR+ 也可以有效率地為 softmax 之外的可核化注意力機制建模。”
自我注意力不需要 O(n^2) 記憶體#
研究論文
自我注意力不需要 O(n^2) 記憶體（Rabe 和 Staats，2021 年）
“我們提供了一個非常簡單的演算法以實現注意力，這個演算法關於序列長度需要 O(1) 記憶體，而其在自我注意力上的延伸則需要 O(log n) 記憶體。這與常常被聲稱的自我注意力需要 O(n^2) 記憶體的說法相違。雖然時間複雜度還是 O(n^2)，但裝置記憶體而不是運算能力通常是現代加速器的限制因素。因此，減少注意力的記憶體需求允許處理比其他情況可行時更長的序列……”
電腦視覺#
著色變壓器 (ColTran)#
研究論文
著色變換器（Kumar 等人，2020 年）
「我們介紹著色變換器 (ColTran)，一種完全依賴自注意力來進行影像著色的架構。我們引入了條件變換器層，這是一個基於自注意力的條件式產生模型的新穎建構模組。我們的消融研究顯示，採用此機制優於多種不同的基線。最後，我們證明 ColTran 可以產生多樣化、高保真度的 ImageNet 著色，即使對於人類評分者而言，這些著色也與地面真相幾乎無法區分。」
視覺變換器 (ViT)、MLP 混合器架構及大型視覺#
GitHub 上的程式碼
-
「此程式碼庫設計用於使用雲端 TPU VM 或 GPU 機器訓練大型視覺模型。它基於 Jax/Flax 函式庫，並使用 tf.data 和 TensorFlow 資料集來建立可擴充且可重複的輸入管線。」
-
視覺變換器和 MLP 混合器的 JAX 程式碼
超過 5 萬個視覺變換器和 hybrid 檢查點，用於產生「如何訓練你的 ViT？」的資料
研究論文
一張圖片值 16x16 個字詞：大規模影像辨識變換器（Dosovitskiy 等人，2020 年）
「在視覺中，注意力時與卷積神經網路一起應用，時用於取代卷積神經網路的特定元件，同時保持其整體結構。我們證明了對 CNN 的依賴並非必要，且直接應用於影像貼片的序列的純變換器在影像分類任務中能表現得非常好。在大量資料上進行預訓練並轉移到多個中型或小型影像辨識基準（ImageNet、CIFAR-100、VTAB 等），與最先進的卷積神經網路相比，視覺變換器 (ViT) 取得了極佳的結果，同時需要的訓練運算資源大幅減少。」
MLP-Mixer：一種適用於視覺的純 MLP 架構（Tolstikhin 等人，2021 年）
「在本文中，我們證明了儘管卷積和注意力對於良好的效能都已足夠，但兩者皆非必要條件。我們介紹了 MLP-Mixer，一種完全基於多層感知器 (MLP) 的架構。MLP-Mixer 包含兩種類型的層：一種是將 MLP 獨立應用於影像貼片（即「混合」按位置的特色），一種是將 MLP 應用於各個貼片（即「混合」空間資訊）。在大型資料集上進行訓練，或使用現代化正則化方案時，MLP-Mixer 能在影像分類基準上達到具競爭力的分數，其預訓練和推論成本與最先進的模型相當。」
如何訓練你的 ViT？視覺變換器中的資料、擴充和正則化（Steiner 等人，2021 年）
「視覺變壓器 (ViT) 已證實可針對範圍廣泛的視覺應用（例如影像分類、物件偵測和語意影像分割）達成極具競爭力的效能。與卷積神經網路相較，視覺變壓器的較弱歸納偏差通常會增加對模型正規化或資料擴充（簡稱「AugReg」）的依賴性，尤其在較小的訓練資料集上訓練時。我們執行系統性的經驗研究，以進一步了解訓練資料、AugReg、模型大小和運算預算之間的相互作用。
當視覺變壓器不需預訓練或強資料擴充就能超越 ResNet (X. Chen 等人，2021 年)
「視覺變壓器 (ViT) 和 MLP 意味著進一步努力用通用神經架構取代手工連接的特徵或歸納偏差。既有作品透過大量資料（例如大規模預訓練和/或重複的強資料擴充）來強化模型，但仍會回報優化相關的問題（例如對初始化和學習率的敏感度）。因此，本文從損失幾何的角度探討 ViT 和 MLP 混合器，旨在改善訓練模型資料效率和推理廣義化的能力。
LiT：鎖定影像文字調整後的零次轉移 (X. Zhai 等人，2021 年)
「本文介紹對比調整，這是一種採用對比訓練，在影像和文字模型對齊時，仍然善用其預訓練的簡單方法。在我們的經驗研究中，我們發現鎖定預先訓練的影像模型，搭配解鎖的文字模型，效果最好。我們將此對比調整的實例稱作「鎖定影像調整」(LiT)，只教導文字模型從預先訓練的影像模型中讀取優秀的表徵，以進行新的任務。LiT 模型取得零次轉移到新視覺任務的能力，例如影像分類或檢索。所提出的 LiT 廣泛適用，將三種不同的影像文字資料集與多種預先訓練方法（監督式和無監督式）搭配使用，以及運用不同的架構（ResNet、視覺變壓器和 MLP 混合器），都展現穩定的表現。
運用稀疏專家混合 (MoE) 縮放視覺#
研究論文
運用稀疏專家混合縮放視覺 (Riquelme 等人，2021 年)
「稀疏閘控專家混合網路 (MoE) 已展示在自然語言處理中的絕佳可擴充性。然而，在電腦視覺上，幾乎所有效能良好的網路都是「密集」的，也就是每個輸入都由每個參數處理。我們提出視覺 MoE (V-MoE)，這是視覺變壓器的稀疏版本，可與最大的密集網路競爭…我們展示 V-MoE 在縮放視覺模型的潛力，並訓練了一個 15B 參數的模型，在 ImageNet 上達成 90.35% 的表現…
擴散#
變異擴散模型#
研究論文
變分擴散模型（Kingma 等人，2021 年）
「基於擴散的生成模型展示了在感官上令人印象深刻的合成能力，但它們是否也能成為出色的基於機率的模型？我們肯定地回答這個問題，並介紹了一系列基於擴散的生成模型，可在標準圖像密度估計基準上獲得最先進的機率。與其他基於擴散的模型不同，我們的模型允許與其他模型一起有效地最佳化雜訊時程。我們展示了變分下界（VLB）簡化為一個非常簡短的表達式來表示擴散資料的訊雜比，進而改善了我們對此類模型的理論理解。利用此見解，我們證明了文獻中提出的幾個模型之間存在的等價性。此外，我們展示了連續時間 VLB 對於雜訊時程的不變性，僅端點的訊雜比除外。這使我們能夠學習最小化所得 VLB 估計機方差的雜訊時程，進而加快最佳化…」
領域適應#
GIFT（特徵漸進插值朝向目標）#
研究論文
野外的漸進領域適應：沒有中間分佈時（Abnar 等人，2021 年）
「我們關注領域適應問題，目標是將模型轉移到目標分佈，而不是學習不變區域表示。已顯示在以下兩個假設下：a）可以取得來自中間分佈的樣本，以及 b）樣本已標示來自來源分佈變化的數量，自訓練可以成功應用於漸進偏移的樣本，以將模型調整到目標分佈。我們假設有 (a) 足以透過使用隱式課程，讓反覆自訓練慢慢地將模型調整到目標分佈。在 (a) 不成立的情況下，我們觀察到反覆自訓練無法完成。我們提出 GIFT，一種透過插值來源和目標領域範例的表示來建立虛擬樣本作為中間分佈的方法…」
概括 #
代理差距最小化改善清晰度感知訓練#
研究論文
代理差距最小化改善清晰度感知訓練（J. Zhuang 等人，2022 年）
「最近提出的針對銳利度的最小化 (SAM) 透過最小化定義為參數空間中鄰域內最大損失的擾動損失來提升泛化能力。然而，我們顯示銳利與平坦的最小值都可以有較低的擾動損失，表示 SAM 並非總是偏好平坦的最小值。反之，我們定義了代理差距，一種等於在鄰域半徑（用於衍生擾動損失）較小時局部最小值中 Hessian 的顯著特徵值的測量。代理差距在訓練期間易於計算，且適用於直接最小化。以上述觀察為基礎，我們提出代理差距引導的銳利度感知最小化 (GSAM)，一種對 SAM 的創新優化，只有極小的計算開銷…」
元學習#
learned_optimization#
GitHub 上的程式碼：learned_optimization
研究論文
使用持續進化策略在展開計算圖中進行無偏梯度估計（Vicol 等人，2021 年）
「我們引進一種稱為持續進化策略 (PES) 的方法，其將計算圖分割為一系列的截斷展開，並在每次展開後執行基於進化策略的更新步驟。PES 透過在整個展開序列中累積校正項，消除了這些截斷的偏差。PES 能夠進行快速參數更新，具有較低記憶體使用率，無偏且方差特性合理。」
梯度並非你所需要的一切（Metz 等人，2021 年）
「…在這份簡報中，我們討論在各種可微分環境中出現的常見的基於混亂的失敗模式，包括從遞迴神經網路和數值物理模擬到訓練已學習的最佳化器。我們追溯這種失敗到所研究系統的雅可比矩陣的頻譜，並提供實務工作者可能會期望此失敗損壞他們的基於微分的最佳化演算法的準則。」
模型效能#
神經渲染 / NeRF#
可概化的 Patch-Based 神經渲染#
研究論文
可概化的 Patch-Based 神經渲染（Suhail 等人，2022 年）
「…我們提出一個不同的典範，在這個典範中，不需要深入特徵和類似 NeRF 的體積渲染。我們的演算法只要取得從場景採樣的貼片集合，就能預測全新場景中目標光線的顏色。」
最佳化#
Amos 最佳化器與 JEstimator#
GitHub 上的程式碼
-
「…實作相容於 optax 函式庫的最佳化器 Amos，和一個輕量級函式庫 JEstimator，它配備類似於 tf.Estimator 的介面，用於管理 JAX 中機器學習程式 T5X 相容的檢查點，我們使用這個介面，來執行本文的實驗。」
-
研究論文
Amos：一種採用自適應權重衰減方式，朝向模型導向縮放的 ADAM 式最佳化器（Tian 和 Parikh，2022 年）
提出「一種相容於 optax 函式庫的最佳化器 Amos，和一個輕量級函式庫 JEstimator，它配備類似於 tf.Estimator 的介面，用於管理 JAX 中機器學習程式 T5X 相容的檢查點」。「當用於 BERT 變體和 T5 的預訓練時，Amos 持續比最先進的 AdamW 設定收斂得更快，在 ≤70% 的訓練步驟和時間內獲得更佳驗證損失，同時插槽變數只需 ≤51% 的記憶體。」
量化#
Pareto 最佳量化 ResNet 大都是 4 位元與 AQT：有精準量化訓練#
GitHub 上的程式碼
研究論文
Pareto 最佳量化 ResNet 主要為 4 位元（Abdolrashidi 等，2021）
“此研究利用 ResNet 作為案例分析量化對推論運算成本品質折衷曲線的影響。我們的結果顯示，對於每個 bfloat16 ResNet 模型，都有更低成本且精確度更高的量化模型；換句話說，bfloat16 運算成本品質折衷曲線由 4 位元和 8 位元的曲線帕雷托主宰，而主要量化為 4 位元的模型能產生最佳帕雷托曲線... 主要量化為 4 位元的模型能產生最佳帕雷托曲線... 我們使用的量化方法針對實用性而最佳化：所需調整不多，且設計時已考量硬體能力... 作為本研究的一部分，我們提供一個以 JAX 撰寫的量化函式庫...”
強化學習#
利用示範進行動作量化之連續控制 (AQuaDem)#
研究論文
利用示範進行動作量化之連續控制（Dadashi 等，2021）
提出“一種針對具備連續動作空間的問題而設計的強化學習 (RL) 架構：利用示範進行動作量化 (AQuaDem)。提議方法在於由人類示範學習連續動作空間離散化。此離散化會為每個輸入狀態回傳一組合理動作（根據示範），進而掌握示範者的先驗知識和其多模態行為。透過離散化動作空間，任何離散動作深度學習技術都能輕鬆應用於連續控制問題中。實驗顯示，提議方法在 RL 設定中優於最先進的方法，例如在 RL 設定中的 SAC，以及在模仿學習設定中的 GAIL。”
序列模型 / 模型平行化#
T5X：使用t5x和seqio放大模型和資料#
-
“T5X 是一個模組化、可組合、研究友善的框架，可用於高效能、可組態、自服務訓練、評量和推論序列模型（從語言開始），支援各種規模。”
研究論文
T5X：使用 t5x 和 seqio 放大模型和資料（Roberts 等，2022）
「近期的基於神經網路的語言模型受益於訓練資料集規模和模型本身參數數目的擴大。由於多種因素的影響，擴大可能很複雜，包括將計算分散在超級電腦叢集（例如 TPU）、在輸入數據時防止出現瓶頸，以及確保可複制的結果。在這項研究中，我們介紹了兩個簡化這些問題的軟體函式庫：t5x 簡化了大規模構建和訓練語言模型的流程，同時保持易用性，而 seqio 則提供了一個任務型的 API，用於簡單建立快速且可重複的訓練資料和評估管道。這些開源函式庫已用於訓練擁有數百億個參數的模型，這些模型的訓練資料集有數 TB。除了函式庫，我們還釋出 T5 型編碼器-解碼器模型以及 GPT 型專用解碼器架構的組態和說明。
模擬#
Brax - 大型剛體模擬的可微分物理引擎#
研究論文
Brax - 大型剛體模擬的可微分物理引擎 (Freeman 等人，2021 年)
「我們介紹 Brax，一個用於剛體模擬的開源函式庫，重點放在加速器上的效能和並行處理，並使用 JAX 編寫。我們關於現有強化學習文獻啟發的一系列任務呈現結果，但在我們的引擎中重新製作。此外，我們在 JAX 中提供 PPO、SAC、ES 和直接策略最佳化的重新實作，這些會和我們的環境一起編譯，讓學習演算法和環境處理可以在相同的裝置上執行，並在加速器上無縫擴展。