NeMo：面向深度神經(jīng)網(wǎng)絡(luò)分解的神經(jīng)元級訓(xùn)練中模塊化方法

NeMo: A Neuron-Level Modularizing-While-Training Approach for Decomposing DNN Models

https://arxiv.org/pdf/2508.11348

摘要

隨著深度神經(jīng)網(wǎng)絡(luò)（DNN）模型日益融入現(xiàn)代軟件系統(tǒng)，DNN 模型高昂的構(gòu)建成本已成為軟件開發(fā)中的一個(gè)重大挑戰(zhàn)。為了應(yīng)對這一挑戰(zhàn)，模型復(fù)用已被廣泛應(yīng)用以減少模型訓(xùn)練成本；然而，不加區(qū)分地復(fù)用整個(gè)模型可能會產(chǎn)生顯著的推理開銷。因此，DNN 模塊化——借鑒軟件工程中的模塊化思想——日益受到關(guān)注，通過將 DNN 模型分解為模塊來實(shí)現(xiàn)模塊復(fù)用。特別是，新興的訓(xùn)練中模塊化（MwT）范式，通過將模塊化納入模型的訓(xùn)練過程，優(yōu)于訓(xùn)練后模塊化，已被證明是 DNN 模塊化的一種更有效方法。然而，現(xiàn)有的 MwT 方法專注于卷積核級別的小規(guī)模卷積神經(jīng)網(wǎng)絡(luò)（CNN）模型。它們難以處理多樣的 DNN 和大規(guī)模模型，特別是基于 Transformer 的模型，后者在各種任務(wù)中持續(xù)取得最先進(jìn)的結(jié)果。

為解決這些局限性，我們提出了 NeMo，一種可擴(kuò)展且更具通用性的 MwT 方法。NeMo 在神經(jīng)元級別操作——這是所有 DNN 共有的基本組件——從而確保適用于 Transformer 和各種 DNN 架構(gòu)。此外，我們設(shè)計(jì)了一種基于對比學(xué)習(xí)的模塊化訓(xùn)練方法，配備有效的復(fù)合損失函數(shù)，因此可擴(kuò)展至大規(guī)模模型。在兩個(gè)廣泛使用的分類數(shù)據(jù)集上對兩個(gè)基于 Transformer 的模型和四個(gè) CNN 模型進(jìn)行的綜合實(shí)驗(yàn)，證明了 NeMo 優(yōu)于最先進(jìn)的 MwT 方法。結(jié)果顯示，模塊分類準(zhǔn)確率平均性能提升 1.72%，模塊大小減少 58.10%。我們的發(fā)現(xiàn)表明，NeMo 在 CNN 和大規(guī)?；?Transformer 的模型中均表現(xiàn)出有效性。此外，基于開源項(xiàng)目的案例研究展示了 NeMo 在實(shí)際場景中的潛在益處，為實(shí)現(xiàn)可擴(kuò)展且通用化的 DNN 模塊化提供了一種有前景的方法。

CCS 概念：? 軟件及其工程； 其他關(guān)鍵詞和短語：模型復(fù)用，transformer，視覺 transformer，卷積神經(jīng)網(wǎng)絡(luò)，模塊化

1 引言

深度神經(jīng)網(wǎng)絡(luò)（DNN）模型日益融入當(dāng)代軟件系統(tǒng) [36, 39]，使得 DNN 訓(xùn)練成為軟件開發(fā)生命周期的一個(gè)關(guān)鍵組成部分。然而，DNN 的訓(xùn)練過程，特別是那些具有數(shù)十億參數(shù)和大數(shù)據(jù)集的模型，可能昂貴得令人望而卻步。為了減輕開發(fā)和訓(xùn)練成本，模型復(fù)用 [8, 32, 42, 52, 66] 已在實(shí)踐中被廣泛采用，引起了 AI 和軟件工程社區(qū)的極大關(guān)注。研究人員探索了諸如遷移學(xué)習(xí) [9, 10, 32, 66] 和模型檢索 [26] 等技術(shù)以促進(jìn)有效的模型復(fù)用，同時(shí)開發(fā)了工程最佳實(shí)踐——例如模型重構(gòu) [21] 和優(yōu)化模型供應(yīng)鏈 [60]——以提高效率和安全性。然而，這些方法通常將 DNN 模型視為單體工件，不加區(qū)分地復(fù)用整個(gè)模型可能會產(chǎn)生額外的推理開銷并引入安全漏洞。

DNN 模型和軟件程序共享根本性的相似之處 [30, 36]，前者通過數(shù)據(jù)驅(qū)動的自動訓(xùn)練過程編碼功能，后者通過顯式編碼實(shí)現(xiàn)功能。在軟件工程中，程序開發(fā)通常遵循模塊化原則，促進(jìn)程序組件的按需復(fù)用 [33, 34, 53]。借鑒軟件工程中的模塊化思想，DNN 模型模塊化和按需模型復(fù)用正受到越來越多的關(guān)注。例如，Pan 等人和 Qi 等人開創(chuàng)了全連接神經(jīng)網(wǎng)絡(luò)（FCNNs）[30, 38]、卷積神經(jīng)網(wǎng)絡(luò)（CNNs）[3, 31, 37, 38, 40, 41] 和循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）[20] 模型的 DNN 模塊化，通過識別模型每個(gè)功能的相關(guān)權(quán)重或特殊結(jié)構(gòu)組件。每個(gè)模塊擁有模型功能的一部分，僅保留負(fù)責(zé)其特定功能的權(quán)重。例如，一個(gè) 10 類圖像分類模型可以分解為 10 個(gè)模塊，每個(gè)模塊僅保留負(fù)責(zé)識別單一類的權(quán)重。除了模型復(fù)用之外，DNN 模塊化通過識別相關(guān)模塊，為模型維護(hù) [43]、部署 [8] 和供應(yīng)鏈管理 [60] 提供了潛在益處。

根據(jù)模塊化執(zhí)行的時(shí)間，現(xiàn)有的 DNN 模塊化方法可以分為訓(xùn)練后模塊化 [20, 37, 38, 40, 43, 64, 70] 和訓(xùn)練中模塊化 [41]，前者分解一個(gè)訓(xùn)練好的 DNN 模型，后者將模塊化集成到隨機(jī)初始化模型的訓(xùn)練過程中，然后分解模塊化訓(xùn)練好的模型。在我們之前的工作 [41] 中，我們提出了一種模塊化訓(xùn)練方法 MwT，它將內(nèi)聚性和耦合性這兩個(gè)因素集成到訓(xùn)練損失中。“內(nèi)聚性”評估用于同類樣本的卷積核集合之間的重疊，而“耦合性”評估用于不同類樣本的卷積核集合之間的重疊。通過優(yōu)化這些因素，MwT 最小化了負(fù)責(zé)不同功能的權(quán)重集之間的重疊。由于模塊化模型是專門為模塊化訓(xùn)練的，訓(xùn)練中模塊化在模塊大小、推理性能和模塊化時(shí)間成本方面顯著優(yōu)于訓(xùn)練后模塊化。然而，將模塊化納入訓(xùn)練過程是非常具有挑戰(zhàn)性的，因?yàn)樗枰紤]各種架構(gòu)和規(guī)模模型的訓(xùn)練細(xì)節(jié)，并避免對模型性能和訓(xùn)練時(shí)間成本產(chǎn)生副作用。我們的實(shí)驗(yàn)觀察表明，即使是最先進(jìn)的訓(xùn)練中模塊化方法 MwT [41]，由于模型架構(gòu)和模型大小的通用性問題，也離實(shí)際應(yīng)用還很遠(yuǎn)。

當(dāng)前的 DNN 模塊化技術(shù)，包括 MwT，主要局限于小規(guī)模模型，例如最多只有 1470 萬參數(shù)（用于模塊化模型的可訓(xùn)練參數(shù)）的 FCNN 和 CNN 模型。隨著基于 Transformer 的模型繼續(xù)在圖像 [23, 27, 56]、自然語言 [29, 59] 和代碼相關(guān)任務(wù) [11, 44] 上取得最先進(jìn)結(jié)果，它們在許多應(yīng)用中已基本取代了傳統(tǒng)的 CNN 和 RNN 架構(gòu)。基于 Transformer 模型的優(yōu)越性不僅源于先進(jìn)的架構(gòu) [59]，也得益于大規(guī)模參數(shù) [29]。另一方面，如此大規(guī)模的參數(shù)導(dǎo)致訓(xùn)練和推理成本激增，因此一種針對大規(guī)模基于 Transformer 模型的 DNN 模塊化方法更為必要。盡管 MwT 提出了一個(gè)適用于多樣 DNN 模型的通用框架，但其具體實(shí)現(xiàn)局限于 CNN 模塊化，不支持基于 Transformer 的模型。主要原因可能包括：(1) MwT 在卷積核級別執(zhí)行模塊化，這是 CNN 特有的結(jié)構(gòu)組件，以及 (2) MwT 面臨大規(guī)模模型的性能可擴(kuò)展性挑戰(zhàn)。它利用多個(gè)損失（即內(nèi)聚性和耦合性損失）的直接求和進(jìn)行評估和優(yōu)化。這種方式已通過理論分析 [7, 22, 45] 和實(shí)驗(yàn)（見表 3）被證明對于優(yōu)化是低效的。稿件已提交至 ACM

