何愷明新作出爐！異構預訓練Transformer顛覆本體視覺學習範式，AI性能暴漲超20%

新智元報道

編輯：桃子 喬楊

【新智元導讀】通用機器人模型，如何解決異構性難題？來自MIT、Meta FAIR團隊全新提出異構預訓練Transformer（HPT），不用從頭訓練，即可破解。

通用機器人模型，目前最大的障礙便是「異構性」。

也就是說，必須收集全方位——每個機器人、任務和環境的特定數據，而且學習後的策略還不能泛化到這些特定設置之外。

由此，AI大神何愷明帶隊的MIT、Meta FAIR團隊，提出了異構預訓練Transformer（HPT）模型。

即預訓練一個大型、可共享的神經網絡主幹，就能學習與任務和機器人形態無關的共享表示。

簡單講，就是在你的策略模型中間放置一個可擴展的Transformer，不用從頭開始訓練！

論文地址：https://arxiv.org/pdf/2409.20537

研究人員將不同本體視覺輸入對齊到統一的token序列，再處理這些token以控制不同任務的機器人。

最後發現，HPT優於多個基準模型，並在模擬器基準和真實世界環境中，將未見任務微調策略性能，提升20%。

值得一提的是，這項研究被NeurIPS 2024接收爲Spotlight。

在真實環境中，HPT加持下的機器人本體，能夠自主向柴犬投食。

而且， 即便是灑了一地狗糧，機器人也能用抹布，將其收到一起。

而在模擬環境中，HPT架構讓機器人任務操作，更加精準。

接下來，一起深度瞭解下異構預訓練Transformer（HPT）模型的核心要素吧。

搭建「異構性」橋樑

如今，構建特定的機器人策略很困難，其中最大的難題就是數據收集和缺少泛化性。

不同硬件的機器人在物理上具有不同的本體（embodiment），每種實例可以有不同的「本體感覺」（proprioception），包括不同的自由度、末端執行器、運動控制器和爲特定應用構建的工作空間配置。

此外，另一種常見的異構性就是視覺異構性。

不同機器人搭載了不同的視覺傳感器，而且通常配備在不同位置（比如手腕/第三視角）；每個機器人的外觀也會因環境和任務而有很大差異。

正是由於這些難以跨越的異構性障礙，因此通常需要收集每個機器人、任務和環境的特定數據，並且學習到的策略不能泛化到這些特定設置之外。

雖然機器人領域已經積累了海量的開源數據，但異構性讓數據集很難被共同利用。

從圖4中就可以看出，僅僅是按環境分類，機器人領域的數據就能被「瓜分」爲遠程遙控、模擬、野外、人類視頻等接近4等份。

機器人領域數據集的異質性

近些年來NLP和CV領域的突飛猛進，讓我們看到了徹底改變機器學習領域的一個歷史教訓：對大規模、高質量和多樣化數據進行預訓練，可以帶來通常優於特定模型的通用模型。

話至此處，當今機器人領域的一箇中心問題浮出水面：如何利用異構數據來預訓練機器人基礎模型？

除了更多數據帶來的好處之外，不同任務的訓練還可以增強表示（representation）的通用性。

這類基礎模型將會在各種任務上實現高成功率、對異常值更加穩健，並且能夠靈活地適應新任務。

那麼，到底應該如何充分利用異構化的數據集？

如圖1所示，一個基本的思路是，將來自不同領域和任務的輸入信號映射到高維表示空間，並讓它們表現出一致的縮放行爲。

之後，只需要最少的微調，就可以將得到的高維表示遷移到特定的下游任務，同時獲得良好的性能。

HPT概念示意圖

HPT所要做的，就是找到一種共享的策略「語言」，能夠對齊來自不同預訓練的異質的本體感覺和視覺信息，將自己的信號映射到共享的潛在空間。

HPT模型架構

HPT全稱爲Heterogeneous Pre-trained Transformers，是一個架構系列，採用了模塊化的設計思路，從異構本體的數據中進行可擴展學習。

