利用深度學習的力量來搜尋更好的引數空間來設計電路，比人腦的窮盡更有效率。

近日，Google在arxiv上發表一篇名為《Apollo: Transferable Architecture Exploration》的論文。

主要作者之一Amir Yazdanbakhsh在其部落格中詳細介紹了一個名為Apollo的研究專案。

這項工作可能標誌著使用機器學習技術來最佳化晶片「結構」的開始。

Apollo專案

Jeff Dean 在一年前的國際固態電路大會(International Solid State Circuits Conference)中發表過一次演講，認為「機器學習可以用於一些低層次的設計決策，也就是說在佈局和走線(place and route)上，晶片設計人員可以使用軟體來確定晶片的電路和佈局，類似於設計建築平面圖。」

相比之下，Apollo的目標不是要構建一個平面圖，而是要完成晶片的體系結構的設計。作者稱其為「架構探索」(architecture exploration)。

架構探索比佈局和走線的設計層級要更高，並且架構的改進能帶來比佈局走線更多的效能提升。

晶片的體系結構是一個晶片的功能元素設計，包括如何互動，軟體開發人員如何獲取硬體介面。經典的Intel x86處理器的片上儲存器大小、專用算術邏輯單元和暫存器等等全部的片上元素，以及他們組合起來的方式一起被稱為Intel架構。

作者稱Apollo為「第一個可遷移的架構探索系統」，隨著它設計更多不同的晶片，它的架構設計能力也會更好。也就是之前學習到的知識能夠遷移到新的晶片設計任務上。

作者團隊正在開發的晶片也是AI加速運算的晶片，和基於安培架構的GPU英偉達A100，Cerebras系統WSE晶片和其他創業公司的晶片類似。

「利用人工智慧來設計人工智慧晶片」，成為正反饋迴圈。

晶片的運算

由於任務是設計人工智慧晶片，Apollo正在探索的架構是適合執行神經網路的架構。神經網路的運算由大量簡單的數學單元、線性代陣列成，主要的運算是乘法和求和。

作者團隊認為設計任務的目的是找到一種正確的方式來混合這些簡單的數學算數來適配一個給定的AI任務。

他們首先選擇了一個相當簡單的AI任務，一個卷積神經網路MobileNet，2017年由Google的Andrew G. Howard和其團隊設計，是一種不需要太多資源的神經網路。此外，他們還測試了一些其他網路在不同任務（如目標檢測、語義分割）的工作負載情況。

透過這種方式，目標重新被定義為：給定一個神經網路，什麼晶片的引數結構能夠讓運算速度達到最快？

這個搜尋過程對超過4.52億個引數進行重排序，來尋找對於一個模型來說最優的引數。這些引數包括使用多少個數學單元（處理器），記憶體大小和啟用函式的記憶體。

Apollo是一個框架，它可以把每個部分換成其他論文中描述的黑盒最佳化方法，並根據特定的工作負載來調整這些方法，然後比較每種方法在解決特定目標下的表現。

作者使用了一些用來最佳化神經網路結構的最佳化方法，包括2019年Google的Quoc V.Le開發的進化方法；基於模型的強化學習方法；由Christof Angermueller開發用來設計DNA序列的最佳化方法population-based的整合方法；和一個基於貝葉斯的最佳化方法。

Apollo將神經網路設計和生物合成上的設計方法結合來設計電路，電路上取得更快的速度又會反過來促進神經網路設計和生物合成的發展。

Apollo框架將所有的最佳化方法都進行了比較，實驗結果詳細說明了「基於進化」和「基於模型」的方法優於隨機選擇和其他候選方法。

Apollo另一個重要的發現是執行這些最佳化方法的方式要比暴力搜尋更有效率。例如population-based的整合方法和半窮舉「semi-exhausitive」搜尋的方法更好。

population-based的方法能夠找到權衡的解決方案，例如計算機和記憶體的設計需要領域知識，而半窮舉搜尋是不具有學習能力的，所以population-based能夠找到半窮舉找不到的答案。

因此Apollo能夠指出不同最佳化方法之間在晶片設計上的差異。除此之外，Apollo能夠利用遷移學習提升最佳化方法的能力。

透過給Apollo不同的設計方案限制，例如最大的晶片尺寸不能超過毫米級，這些實驗結果可以作為一個後續的最佳化樣本作為輸入。作者團隊發現利用這種面積約束的電路設計能夠影響最終設計效能。

為MobileNet設計晶片中，負載的約束也是要在設計過程中考慮到的。作者團隊中Berkin Akin已經開發了一個新版本的MobileNet，命名為MobileNet Edge，指出最佳化需要晶片和神經網路最佳化的協作。

Akin認為「神經網路架構設計者必須瞭解目標的硬體架構，才能夠最佳化整體的系統性能和能效。」

Apollo對於給定負載約束的情況已經可以完成任務了，但是晶片設計和神經網路之間的協同最佳化，也許是下一個重點研究的目標。

參考連結：

https://www.zdnet.com/article/googles-deep-learning-finds-a-critical-path-in-ai-chips/

https://ai.googleblog.com/2021/02/machine-learning-for-computer.html

https://www.zdnet.com/article/google-experiments-with-ai-to-design-its-in-house-computer-chips/

https://arxiv.org/pdf/2102.01723.pdf