解決上述局限性提出了重大挑戰(zhàn)。首先，選擇合適的模塊化粒度是困難的。雖然 MwT 可以通過將卷積核級別改為 Transformer 的特殊結(jié)構(gòu)組件（如注意力頭級別）來適應(yīng) Transformer，但注意力頭數(shù)量通常很少這一事實(shí)使得將注意力頭劃分為眾多模塊并實(shí)現(xiàn)高模塊性能和低模塊間重疊變得不可行。例如，在 CIFAR10 數(shù)據(jù)集上微調(diào)的 Vision Transformer (ViT) [23] 模型只有 96 個(gè)注意力頭。假設(shè)我們將其分解為十個(gè)模塊，每個(gè)對應(yīng)一個(gè)類，并強(qiáng)制模塊間無重疊，每個(gè)模塊將包含大約 10 個(gè)注意力頭。鑒于 ViT 模型中有 12 個(gè)注意力層，這意味著那兩層不包含任何注意力頭，這是不可接受的。相比之下，即使小 CNN 模型也擁有大量卷積核，例如用于 10 類分類的 VGG16 有 4226 個(gè)核，使得卷積核級別的模塊化可行。雖然在權(quán)重級別執(zhí)行模塊化的現(xiàn)有工作 [30, 31, 38] 提供了一個(gè)替代思路，但這種模塊化粒度會加劇第二個(gè)挑戰(zhàn)。其次，設(shè)計(jì)一個(gè)包含多個(gè)目標(biāo)的有效損失函數(shù)是困難的。訓(xùn)練損失應(yīng)涉及三個(gè)優(yōu)化目標(biāo)，包括推理性能（例如交叉熵?fù)p失）、內(nèi)聚性和耦合性。直接求和這三部分是一種直接的方式；然而，這種方式已通過理論分析 [7, 22, 45] 被證明對于優(yōu)化是無效的。我們的實(shí)證發(fā)現(xiàn)表明，隨著模型權(quán)重規(guī)模和模塊化相關(guān)可訓(xùn)練參數(shù)的增加，這一挑戰(zhàn)被進(jìn)一步放大。大規(guī)模模型的有效模塊化需要更復(fù)雜的訓(xùn)練損失函數(shù)。

為了應(yīng)對上述挑戰(zhàn)，我們提出了 NeMo，一種開創(chuàng)性的神經(jīng)元級訓(xùn)練中模塊化（NEuron-level MOdularizing-while-training）方法。不同于在權(quán)重 [38] 或特殊結(jié)構(gòu)組件級別 [37, 40, 41] 實(shí)現(xiàn)模塊化的現(xiàn)有方法，NeMo 在神經(jīng)元級別操作。這種粒度有效地解決了第一個(gè)挑戰(zhàn)，并提供優(yōu)于 MwT 的三個(gè)優(yōu)勢：(1) 神經(jīng)元是 DNN 模型的基礎(chǔ)組件，使 NeMo 能夠無縫應(yīng)用于 Transformer 架構(gòu)并容易擴(kuò)展到其他 DNN。(2) 模型中的神經(jīng)元數(shù)量介于權(quán)重?cái)?shù)量和特殊結(jié)構(gòu)組件數(shù)量之間，促進(jìn)模塊化和推理性能之間的平衡。(3) 類似于卷積核，無關(guān)神經(jīng)元可以從模塊中物理移除。此外，由于神經(jīng)元包含的權(quán)重少于卷積核，NeMo 是比 MwT 更細(xì)粒度的方法，因此在模塊化方面表現(xiàn)更好。為解決第二個(gè)挑戰(zhàn)，我們通過引入對比學(xué)習(xí)技術(shù) [1, 6, 16] 來優(yōu)化內(nèi)聚性和耦合性損失函數(shù)。改進(jìn)的損失函數(shù)以指數(shù)和分?jǐn)?shù)形式結(jié)合推理性能、內(nèi)聚性和耦合性，而不是簡單地將它們相加，從而實(shí)現(xiàn)更好的模塊化性能和模塊推理性能。此外，改進(jìn)的損失函數(shù)只有一個(gè)超參數(shù)，而 MwT 有三個(gè)超參數(shù)，使 NeMo 更容易適應(yīng)不同模型。

我們進(jìn)行了綜合實(shí)驗(yàn)來評估 NeMo，利用兩個(gè)基于 Transformer 的模型（ViT 和 DeiT）和四個(gè)主流 CNN 架構(gòu)，跨越兩個(gè)廣泛使用的圖像分類數(shù)據(jù)集。為了確保公平比較，我們也調(diào)整了 MwT 使其在神經(jīng)元級別操作以用于基于 Transformer 的模型。實(shí)驗(yàn)結(jié)果證明了 NeMo 在實(shí)現(xiàn)基于 Transformer 和 CNN 模型的模塊化訓(xùn)練和分解方面的有效性。與 MwT 相比，NeMo 不僅提高了模塊化訓(xùn)練的準(zhǔn)確率，還顯著降低了神經(jīng)元保留率（NRR）或核保留率（KRR）。對于 ViT 模型，與 MwT 相比，NeMo 實(shí)現(xiàn)了高達(dá) 64.76% 的 NRR 降低。關(guān)于按需模型復(fù)用，NeMo 在減少模型大小方面取得了顯著改進(jìn)，與 MwT 相比平均權(quán)重減少 57.85%。此外，基于開源倉庫 [13, 46, 68] 的案例研究進(jìn)一步展示了 NeMo 在現(xiàn)實(shí)場景中的實(shí)際益處。在此案例研究中，NeMo 允許開發(fā)者僅復(fù)用預(yù)訓(xùn)練模型中的相關(guān)模塊，從而減少微調(diào)模型在其下游任務(wù)上的推理開銷。

這項(xiàng)工作的主要貢獻(xiàn)如下：

• 據(jù)我們所知，NeMo 是第一個(gè)支持基于 Transformer 模型的模塊化訓(xùn)練和結(jié)構(gòu)化分解的神經(jīng)元級訓(xùn)練中模塊化方法。
• 我們提出了一種基于對比學(xué)習(xí)的方法來優(yōu)化內(nèi)聚性和耦合性，這使得大規(guī)模模型的模塊化訓(xùn)練成為可能。
• 我們使用兩個(gè)代表性的基于 Transformer 模型和四個(gè) CNN 模型在兩個(gè)廣泛使用的數(shù)據(jù)集上進(jìn)行了廣泛實(shí)驗(yàn)。結(jié)果表明，NeMo 在模塊分類準(zhǔn)確率和模塊大小方面可以優(yōu)于最先進(jìn)的方法。此外，我們進(jìn)行了案例研究以展示 NeMo 在現(xiàn)實(shí)場景中的實(shí)際益處。我們發(fā)布了 NeMo 的復(fù)現(xiàn)包 [2]。

2 背景

本節(jié)簡要介紹本研究的一些預(yù)備信息，包括對比學(xué)習(xí)（CL）和一些主流神經(jīng)網(wǎng)絡(luò)模型。

2.1 主流神經(jīng)網(wǎng)絡(luò)模型

神經(jīng)網(wǎng)絡(luò) [14, 19, 24] 是由分層組織的互連節(jié)點(diǎn)（神經(jīng)元）組成的計(jì)算模型，旨在從數(shù)據(jù)中學(xué)習(xí)復(fù)雜模式。其基本結(jié)構(gòu)包括輸入層、一個(gè)或多個(gè)隱藏層和輸出層。在學(xué)習(xí)過程中，與神經(jīng)元間連接相關(guān)的權(quán)重會被調(diào)整。卷積神經(jīng)網(wǎng)絡(luò)（CNN）[25, 50, 51] 代表了一類專門優(yōu)化的神經(jīng)網(wǎng)絡(luò)，用于處理網(wǎng)格狀數(shù)據(jù)，特別是圖像。CNN 引入了卷積層，將濾波器應(yīng)用于輸入數(shù)據(jù)，使網(wǎng)絡(luò)能夠捕捉局部模式和空間層次結(jié)構(gòu)。CNN 的關(guān)鍵組件包括：

卷積層，其使用可學(xué)習(xí)的濾波器提取特征；池化層，其減少空間維度和計(jì)算復(fù)雜度；以及全連接層，其基于提取的特征執(zhí)行高層推理。

Vision Transformer (ViT) [23] 模型將最初為自然語言處理設(shè)計(jì)的 Transformer 架構(gòu)改編用于計(jì)算機(jī)視覺任務(wù)。ViT 將圖像分割為固定大小的補(bǔ)丁，線性嵌入這些補(bǔ)丁，并使用標(biāo)準(zhǔn) Transformer 編碼器處理它們。ViT 的關(guān)鍵組件包括：補(bǔ)丁嵌入，它將圖像補(bǔ)丁轉(zhuǎn)換為線性嵌入；位置編碼，它添加有關(guān)補(bǔ)丁空間位置的信息；多頭自注意力，它允許模型關(guān)注輸入的不同部分；以及前饋網(wǎng)絡(luò)，它處理注意力輸出。

ViT 已在圖像分類任務(wù)中表現(xiàn)出卓越的性能，通常只需最小的特定領(lǐng)域適配即可超越基于 CNN 的模型?；谶@一成功，Touvron 等人（2021）提出了數(shù)據(jù)高效圖像 Transformer（DeiT）[56]，旨在解決原始 ViT 模型的數(shù)據(jù)效率限制。