受到多模態數據學習的啓發，HPT使用了特定於本體的分詞器（stem）來對齊各種傳感器輸入，映射爲固定數量的token，之後送入Transformer結構的共享主幹（trunk），將token映射爲共享表示並進行預訓練。

在對每種本體的輸入進行標記化（tokenize）之後，HPT就運行在一個包含潛在token短序列的共享空間上運行。

論文提到，這種層次結構的動機，也是來源於人類身體的脊髓神經迴路層面中，特定運動反應和感知刺激之間的反饋循環。

預訓練完成後，使用特定於任務的動作解碼器（head）來產生下游動作輸出，但所用的實例和任務在預訓練期間都是未知的。

預訓練包含了超過50個單獨的數據源，模型參數超過1B，模型的代碼和權重都已公開發布。

HPT架構

stem結構

從上面的描述來看，要解決異構性問題，最直接和最關鍵的就是如何訓練stem，將來自異構的本體和模態的傳感器輸入對齊到共享表示空間中。

如圖3所示，stem包含兩個主要部分，即本體感受分詞器和視覺分詞器，將來自不同本體的異構輸入映射爲固定維度、固定數量的token，讓trunk能夠以相同的方式處理。

其中的關鍵思想，是利用cross-attention機制，讓固定數量的可學習token關注到各種特徵。

雖然這篇論文主要處理本體感覺和視覺，但處理觸覺、3D和動作輸入等其他類型的異構傳感器信號也可以在stem中靈活擴展。

HPT中的stem架構

按照時間順序單獨處理每個模態後，將所有token拼接在一起並添加額外的模態嵌入和正弦位置嵌入，就得到了trunk的輸入序列。

爲了避免過擬合，stem被設計爲僅有少量參數，只包含一個MLP和一個注意力層。

trunk結構

作爲預訓練的核心組件，trunk是一個有潛在d維空間的Transormer結構，參數量固定，在不同的本體和任務之間共享，以捕獲複雜的輸入-輸出關係。

預訓練

給定從不同分佈中採樣的異構本體的數據集_1,…,_k,…,_K ，令_k={τ^(i)}_{1≤i≤M_k} 表示_k中一組軌跡M_k，τ^(i)={o_t^(i), a_t^(i)}_{1≤t≤T}表示第i個最大長度爲T的軌跡，每個元組包含observation變量和action變量。

訓練目標如公式（1）所示，需要最小化數據集中的以下損失：

其中ℒ是行爲克隆損失，計算爲預測結果和真實標籤之間的Huber 損失。

該訓練過程有兩個數據縮放軸：單個數據集D_k的體量M_k，以及數據集總數K。

在預訓練階段，每次迭代時僅更新trunk部分參數，並且基於訓練批次採樣更新特定於每個異構本體和任務的stem和head部分。

論文進行了一系列預訓練實驗，包括不同規模的網絡參數和數據集大小，旨在回答一個問題：HPT預訓練在跨域異構數據中是否展現出了擴展能力？

總體而言，某種程度上，HPT隨着數據集數量、數據多樣性、模型體量和訓練計算量呈現出縮放行爲。

HPT網絡詳細信息，寬度表述turnk transformer的潛在維度，深度表示block數量，默認設置爲HPT-Small型號

預訓練數據集詳細信息，默認使用來自RT-X的27個數據集的16k個軌跡進行訓練

數據縮放

數據方面，如圖5所示，即使在異構程度逐漸增大的本體中也具有穩定且可擴展的驗證損失。

此外，作者還發現，計算量（相當於每次訓練運行看到的樣本量）和數據量需要共同擴展，才能在訓練過程中更接近收斂。

epoch縮放

如圖6所示，增加批大小（左）相當於有效地擴展訓練token數（右），通常可以提高模型性能，直至最後收斂。

另一個觀察結果是，使用分佈式方法，在每個訓練批中聚合儘可能更多的數據集，用更大的批大小來彌補異構訓練中的較大方差。

模型縮放

如圖7所示，固定數據集和軌跡數量，沿着模型大小（從1M到1B）進行縮放，並逐漸將批大小從256增加到 2048（模型大小每增加一倍），並使用具有170k軌跡的更大數據集。