現(xiàn)有的結(jié)構(gòu)化模塊化方法 [37, 40, 41] 主要針對 CNN 模型中的卷積核。這些方法通常通過分析其輸出通道來評估每個(gè)核的重要性，隨后在分解階段移除不相關(guān)的核。然而，這種方法論不適用于基于 Transformer 的模型，它們在結(jié)構(gòu)和傳播機(jī)制上存在根本差異（見第 3.2 節(jié)）。

2.2 對比學(xué)習(xí)

對比學(xué)習(xí)及其最近的發(fā)展旨在訓(xùn)練編碼器，以捕捉高維信號不同部分之間的共享信息表示 [57, 63]。對比學(xué)習(xí)的核心思想是將相似樣本（正樣本對）拉在一起，同時(shí)將不相似樣本（負(fù)樣本對）推開。這種方法在自監(jiān)督學(xué)習(xí)設(shè)置中特別有價(jià)值，因?yàn)樗灰蕾囀謩訕?biāo)記的數(shù)據(jù)。

為了使編碼器能夠充分學(xué)習(xí)特征，Wu 等人 [63] 引入了一個(gè)大型記憶庫來存儲數(shù)據(jù)集中所有樣本的特征表示。He 等人 [16] 引入了動量對比（MoCo），它通過維護(hù)一個(gè)負(fù)樣本隊(duì)列并使用動量編碼器來確保特征表示的一致性，從而解決了負(fù)樣本的動態(tài)性質(zhì)問題。Chen 等人 [6] 通過移除對記憶庫或?qū)Ｓ眉軜?gòu)的需求，簡化了對比學(xué)習(xí)框架。相反，它使用大批量大小和數(shù)據(jù)增強(qiáng)策略來即時(shí)生成正樣本對和負(fù)樣本對。Caron 等人 [5] 提出了 SwAV，這是一種將對比學(xué)習(xí)與聚類相結(jié)合的方法。Grill 等人 [15] 引入了 BYOL，這是一種完全消除對負(fù)樣本需求的方法。MwT [41] 中采用的內(nèi)聚性和耦合性指標(biāo)分別量化了同一類別內(nèi)和不同類別間所選神經(jīng)元的相似性。它與各種對比學(xué)習(xí)技術(shù)共享根本性的相似之處，但缺乏針對具有眾多神經(jīng)元模型的有效優(yōu)化損失。因此，我們將對比學(xué)習(xí)納入模塊化訓(xùn)練，并考慮標(biāo)簽以改進(jìn)內(nèi)聚性和耦合性的計(jì)算。

3 方法

本節(jié)詳細(xì)介紹了 NeMo 的方法論。NeMo 旨在進(jìn)行模塊化訓(xùn)練，并將一個(gè) n 類分類模型分解為幾個(gè)模塊，每個(gè)模塊包含神經(jīng)元和功能的子集。

3.1 NeMo 概述

如圖 1 所示，對于一個(gè)隨機(jī)初始化的模型，NeMo 通過兩個(gè)階段實(shí)現(xiàn)按需復(fù)用。(1) 模塊化訓(xùn)練。NeMo 從一個(gè)隨機(jī)初始化的模型開始，并整合了一個(gè)神經(jīng)元標(biāo)識符，它由幾個(gè)掩碼生成器組成。為了實(shí)現(xiàn)模塊化訓(xùn)練，我們引入了一個(gè)組合損失函數(shù)，其中包括準(zhǔn)確率損失和模塊化損失。在訓(xùn)練過程中，NeMo 通過應(yīng)用交叉熵和對比學(xué)習(xí)算法來增加模型的內(nèi)聚性并減少其耦合性。(2) 結(jié)構(gòu)化模塊化。在訓(xùn)練過程之后，NeMo 為每個(gè)功能生成掩碼，然后根據(jù)掩碼從模塊化訓(xùn)練好的模型中結(jié)構(gòu)化地移除不相關(guān)的神經(jīng)元，從而獲得模塊。圖 2 提供了 NeMo 的詳細(xì)工作流程。具體來說，在模塊化訓(xùn)練階段，為了識別相關(guān)神經(jīng)元，一個(gè)神經(jīng)元標(biāo)識符被附加到模型上。它包含幾個(gè)掩碼生成器，每個(gè)對應(yīng)一個(gè)網(wǎng)絡(luò)層。生成的掩碼代表神經(jīng)元對特定子任務(wù)的相關(guān)性。這些掩碼被用于計(jì)算對比損失以進(jìn)行模型優(yōu)化和隨后的剪枝過程。在模型訓(xùn)練完成后，NeMo 利用這些掩碼從每一層物理地移除冗余神經(jīng)元。這個(gè)過程有效地從權(quán)重矩陣中剪枝特定權(quán)重，產(chǎn)生一個(gè)定制的子模型。該子模型的功能與用戶需求精確對齊，確保它僅保留指定子任務(wù)所需的基本組件，同時(shí)消除多余元素。

3.1.2 評估模塊化和分類的性能。 在模塊化訓(xùn)練中，Qi 等人 [41] 引入了內(nèi)聚性（cohesion）和耦合性（coupling）的概念來評估模塊化性能。內(nèi)聚性量化了同類樣本內(nèi)卷積核使用的重疊程度，而耦合性衡量不同類之間的這種重疊。為了優(yōu)化內(nèi)聚性和耦合性，MwT [41] 計(jì)算樣本掩碼之間的相似度，并通過將其與交叉熵?fù)p失直接相加，將相似度整合到損失函數(shù)中。然而，這種損失計(jì)算存在一些缺點(diǎn)。首先，它平等地對待所有樣本，忽略了不同樣本之間難度水平的差異。其次，直接計(jì)算和求和損失可能導(dǎo)致數(shù)值不穩(wěn)定和優(yōu)化困難 [7, 22, 45]。此外，樣本的分布未受控制，使得標(biāo)識符難以學(xué)習(xí)某些特征 [63]。因此，模塊化的評估和優(yōu)化需要一種更復(fù)雜的算法。

對比學(xué)習(xí)（CL）旨在通過區(qū)分相似和不相似的數(shù)據(jù)點(diǎn)對來開發(fā)有用的表示。傳統(tǒng)的 CL 方法 [5, 6, 15, 16, 69]，例如 SimCLR [6] 和 MoCo [16]，通常通過對同一實(shí)例應(yīng)用數(shù)據(jù)增強(qiáng)來創(chuàng)建正樣本對，而負(fù)樣本對則從批次內(nèi)的不同實(shí)例生成。這些方法利用對比損失函數(shù)（例如 InfoNCE 損失 [57]）來最大化正樣本對之間的相似性，并最小化負(fù)樣本對之間的相似性。然后，它們擁有一個(gè)具有訓(xùn)練良好特征的神經(jīng)網(wǎng)絡(luò)編碼器，用于圖像分類或其他下游任務(wù)。

其中 α 表示對比損失的加權(quán)因子。基于設(shè)計(jì)的損失函數(shù)，模型通過梯度下降進(jìn)行訓(xùn)練，使其學(xué)會使用相應(yīng)的神經(jīng)元組（即模塊）來識別不同類別的樣本。

3.2 結(jié)構(gòu)化模塊化

在模塊化訓(xùn)練之后，我們使用神經(jīng)元標(biāo)識符生成的神經(jīng)元掩碼來分解模塊化模型。具體來說，我們首先確定神經(jīng)元在模型中對應(yīng)什么結(jié)構(gòu)，然后根據(jù)神經(jīng)元掩碼對其進(jìn)行分解。

3.2.1 確定神經(jīng)元的結(jié)構(gòu)。 在卷積層中，神經(jīng)元通常代表一個(gè)卷積核及其在輸入特征圖上的感受野，在輸出特征圖通道之一中產(chǎn)生一個(gè)輸出元素。鑒于設(shè)計(jì)的復(fù)雜性，我們從某個(gè)通道形成神經(jīng)元集合用于模塊化訓(xùn)練，即將一個(gè)卷積核作為一個(gè)獨(dú)立的子結(jié)構(gòu)。關(guān)于卷積核的移除，請參考 MwT [41]。

在線性（或全連接）層中，神經(jīng)元聚合來自所有前驅(qū)神經(jīng)元的加權(quán)輸入，添加偏置項(xiàng)，并應(yīng)用激活函數(shù)。這些神經(jīng)元主要進(jìn)行線性變換和特征組合，通過偏置項(xiàng)實(shí)現(xiàn)非線性映射 [67]。與卷積層中的神經(jīng)元不同，線性層中的神經(jīng)元不構(gòu)成特殊的結(jié)構(gòu)（例如，卷積核）。與卷積層神經(jīng)元不同，線性層中的那些神經(jīng)元缺乏特定的結(jié)構(gòu)單元（例如，卷積核）。為了實(shí)現(xiàn)神經(jīng)元移除，我們提出了一種新穎的分解方法，該方法消除了線性變換操作中權(quán)重矩陣的部分內(nèi)容。

在 Transformer 模型中，所有層——包括 MLP 和注意力層——本質(zhì)上都是由線性運(yùn)算組成的。因此，輸出通道的移除統(tǒng)一應(yīng)用于所有此類層。此外，每個(gè)線性層的輸入通道必須進(jìn)行調(diào)整，以匹配其前一層的修改后的輸出通道。對于 MLP 層，其中一層的輸出直接作為下一層的輸入，我們根據(jù)前一層的掩碼移除后續(xù)層的輸入通道。至于注意力層，更新輸入通道需要考慮查詢（Q）、鍵（K）和值（V）矩陣之間的計(jì)算依賴關(guān)系。因此，我們應(yīng)用適當(dāng)?shù)难诖a，基于這些關(guān)系移除不相關(guān)的輸入通道，如圖 3 所示。