可以觀察到，當我們擴展到具有更大計算量（紅線）的更大模型時，預訓練可以實現較低的驗證損失，直到達到穩定水平，但沒有發現縮放模型深度和模型寬度之間存在顯著差異。

圖8中的實驗結果表明，HPT可以相當有效地處理異構數據。儘管與真實機器人存在很大的差距，但對其他本體的數據集（例如模擬環境和人類視頻數據集）進行預訓練是可能的。

遷移學習

如上，作者使用了最後一次迭代中驗證集上的損失來評估預訓練。

接下來，他們將通過實驗，去驗證機器人在遷移學習中，任務成功率的問題：

預訓練的HPT模型，是否可以遷移到模擬和現實世界中的全新本體、任務、以及環境中？

模擬環境

如下圖10（a）中，研究人員在閉環模擬中測試了下游任務的模型，並觀察到使用HPT-B到HPTXL預訓練模型，提到的任務成功率。

在圖10（b）中，他們在最近發佈的Simpler基準上運行HPT，它允許在高保真模擬上與Octo、RT1-X、RT2-X進行比較。

在Google EDR機器人中，研究人員重點關注三個不同的任務「關閉抽屜」、「選可樂罐」。

對於每個任務，他們測試了幾種不同的初始化，所有任務總共有300+ episode。

現實世界

這裡，作者採用了與前一節類似的遷移學習方法，並在真實世界的評估協議下，評估預訓練的HPT表示。

他們以256批大小和訓練率訓練策略20000次迭代。

圖12顯示的定量結果，研究人員觀察到，預訓練策略相比No-Trunk和From-Scratch基準獲得了更好的成功率。

特別是在倒水的任務中，From-Scratch基準使用了最先進的擴散策略架構，以展示預訓練表示的靈活性。

圖11定性結果顯示，作者觀察到預訓練的HPT在面對不同姿勢、物體數量、相機配置、光照條件時，表現出更好的泛化能力和魯棒性。

在表3中，作者對Sweep Leftover任務進行了消融研究。

儘管最近數據規模激增，但由於異構性的存在，機器人學習的通用性仍然受到限制。

研究人員提出的HPT——一種模塊化架構和框架，通過預訓練來應對這種異構性。

他希望這一觀點能夠啓發未來的工作，以處理機器人數據的異構性本質，從而爲機器人基礎模型鋪平道路。

作者介紹

Lirui Wang

Lirui Wang是MIT CSAIL的博士生，導師是Russ Tedrake教授。

在此之前，他曾在華盛頓大學獲得學士和碩士學位，導師是Dieter Fox教授。

他的研究興趣在於機器學習和機器人學。尤其是，他對開發能夠在複雜和非結構化的真實世界環境中，泛化的算法和系統感興趣。

爲了實現這一點，他一直致力於研究能夠隨着異類數據進行擴展的「艦隊學習」（fleet learning）。

Xinlei Chen

Xinlei Chen是舊金山Meta Fair實驗室的研究科學家。目前的研究興趣是預訓練，特別是自監督、多模態視覺表徵的預訓練。

他曾在CMU語言技術研究所獲得博士學位，就讀期間也在機器人研究所工作。此前，他獲得了浙大的學士學位。

Jialiang Zhao

Jialiang Zhao目前是 MIT CSAIL感知科學小組的博士生，導師是Edward H. Adelson教授，並與Russ Tedrake 、何愷明合作。

Kaiming He

何愷明目前是麻省理工學院電子工程與計算機科學系副教授。

他提出的最爲著名的研究是深度殘差網絡（ResNets），並被廣泛應用到現代深度學習模型當中，比如Transformer（GPT、ChatGPT）、AlphaGo Zero、AlphaFold、擴散模型等。

在加入MIT之前，何愷明於2016年至2024年擔任Facebook AI Research的研究科學家，並於2011年-2016年擔任微軟亞洲研究院（MSRA）的研究員。

他曾在2011年在香港中文大學獲得博士學位，並於2007年在清華大學獲得學士學位。

參考資料：

https://liruiw.github.io/hpt/

https://x.com/LiruiWang1/status/1841098699436351742