MLP 層。 在每一層的線性映射中，神經(jīng)元移除是通過根據(jù)掩碼移除對應(yīng)于輸出維度的權(quán)重來實(shí)現(xiàn)的。然而，對于相鄰層，這個(gè)過程變得復(fù)雜。在移除圖 3 中層 A 的權(quán)重矩陣后，輸出維度已經(jīng)改變，無法與層 B 對齊。雖然零填充層 A 的輸出可以實(shí)現(xiàn)對齊，但這會使層 B 輸入維度的一部分變得冗余。為了優(yōu)化計(jì)算效率，我們改為根據(jù)層 A 的掩碼剪枝層 B 的輸入維度。

殘差連接。 網(wǎng)絡(luò)層之間的殘差連接 [17] 給模塊化過程引入了復(fù)雜性。權(quán)重矩陣的選擇性剪枝可能導(dǎo)致殘差連接的輸入和輸出之間出現(xiàn)維度不匹配，從而阻礙正常的數(shù)據(jù)傳播。例如，在圖 4 的 ViT 模型中，殘差連接鏈接了注意力層和前饋層。默認(rèn)情況下，它們連接兩個(gè)維度相同的層輸出，如圖 4-(a) 所示。然而，在移除神經(jīng)元后，這兩層的輸出維度不匹配；見圖 4-(b)。CNNSplitter [37] 中的解決方案通過添加額外的核來匹配圖 4-(c) 中的兩個(gè)殘差連接層，這引入了額外的內(nèi)存和計(jì)算開銷。MwT [41] 提供了另一種解決方案，即通過填充每個(gè)殘差連接層的輸出以匹配其維度，如圖 4-(d) 所示。然而，填充后層是冗余的，因?yàn)樯婕按颂畛涞挠?jì)算結(jié)果為零。只有前層的填充是有效的，因?yàn)樗枰c后層匹配以進(jìn)行下一次傳播。在圖 4-(e) 中，NeMo 提供了按需填充（On-demand Padding），僅對前層應(yīng)用填充并移除無效計(jì)算，從而避免了額外的開銷。

3.3 按需復(fù)用

NeMo 通過結(jié)構(gòu)化神經(jīng)元移除實(shí)現(xiàn)高效的按需復(fù)用，最大限度地減少了內(nèi)存和計(jì)算開銷。例如，當(dāng) tasked with identifying “cat” and “dog” categories（任務(wù)為識別“貓”和“狗”類別）時(shí)，NeMo 可以從一個(gè) 10 類分類模型中提取相關(guān)模塊，消除無關(guān)神經(jīng)元，從而降低推理成本。

4 實(shí)驗(yàn)

我們通過回答以下研究問題來評估 NeMo 的有效性：

• RQ1：NeMo 在訓(xùn)練和模塊化 DNN 模型方面有多有效？
• RQ2：NeMo 在復(fù)用 DNN 模塊方面有多有效？
• RQ3：與 MwT 相比，NeMo 對于 CNN 模型有多有效？
• RQ4：主要超參數(shù)如何影響 NeMo 的性能？
• RQ5：與 MwT 相比，變化的模型規(guī)模如何影響 NeMo 的模塊化效率？

4.1 實(shí)驗(yàn)設(shè)置

模型。 (i) Vision Transformer (ViT) [23] 和 Data Efficient Image Transformer (DeiT) [56]，包含 12 個(gè)編碼器，且每層大小為 384。(ii) 四個(gè)具有代表性的 CNN 模型，包括 ResNet18 [17]、VGG16 [48]、SimCNN 和 ResCNN，這些模型也被基線方法 MwT [41] 所使用。

數(shù)據(jù)集。 三個(gè)公開的分類數(shù)據(jù)集被用于標(biāo)準(zhǔn)和模塊化訓(xùn)練，包括 CIFAR10、Street View House Number (SVHN) 和 FashionMNIST，這些數(shù)據(jù)集常用于 DNN 模塊化工作 [31, 37, 38, 41]。CIFAR10 數(shù)據(jù)集包含 50,000 張用于訓(xùn)練的自然圖像和 10,000 張用于測試的圖像，分辨率為 32 × 32。每個(gè)樣本包含 10 個(gè)類別中的一個(gè)：airplanes（飛機(jī)）、cars（汽車）、birds（鳥）、cats（貓）、deer（鹿）、dogs（狗）、frogs（青蛙）、horses（馬）、ships（船）和 trucks（卡車）。SVHN 數(shù)據(jù)集包含 604,388 個(gè)用于訓(xùn)練的 0 到 9 的門牌號和 26,032 個(gè)用于測試的門牌號，分辨率為 32 × 32。FashionMNIST 包含 60,000 張用于訓(xùn)練的圖像和 10,000 張用于測試的圖像，分辨率為 28 × 28。每個(gè)樣本包含 10 個(gè)類別中的一個(gè)：t-shirt（T 恤）、trouser（褲子）、pullover（套頭衫）、dress（連衣裙）、coat（外套）、sandal（涼鞋）、shirt（襯衫）、sneaker（運(yùn)動鞋）、bag（包）和 ankle boot（短靴）。

基線方法。 (i) 標(biāo)準(zhǔn)訓(xùn)練。標(biāo)準(zhǔn)訓(xùn)練使用帶有交叉熵?fù)p失的小批量隨機(jī)梯度下降來優(yōu)化 ViT、DeiT 和 CNN 模型。(ii) MwT [41]。MwT 是針對 CNN 模型的最先進(jìn)的模塊化訓(xùn)練方法。為了在 ViT 和 DeiT 模型上比較 NeMo 和 MwT，我們將 MwT 適配到了神經(jīng)元級別，用神經(jīng)元標(biāo)識符替換了核識別器。

評估指標(biāo)。 (i) 準(zhǔn)確率 (ACC)，計(jì)算為在整個(gè)測試集上正確預(yù)測的百分比。(ii) 神經(jīng)元保留率 (NRR) 和卷積核保留率 (KRR)，表示目標(biāo)模塊中保留的平均神經(jīng)元數(shù)量或卷積核數(shù)量除以原始模型中的總數(shù)。(iii) 內(nèi)聚性 (Cohesion)，即所有模塊的平均內(nèi)聚性（公式 (1)）[41]。(iv) 耦合性 (Coupling)，即所有模塊對的平均耦合性（公式 (2)）[41]。

超參數(shù)。 在標(biāo)準(zhǔn)訓(xùn)練中，ViT、DeiT、ResNet18、SimCNN 和 ResCNN 使用 128 的小批量大小訓(xùn)練 200 個(gè)輪次（epochs），而 VGG16 使用 256 的小批量大小進(jìn)行訓(xùn)練。我們將學(xué)習(xí)率設(shè)置為 0.05，Nesterov 動量設(shè)置為 0.9。至于數(shù)據(jù)增強(qiáng) [47]，我們對所有模型應(yīng)用隨機(jī)裁剪和翻轉(zhuǎn)。對比損失權(quán)重 α α 的設(shè)置如表 1 所示，所有模型對比損失中的溫度 τ τ 均設(shè)置為 0.2。

實(shí)驗(yàn)階段。 對于 NeMo，它由兩個(gè)階段組成。第一階段是模塊化訓(xùn)練，NeMo 通過向相應(yīng)的線性層（包括 MLP 和注意力層）添加神經(jīng)元標(biāo)識符，在訓(xùn)練過程中解耦神經(jīng)元。第二階段是分解，NeMo 使用訓(xùn)練集執(zhí)行前向傳播，以獲取表示每一層神經(jīng)元與相應(yīng)類別之間關(guān)系的掩碼。在這些掩碼中，0 表示神經(jīng)元與目標(biāo)類別無關(guān)，而非零值表示相關(guān)。利用這些掩碼，NeMo 將模型分解為所有模塊（例如，在 CIFAR-10 分類模型中，獲得 10 個(gè)模塊），并在其目標(biāo)類別任務(wù)上對每個(gè)模塊進(jìn)行幾個(gè)輪次的微調(diào)。MwT 遵循類似的過程。在 RQ1 中，為了證明模塊化訓(xùn)練不會顯著降低性能，我們設(shè)計(jì)了一個(gè)標(biāo)準(zhǔn)訓(xùn)練基線，其中模型在沒有任何模塊化方法的情況下進(jìn)行訓(xùn)練，作為準(zhǔn)確率基線。

4.2 實(shí)驗(yàn)結(jié)果

4.2.1 RQ1 - 模塊化訓(xùn)練與模塊化的有效性。 為了評估 NeMo 在訓(xùn)練和模塊化模型方面的有效性，我們將 NeMo 應(yīng)用于兩個(gè)基于 Transformer 的模型。此外，我們將 MwT 中的相關(guān)核識別替換為神經(jīng)元標(biāo)識符，以使 MwT 能夠在基于 Transformer 的模型上進(jìn)行比較。我們評估 (1) 模塊化模型的分類準(zhǔn)確率，(2) 基于 Transformer 模型的神經(jīng)元保留率（NRR），以及 (3) 訓(xùn)練模塊的內(nèi)聚性和耦合性。表 2 展示了基于 Transformer 模型上標(biāo)準(zhǔn)訓(xùn)練、模塊化訓(xùn)練和模塊化的結(jié)果。""列顯示了 ViT 和 DeiT 模型的神經(jīng)元數(shù)量。"Standard Model ACC"和"Modular Model ACC"分別顯示了標(biāo)準(zhǔn)訓(xùn)練模型和模塊化訓(xùn)練模型在測試集上的準(zhǔn)確率。對于基于 Transformer 的模型 ViT 和 DeiT，標(biāo)準(zhǔn)模型和模塊化模型分別實(shí)現(xiàn)了 88.27% 和 87.04% 的準(zhǔn)確率，表明模塊化訓(xùn)練僅損失 1.23 個(gè)百分點(diǎn)，說明模塊化訓(xùn)練不會導(dǎo)致分類準(zhǔn)確率的大幅損失。至于神經(jīng)元保留率（NRR），ViT 和 DeiT 模型的平均 NRR 為 11.59%，這意味著在閾值為 0.9 的情況下，模塊化后每個(gè)模塊平均保留 11.59% 的神經(jīng)元?；?Transformer 模型的內(nèi)聚度為 0.9294，表明一個(gè)模塊使用幾乎相同的神經(jīng)元來預(yù)測屬于相應(yīng)類別的樣本。另一方面，基于 Transformer 模型的耦合度為 0.1410，表明不同模塊僅共享少量神經(jīng)元。

此外，如表 3 所示，我們還將 NeMo 與最先進(jìn)的模塊化訓(xùn)練方法 MwT 在 ViT 和 DeiT 模型上進(jìn)行比較。類似于 MwT 中的 CNN 模型在卷積層的輸出通道上應(yīng)用掩碼以表示應(yīng)選擇哪些卷積核并生成模塊，基于 Transformer 的模型可以在線性層的輸出上使用掩碼來表示應(yīng)選擇哪些神經(jīng)元?；谶@些掩碼，我們可以將 MwT 應(yīng)用于基于 Transformer 的模型，并使用掩碼結(jié)構(gòu)化地分解模型。平均而言，MwT 和 NeMo 的模塊化訓(xùn)練模型準(zhǔn)確率分別為 80.03% 和 87.04%。MwT 導(dǎo)致準(zhǔn)確率損失 6.28 個(gè)百分點(diǎn)，而 NeMo 僅導(dǎo)致 1.48 個(gè)百分點(diǎn)的準(zhǔn)確率損失。另一方面，在 NRR 方面，NeMo 也優(yōu)于 MwT（11.59% vs 26.37%），降低了 58.10%。總體而言，與 MwT 相比，使用 NeMo 訓(xùn)練的基于 Transformer 的模型表現(xiàn)出更高的內(nèi)聚性和更低的耦合性，同時(shí)具有更高的準(zhǔn)確率和更低的 NRR。

表 4 展示了 ViT 模型在不同分類子任務(wù)上的神經(jīng)元保留率和權(quán)重保留率。例如，在表 4 的左表中，對于一個(gè) 2 類 CIFAR10 分類子任務(wù)，來自 NeMo 的相應(yīng) ViT-CIFAR10 模塊僅使用模型 18.80% 的神經(jīng)元，而來自 MwT 的相同模塊則擁有模型 44.48% 的神經(jīng)元。平均而言，與 MwT 相比，NeMo 的整體神經(jīng)元保留率降低了 43.68%（43.61% vs 77.34%，簡單計(jì)算所有模塊的平均值），顯著減少了模塊大小和操作開銷。值得注意的是，我們研究中采用的 ViT 和 DeiT 模型由 12 個(gè)編碼器組成。為了確保模塊的準(zhǔn)確率，我們保留了完整的第一個(gè)編碼器，并分解了后續(xù)的 11 個(gè)編碼器。因此，所有 NRR 值表示后 11 個(gè)編碼器中的神經(jīng)元保留率。為了更準(zhǔn)確地評估模塊保留的參數(shù)數(shù)量，我們使用 WRR 來量化所有 12 個(gè)編碼器的權(quán)重保留率。在表 4 的右表中，我們考慮了 ViT 和 DeiT 模型中的所有編碼器，并計(jì)算了權(quán)重保留率。對于這些模型，與神經(jīng)元保留率相比，權(quán)重保留率更直接地反映了參數(shù)的情況。平均而言，來自 NeMo 的模塊僅保留了原始模型 33.14% 的權(quán)重，比來自 MwT 的模塊（擁有模型 65.06% 的權(quán)重）降低了 49.06%。對于每個(gè)子任務(wù)，與 MwT 相比，NeMo 能夠生成更小的模塊。

我們在表 5 中進(jìn)一步比較了 NeMo 與 MwT 在按需模型復(fù)用準(zhǔn)確率方面的表現(xiàn)。表中的"loss"行表示與 MwT 相比，NeMo 損失準(zhǔn)確率的程度。正值表示 NeMo 的準(zhǔn)確率低于 MwT，而負(fù)值表示 NeMo 的準(zhǔn)確率高于 MwT。本質(zhì)上，在所有模型和子任務(wù)中，NeMo 保持了與 MwT 相同的準(zhǔn)確率。考慮到 NeMo 為所有子任務(wù)生成的模塊保留的神經(jīng)元數(shù)量幾乎是 MwT 的一半，NeMo 在準(zhǔn)確率方面的表現(xiàn)值得稱贊。

為了評估 NeMo 在簡化按需模型復(fù)用方面的程度，我們測量了所有模塊在目標(biāo)任務(wù)上的 FLOPs(M)，并在表 6 中將 NeMo 的結(jié)果與 MwT 和標(biāo)準(zhǔn)訓(xùn)練模型進(jìn)行了比較。模型名稱行中的值表示直接復(fù)用時(shí)標(biāo)準(zhǔn)訓(xùn)練模型的 FLOPs。我們的實(shí)驗(yàn)表明，NeMo 顯著簡化了按需模型復(fù)用，將 FLOPs 降低了高達(dá) 82.99%。與 MwT 相比，NeMo 在簡化按需復(fù)用方面也顯示出實(shí)質(zhì)性改進(jìn)，F(xiàn)LOPs 降低了高達(dá) 66.73%。

4.2.3 RQ3 - 與 MwT 在 CNN 上的比較。 為了評估 NeMo 的通用性，我們在四個(gè) CNN 模型上對 NeMo 和 MwT 進(jìn)行了全面的比較分析。此次評估主要有兩個(gè)目的：(1) 證明對比學(xué)習(xí)在優(yōu)化內(nèi)聚性和耦合性損失方面的有效性。(2) 確立 NeMo 的有效性擴(kuò)展到基于 Transformer 的架構(gòu)之外，顯示其在不同神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)上的通用性。為了確保公平比較，我們使用了與原始 MwT 研究相同的四個(gè) CNN 模型。所有 MwT 數(shù)據(jù)直接來源于原始論文或其源代碼。這種方法保證了模型架構(gòu)和數(shù)據(jù)的一致性，使得能夠在各種 CNN 模型上對 NeMo 相對于 MwT 的性能進(jìn)行無偏評估。具體來說，我們評估了 NeMo 在模塊化訓(xùn)練和按需復(fù)用方面的有效性。

模塊化訓(xùn)練。 為了評估 NeMo 在訓(xùn)練 CNN 模型方面的有效性，我們將神經(jīng)元標(biāo)識符替換為相關(guān)核識別 [41]，使用卷積核掩碼計(jì)算內(nèi)聚性和耦合性，以使 NeMo 適用。我們評估 (1) 模塊化模型的分類準(zhǔn)確率，(2) 四個(gè) CNN 模型的核保留率 (KRR)，以及 (3) 訓(xùn)練模塊的內(nèi)聚性和耦合性。表 7 顯示了 CNN 模型上標(biāo)準(zhǔn)訓(xùn)練、模塊化訓(xùn)練和模塊化的結(jié)果。""列指示不同 CNN 模型的卷積核數(shù)量。"Standard Model ACC"和"Modular Model ACC"分別顯示標(biāo)準(zhǔn)訓(xùn)練模型和模塊化訓(xùn)練模型在測試集上的準(zhǔn)確率。標(biāo)準(zhǔn)模型和模塊化模型分別實(shí)現(xiàn)了 93.50% 和 92.59% 的準(zhǔn)確率，模塊化訓(xùn)練的平均準(zhǔn)確率損失為 0.91 個(gè)百分點(diǎn)，表明 NeMo 的模塊化訓(xùn)練過程不會導(dǎo)致分類準(zhǔn)確率的大幅損失。在表 8 中，我們將 NeMo 與 MwT 進(jìn)行比較。MwT 和 NeMo 的模塊化訓(xùn)練模型準(zhǔn)確率分別為 92.41% 和 92.59%，表明 NeMo 實(shí)現(xiàn)了 0.18 個(gè)百分點(diǎn)的提升。此外，在 KRR 方面，NeMo 相比 MwT 取得了不錯(cuò)的性能增益（17.55% vs 12.80%），提升了 27.07%。關(guān)于內(nèi)聚性和耦合性，我們采用 MwT 評估算法并評估了使用 NeMo 訓(xùn)練的模型的內(nèi)聚性和耦合性。NeMo 在內(nèi)聚性上達(dá)到 0.9710，耦合性上達(dá)到 0.1950，而 MwT 分別為 0.9518 和 0.2302。總體而言，與 MwT 相比，使用 NeMo 訓(xùn)練的模型表現(xiàn)出更高的內(nèi)聚性和更低的耦合性，同時(shí)具有更高的準(zhǔn)確率和更低的 NRR/KRR。

表 9 從核保留率（KRR）方面比較了 NeMo 與 MwT 在復(fù)用 CNN 模型時(shí)的表現(xiàn)。平均而言，NeMo 顯著降低了模塊中的 KRR。對于每個(gè)子任務(wù)，與 MwT 相比，NeMo 能夠生成更小的模塊。我們在表 10 中進(jìn)一步比較了按需復(fù)用的準(zhǔn)確率。平均而言，NeMo 保持了與 MwT 相當(dāng)?shù)男阅埽谟行苑矫姹憩F(xiàn)出最小甚至沒有退化。值得注意的是，NeMo 生成的模塊明顯小于 MwT 生成的模塊。鑒于模塊大小的顯著減少和具有競爭力的性能，我們可以得出結(jié)論：與 MwT 相比，NeMo 在 CNN 模型的按需復(fù)用方面表現(xiàn)出更優(yōu)越的有效性。這種在不犧牲性能的情況下效率的提升，凸顯了 NeMo 在推進(jìn)深度學(xué)習(xí)應(yīng)用中模塊化方法的潛力。

表 11 評估了 NeMo 和 MwT 實(shí)現(xiàn)的計(jì)算開銷減少。模型名稱行中的數(shù)字表示直接復(fù)用標(biāo)準(zhǔn)訓(xùn)練模型時(shí)的 FLOPs(M) 數(shù)量。我們的實(shí)驗(yàn)表明，NeMo 顯著簡化了 CNN 模型的按需復(fù)用，將 FLOPs 降低了高達(dá) 80.29%。與 MwT 相比，NeMo 在減少 FLOPs 方面也表現(xiàn)出改進(jìn)，將計(jì)算成本降低了高達(dá) 25.46%。

4.2.4 RQ4 - 超參數(shù)的影響。 最后，我們研究了訓(xùn)練過程中超參數(shù) α α（對比損失的加權(quán)因子）的影響。此外，我們還討論了 τ τ，即對比損失中的溫度參數(shù)。由于模型架構(gòu)的多樣性，我們僅展示了 VGG16 模型在 CIFAR-10 上的訓(xùn)練過程。更多實(shí)驗(yàn)結(jié)果和訓(xùn)練細(xì)節(jié)可在項(xiàng)目網(wǎng)頁 [2] 上找到。

我們觀察到 α α 直接影響模塊化訓(xùn)練過程。圖 6 展示了模塊化訓(xùn)練過程中的準(zhǔn)確率、耦合度、內(nèi)聚度和 KRR，其中 α α 值分別設(shè)置為 1.0、1.2、1.4、1.6 和 1.8，批量大小保持一致為 128。經(jīng)過 200 個(gè)輪次的訓(xùn)練后，可以明顯看出，隨著 α α 的增加，模型準(zhǔn)確率逐漸下降，但 KRR 也顯著降低。在合理范圍內(nèi)，適度增加 α α 不會導(dǎo)致大幅的準(zhǔn)確率損失，但可以顯著降低 KRR。然而，過大的 α α 值會使模型難以擬合，并導(dǎo)致不可接受的準(zhǔn)確率損失。

4.2.5 RQ5 - NeMo 的可擴(kuò)展性。 與之前的工作 MwT 相比，NeMo 提供了神經(jīng)元級別的識別和模塊化方法。然而，對于大多數(shù)模型，神經(jīng)元的數(shù)量通常遠(yuǎn)大于具有可比參數(shù)數(shù)量的 CNN 模型中的卷積核數(shù)量。例如，在我們的實(shí)驗(yàn)中，CNN 模型平均有 4,160 個(gè)卷積核，而 ViT-small 模型包含超過 40,000 個(gè)神經(jīng)元。這個(gè)數(shù)量直接影響我們優(yōu)化目標(biāo)的數(shù)量。由于神經(jīng)元存在于所有神經(jīng)網(wǎng)絡(luò)模型中，NeMo 將在未來工作中擴(kuò)展到更多模型架構(gòu)，這就要求我們的損失函數(shù)在面對更大優(yōu)化目標(biāo)時(shí)保持有效性。我們基于對比學(xué)習(xí)以及指標(biāo)內(nèi)聚性和耦合性重新設(shè)計(jì)了模塊化訓(xùn)練損失函數(shù)。為了評估這個(gè)新?lián)p失函數(shù)的有效性，我們選擇了三個(gè)具有不同神經(jīng)元數(shù)量的 ViT 模型，并通過僅改變損失函數(shù)來評估性能。

如表 12 所示，隨著 ViT 模型中神經(jīng)元數(shù)量的增加，MwT 在平衡神經(jīng)元保留率 (NRR) 和準(zhǔn)確率 (ACC) 損失方面比 NeMo 困難得多。MwT 要么遭受過度的準(zhǔn)確率下降 (-1.39)，要么保留多幾倍的神經(jīng)元 (+25.33) 以保持可接受的準(zhǔn)確率。鑒于 ACC 和 NRR 代表了一個(gè)明確的權(quán)衡，我們的實(shí)驗(yàn)表明，MwT 在具有更多優(yōu)化目標(biāo)的場景中表現(xiàn)不佳，通常導(dǎo)致更差的準(zhǔn)確率。

4.3 案例研究

NeMo 的一個(gè)潛在應(yīng)用是使模型共享平臺能夠提供按需模型復(fù)用 [3]。不同于僅支持整個(gè)模型復(fù)用的現(xiàn)有模型共享平臺（例如 HuggingFace），配備按需模型復(fù)用功能的平臺允許開發(fā)者僅復(fù)用相關(guān)的模塊。這種方法類似于軟件工程中的實(shí)踐，其中細(xì)粒度模塊化可以減少用戶面臨的復(fù)用開銷。最近的一項(xiàng)工作，ModelFoundry [3]，已經(jīng)探索了這一場景。ModelFoundry 建立了一個(gè)模塊化和組合系統(tǒng)，集成了幾種模塊化算法 [38, 40] 將模型分解為功能模塊。它維護(hù)了一個(gè)這些模塊的倉庫，具有搜索和成本估算能力以用于按需復(fù)用。目前，由于模塊化算法的限制，ModelFoundry 僅支持復(fù)用 CNN 模型。NeMo 可以進(jìn)一步擴(kuò)展該系統(tǒng)以支持 Transformer 模型的按需復(fù)用。

4.3.1 在視覺相關(guān)任務(wù)中的應(yīng)用。 具體來說，在典型的開發(fā)場景中，開發(fā)者經(jīng)常復(fù)用整個(gè)預(yù)訓(xùn)練模型，即使只需要其功能的一個(gè)子集 [13, 46, 68]。例如，Rock-Paper-Scissors [13] 項(xiàng)目復(fù)用在 ImageNet 上預(yù)訓(xùn)練的 ViT 模型，并在包含“石頭”、“布”和“剪刀”圖像的三類分類數(shù)據(jù)集上對其進(jìn)行微調(diào)。盡管下游任務(wù)僅包含三個(gè)類別，該項(xiàng)目復(fù)用了整個(gè) ViT 模型，其中包含大量對應(yīng)于不相關(guān)類別的冗余權(quán)重，產(chǎn)生了額外的復(fù)用開銷。NeMo 可以通過使用模塊化訓(xùn)練預(yù)訓(xùn)練模型來緩解這個(gè)問題。因此，該項(xiàng)目可以僅復(fù)用相關(guān)模塊，例如識別手部相關(guān)圖像的模塊，并在石頭 - 剪刀 - 布數(shù)據(jù)集上微調(diào)該模塊，從而產(chǎn)生一個(gè)小得多的模型，具有更低的推理開銷。

我們在此案例中實(shí)驗(yàn)性地評估 NeMo，展示 NeMo 在實(shí)際場景中的有效性。由于 ViT 模型在 ImageNet 數(shù)據(jù)集上的巨大訓(xùn)練開銷（幾天或幾周），我們構(gòu)建了一個(gè)小型數(shù)據(jù)集用于模塊化訓(xùn)練，包括來自 CIFAR-10 的 5000 張飛機(jī)和 5000 張汽車圖像，以及來自 Sign Language MNIST [54] 的 7500 張手部圖像。我們使用 NeMo 在我們構(gòu)建的數(shù)據(jù)集上預(yù)訓(xùn)練一個(gè) ViT-small 模型，并在 Rock-Paper-Scissors 數(shù)據(jù)集 [13] 上微調(diào)對應(yīng)于“手”類別的模塊。為了公平比較，我們使用標(biāo)準(zhǔn)訓(xùn)練方法預(yù)訓(xùn)練一個(gè) ViT-small 模型，并在 Rock-Paper-Scissors 數(shù)據(jù)集上對其進(jìn)行微調(diào)。表 13 展示了準(zhǔn)確率和 NRR 方面的結(jié)果。與復(fù)用整個(gè)模型相比，復(fù)用模塊可以實(shí)現(xiàn)相當(dāng)?shù)臏?zhǔn)確率，同時(shí)僅保留 47.6% 的神經(jīng)元。62.54% 的 FLOPs 減少也表明推理開銷更低。

4.3.2 在文本相關(guān)任務(wù)中的應(yīng)用。 我們也在文本相關(guān)任務(wù)上評估 NeMo，特別是那些與軟件工程相關(guān)的任務(wù)。例如，考慮復(fù)用預(yù)訓(xùn)練的 CodeBERT [11] 模型進(jìn)行代碼克隆檢測的典型場景。CodeBERT [11] 在覆蓋六種編程語言（Go, Java, JavaScript, PHP, Python, Ruby）的 CodeSearchNet [12] 數(shù)據(jù)集上預(yù)訓(xùn)練。假設(shè)用戶的任務(wù)涉及僅在 Java 程序中檢測代碼克隆。傳統(tǒng)上，用戶必須復(fù)用整個(gè) CodeBERT 模型，即使只有其 Java 知識是相關(guān)的。相比之下，使用 NeMo，開發(fā)者可以以模塊化方式預(yù)訓(xùn)練 CodeBERT 并共享單個(gè)模塊，每個(gè)模塊專門針對特定的編程語言。然后用戶可以僅復(fù)用 Java 模塊進(jìn)行克隆檢測任務(wù)。該模塊明顯小于整個(gè)模型，導(dǎo)致推理成本大幅降低。

我們在此場景中實(shí)驗(yàn)性地評估 NeMo 的實(shí)用性。具體來說，基于預(yù)訓(xùn)練的 CodeBERT 模型，我們使用 NeMo 在 CodeSearchNet 數(shù)據(jù)集上執(zhí)行模塊化訓(xùn)練三個(gè)輪次，產(chǎn)生六個(gè)特定于語言的模塊，平均神經(jīng)元保留率為 39.23%。然后我們復(fù)用 Java 模塊，并在 BigCloneBench [49] 上對其進(jìn)行微調(diào)，這是一個(gè)僅由 Java 代碼組成的克隆檢測數(shù)據(jù)集。微調(diào)后的模塊在測試集上達(dá)到 96.55% 的準(zhǔn)確率。與復(fù)用整個(gè)預(yù)訓(xùn)練 CodeBERT 模型（達(dá)到 98.57% 準(zhǔn)確率）相比，復(fù)用 Java 模塊實(shí)現(xiàn)了相當(dāng)?shù)男阅?，同時(shí)僅保留 43.58% 的神經(jīng)元。這證明了 NeMo 在不犧牲性能的情況下顯著降低推理成本的有效性。

5 討論

5.1 NeMo 的通用性

鑒于深度學(xué)習(xí)的快速發(fā)展和多樣 DNN 模型的出現(xiàn)，DNN 模塊化方法的通用性很重要。我們認(rèn)為 NeMo 可以從兩個(gè)角度泛化到多種 DNN 模型：(1) 模型架構(gòu)：NeMo 在神經(jīng)元級別執(zhí)行模塊化訓(xùn)練，神經(jīng)元是所有 DNN 模型架構(gòu)中的基本組件。這使得 NeMo 適用于廣泛的神經(jīng)網(wǎng)絡(luò)模型。此外，NeMo 提供了神經(jīng)元級別的結(jié)構(gòu)化分解方法，從權(quán)重矩陣中移除不相關(guān)的權(quán)重，支持各種 DNN 模型（例如，CNN、FCNN、RNN 和基于 Transformer 的模型）的靈活和按需復(fù)用。(2) 模型規(guī)模：NeMo 引入了一種基于對比學(xué)習(xí)的新型模塊化訓(xùn)練損失，確保其在更大規(guī)模模型中的有效性和效率。隨著模型復(fù)雜性增加，可訓(xùn)練參數(shù)從大約 4,000 個(gè)卷積核擴(kuò)展到超過 40,000 個(gè)神經(jīng)元——在這個(gè)規(guī)模上，現(xiàn)有方法如 MwT [41] 難以維持訓(xùn)練和模塊化性能。通過結(jié)合對比學(xué)習(xí)優(yōu)化，NeMo 在基于 Transformer 的模型上實(shí)現(xiàn)了卓越的性能，有效地解決了大規(guī)模 DNN 模型模塊化中的可擴(kuò)展性挑戰(zhàn)。

為了評估 NeMo 的通用性，我們將其應(yīng)用于物體檢測任務(wù)。待模塊化的模型是 Detection Transformer 模型 (DETR) [4]，它具有 ResNet50 骨干網(wǎng)絡(luò)和 12 層 transformer 編碼器 - 解碼器組件。多物體檢測任務(wù)是一個(gè)實(shí)際但復(fù)雜的任務(wù)，由于同時(shí)預(yù)測各種目標(biāo)和邊界框，對 NeMo 提出了挑戰(zhàn)。為了簡化實(shí)現(xiàn)，我們使用了遙感數(shù)據(jù)集 RSOD（976 張圖像，包含 6,950 個(gè)物體）[28]，具有四個(gè)獨(dú)占類別：飛機(jī)、操場、立交橋和油箱。特別是，每張圖像包含多個(gè)僅屬于單一類別的物體實(shí)例。解碼器組件有一個(gè)分類頭和一個(gè)框頭。對于每次前向傳播，它有 100 個(gè)序列查詢輸入用于分類和框結(jié)果。在解碼器組件中應(yīng)用模塊化會導(dǎo)致這些結(jié)果的顯著損失。此外，100 個(gè)查詢對于我們的任務(wù)來說太多了，每個(gè)圖像最多有 20 個(gè)物體，導(dǎo)致模塊化訓(xùn)練過程中的額外計(jì)算強(qiáng)度。我們將其設(shè)置為 20，并選擇性地模塊化編碼器組件。對于評估指標(biāo)，我們采用 mAP@50 來量化檢測精度，同時(shí)通過神經(jīng)組件的內(nèi)聚 - 耦合分析來評估模塊化質(zhì)量。

模塊化 DETR 模型實(shí)現(xiàn)了 44.9% 的 mAP，表明與標(biāo)準(zhǔn)訓(xùn)練實(shí)現(xiàn)的 46.6% 相比，性能損失僅為 1.7%。此外，高模塊間內(nèi)聚性 (0.9990) 表明對于同類物體，它使用同一組神經(jīng)元。降低的耦合性 (0.7801) 表明對于不同類別的物體檢測，它使用一些不同的神經(jīng)元。結(jié)果表明 NeMo 可能適用于物體檢測任務(wù)。

我們還觀察到，與分類任務(wù)相比，物體檢測任務(wù)在耦合性方面的模塊化性能有所下降。限制 NeMo 在基于 DETR 的物體檢測中模塊化效能的主要因素包括：1) RSOD 數(shù)據(jù)集的有限規(guī)模（976 張圖像）和分辨率證明不足以進(jìn)行全面的特征學(xué)習(xí)；2) 遙感物體的固有復(fù)雜性甚至挑戰(zhàn)了標(biāo)準(zhǔn)訓(xùn)練范式（基線 mAP=46.6%）；3) 邊界框定位需要密集的上下文語義處理。觀察到的高耦合性（0.78）可能揭示了物體檢測的固有特征在圖像之間存在糾纏。

5.2 對有效性的威脅

5.2.1 外部有效性。 雖然我們主張 NeMo 可泛化到各種 DNN 模型，但由于巨大的時(shí)間和計(jì)算成本，在所有 DNN 模型上評估它是不現(xiàn)實(shí)的。為了減輕對通用性的潛在威脅，我們廣泛的實(shí)驗(yàn)使用了兩類主流 DNN 模型，包括兩個(gè)基于 Transformer 的模型和四個(gè)基于 CNN 的模型。此外，如同現(xiàn)有工作 [30, 31, 38, 41] 一樣，我們在計(jì)算機(jī)視覺任務(wù)上評估 NeMo。NeMo 在語言模型上的有效性，特別是像 CodeBERT 這樣的基于 Transformer 的模型，仍有待調(diào)查。鑒于 NeMo 支持基于 Transformer 的視覺模型的模塊化，它也可能適用于語言模型。我們將這些調(diào)查留待未來工作。

5.2.2 內(nèi)部有效性。 對內(nèi)部有效性的一個(gè)威脅可能來自對象選擇偏差。為了減少這種威脅，我們使用 CIFAR-10 和 SVHN 數(shù)據(jù)集，來自 PyTorch [35] 的 VGG16 和 ResNet18 模型，以及來自 HuggingFace [62] 的 ViT 和 DeiT 模型，這些都是組織良好且廣泛使用的。此外，為了減輕來自深度學(xué)習(xí)模型隨機(jī)性的威脅，我們使用不同的隨機(jī)種子在 SVHN 上重復(fù)了 ViT_s 模型訓(xùn)練實(shí)驗(yàn) 10 次。這些實(shí)驗(yàn)之間準(zhǔn)確率、NRR、內(nèi)聚性和耦合性的結(jié)果標(biāo)準(zhǔn)差分別為 0.005、0.0112、0.0047 和 0.0181，表明 NeMo 的性能是穩(wěn)定且一致的。

5.2.3 構(gòu)造有效性。 構(gòu)造有效性指的是我們使用的評估指標(biāo)。DNN 模塊化的內(nèi)聚性和耦合性指標(biāo)是在 MwT [41] 中提出的，并已被證明在評估 DNN 模塊化方法的模塊化性能方面是有效的。

5.3 局限性與未來工作

5.3.1 生成模型。 盡管 NeMo 的架構(gòu)表現(xiàn)出與各種模型的兼容性，但其在許多其他任務(wù)上的應(yīng)用仍然極具挑戰(zhàn)性。對于諸如 GPT 之類的生成模型，序列是通過自回歸輸出機(jī)制產(chǎn)生的，神經(jīng)元標(biāo)識符需要確定整個(gè)序列中每個(gè)令牌的激活神經(jīng)元。因此，這一過程產(chǎn)生的計(jì)算開銷是分類模型的數(shù)倍。對于其他類型的模型，如檢測、分割和強(qiáng)化學(xué)習(xí)模型，NeMo 仍然存在局限性。固有的復(fù)雜性源于大多數(shù)任務(wù)缺乏明確定義的類別邊界，且神經(jīng)單元表現(xiàn)出緊密的功能依賴關(guān)系——嘗試模塊化分解往往會導(dǎo)致難以承受的性能下降。即使是像物體檢測這樣的任務(wù)，其中邊界框注釋表面上在定義的類別內(nèi)操作，但對全局上下文理解的關(guān)鍵依賴從根本上挑戰(zhàn)了 NeMo?；谶@些挑戰(zhàn)，我們正在進(jìn)行的研究優(yōu)先考慮開發(fā)與標(biāo)簽無關(guān)的模塊化框架，特別強(qiáng)調(diào)其對生成架構(gòu)和其他主流模型的適用性。

5.3.2 模塊化訓(xùn)練開銷。 神經(jīng)元標(biāo)識符的集成在模塊化訓(xùn)練中引入了可衡量的計(jì)算開銷，通常在 30% 到 70% 之間，具體取決于模型架構(gòu)和硬件配置。實(shí)證評估揭示了具體表現(xiàn)：對于 CIFAR-10 上基于 ViT 的 10 類分類，標(biāo)準(zhǔn)訓(xùn)練大約在 5 小時(shí)內(nèi)完成，而模塊化訓(xùn)練需要 7 小時(shí)（額外 40% 的訓(xùn)練時(shí)間）。這一額外的訓(xùn)練時(shí)間與 MwT 實(shí)現(xiàn)中的觀察結(jié)果一致，證明了不同框架間開銷模式的一致性。雖然由于模型架構(gòu)不同，NeMo 和 MwT 之間的直接跨架構(gòu)比較證明具有挑戰(zhàn)性，但它們的相對開銷比率幾乎相同。值得注意的是，當(dāng)在同一模型上比較 NeMo 和 MwT 之間的損失函數(shù)時(shí)，NeMo 表現(xiàn)出更快的收斂速度，需要更少的訓(xùn)練輪次，突出了其固有的優(yōu)化優(yōu)勢。

考慮到額外的訓(xùn)練開銷，NeMo 適用于第 4.3 節(jié)討論的場景，即訓(xùn)練好的模型將被廣泛復(fù)用，并且下游任務(wù)需要特定功能。在這種情況下，NeMo 可以通過選擇性神經(jīng)元復(fù)用減少 FLOPs 數(shù)量，消除冗余神經(jīng)元，從而降低計(jì)算成本。

5.3.3 大規(guī)模數(shù)據(jù)集。 實(shí)證觀察揭示了將 NeMo 應(yīng)用于具有大規(guī)模數(shù)據(jù)集的 ViT 架構(gòu)時(shí)的一些挑戰(zhàn)。雖然研究人員試圖通過使用更大的數(shù)據(jù)集來提高 ViT 的準(zhǔn)確率，但我們嘗試將訓(xùn)練擴(kuò)展到 ImagenNet 數(shù)據(jù)集（1000 類）暴露了兩個(gè)基本約束：(1) 使用 NeMo 在大數(shù)據(jù)集上訓(xùn)練 ViT 總是需要幾天或幾周，以及 (2) 隨著類別數(shù)量的增加，出現(xiàn)性能下降，這是因?yàn)?，在每個(gè)類別中，可分配的神經(jīng)元更少。目前涉及耦合權(quán)重松弛和自適應(yīng)正則化的緩解策略僅實(shí)現(xiàn)了部分補(bǔ)救。我們未來的工作將專注于各種模型中的上述問題，包括判別模型和生成模型。

6 相關(guān)工作

傳統(tǒng)軟件系統(tǒng)的功能是模塊化的，使其更易于維護(hù)和進(jìn)一步開發(fā)。類似地，從軟件系統(tǒng)的角度來看，DNN 模型可以被視為一個(gè)以數(shù)據(jù)驅(qū)動方式構(gòu)建并配備各種功能的"復(fù)雜系統(tǒng)"。DNN 模塊化旨在將這樣一個(gè)具有其功能的"復(fù)雜系統(tǒng)"分解為一組子功能，從而促進(jìn) DNN 模型的維護(hù)和開發(fā)。為此，軟件工程社區(qū)一直在探索兩個(gè)主要方向：訓(xùn)練后模塊化和訓(xùn)練中模塊化技術(shù)。

6.1 訓(xùn)練后模塊化

訓(xùn)練后模塊化方法側(cè)重于通過識別負(fù)責(zé)相應(yīng)類別的權(quán)重組來分解訓(xùn)練良好的 DNN 模型 [30, 31, 37, 38, 40, 64, 65]。具體來說，Pan 等人開創(chuàng)了 DNN 模塊化 [30]，將多類分類全連接神經(jīng)網(wǎng)絡(luò)（FCNN）模型分解為模塊，每個(gè)模塊識別單個(gè)類別。該方法基于神經(jīng)元激活確定權(quán)重與特定類別的相關(guān)性，并將不相關(guān)的權(quán)重設(shè)為零，從而產(chǎn)生具有單一功能但大小與模型相同的模塊。他們隨后的工作將基于神經(jīng)元激活的思想擴(kuò)展到 CNN [31] 和 LSTM [20] 的模塊化。由于這些方法在單個(gè)權(quán)重級別衡量相關(guān)性，并且不物理移除不相關(guān)的權(quán)重，我們將它們歸類為非結(jié)構(gòu)化模塊化。相比之下，Qi 等人首次為 CNN 模型提出了結(jié)構(gòu)化模塊化方法，包括 CNNSplitter [37] 和 GradSplitter [40]。這些方法通過搜索相關(guān)卷積核（CNN 中的特殊結(jié)構(gòu)組件）并物理移除不相關(guān)的核來分解訓(xùn)練好的 CNN 模型，從而產(chǎn)生更小的模塊。然而，由于模型未針對模塊化進(jìn)行顯式訓(xùn)練，模型權(quán)重之間的相互依賴性和權(quán)重之間的高耦合性顯著限制了訓(xùn)練后模塊化技術(shù)的效率和有效性 [41]。

6.2 訓(xùn)練中模塊化

為了解決上述局限性，Qi 等人提出了一種新的 DNN 模塊化范式，訓(xùn)練中模塊化（MwT）[41]。MwT 旨在在訓(xùn)練階段識別和優(yōu)化 CNN 中針對不同功能的卷積核。它為 DNN 模塊化引入了高內(nèi)聚性和低耦合性的概念，并整合設(shè)計(jì)的內(nèi)聚性和耦合性損失以在整個(gè)訓(xùn)練過程中優(yōu)化這些屬性。通過將模塊化集成到訓(xùn)練階段，MwT 在效率（即模塊化的時(shí)間成本）和有效性（即模塊分類性能和大小）方面都實(shí)現(xiàn)了顯著改進(jìn)。然而，MwT 僅針對 CNN 模型實(shí)現(xiàn)，無法直接應(yīng)用于像 ViT 這樣的基于 Transformer 的架構(gòu)。鑒于基于 Transformer 模型的廣泛成功，我們提出了 NeMo，它通過將魯棒的模塊化能力擴(kuò)展到基于 Transformer 的模型，有效地解決了這一挑戰(zhàn)。

7 結(jié)論

在這項(xiàng)工作中，為了克服視覺變壓器模型模塊化訓(xùn)練和分解的困難，我們提出了一種新穎的神經(jīng)元級訓(xùn)練中模塊化框架 NeMo，它實(shí)現(xiàn)了基于 Transformer 模型的模塊化訓(xùn)練和結(jié)構(gòu)化模塊化，并且可以輕松擴(kuò)展到各種 DNN 模型。此外，我們基于對比學(xué)習(xí)方法優(yōu)化了內(nèi)聚性和耦合性損失的計(jì)算，在提高模型準(zhǔn)確率的同時(shí)顯著降低了卷積核和神經(jīng)元保留率。按需復(fù)用實(shí)驗(yàn)表明，我們的方法在基于 Transformer 和 CNN 模型的按需復(fù)用任務(wù)中都取得了更好的結(jié)果。此外，基于對比學(xué)習(xí)方法的改進(jìn)使得損失函數(shù)僅有一個(gè)可調(diào)超參數(shù)，該參數(shù)與準(zhǔn)確率和 NRR/KRR 均呈負(fù)相關(guān)。這使得算法更容易適應(yīng)新模型并實(shí)現(xiàn)更好的性能。

在未來，我們將把 NeMo 擴(kuò)展到更多模型結(jié)構(gòu)，并實(shí)現(xiàn)結(jié)構(gòu)化的模型分解。此外，我們將探索更多按需復(fù)用場景，并構(gòu)建更高效的方法以減少模塊化訓(xùn)練的開銷。

NeMo 的源代碼和實(shí)驗(yàn)結(jié)果可在 https://github.com/XiaohanBi-Hub/NeMo 獲取。

原文鏈接：https://arxiv.org/pdf/2508.11348