我們宣布發(fā)布JAX-Privacy 1.0，這是一個基于高性能計算庫JAX構(gòu)建的差分隱私機(jī)器學(xué)習(xí)工具庫。

從個性化推薦到科學(xué)進(jìn)步，AI模型正在幫助改善生活并改變各個行業(yè)。但這些AI模型的影響力和準(zhǔn)確性往往取決于其使用的數(shù)據(jù)質(zhì)量。大規(guī)模、高質(zhì)量的數(shù)據(jù)集對于開發(fā)準(zhǔn)確且具有代表性的AI模型至關(guān)重要，但必須以保護(hù)個人隱私的方式使用這些數(shù)據(jù)。

這就是JAX和JAX-Privacy發(fā)揮作用的地方。JAX于2020年推出，是一個專為大規(guī)模機(jī)器學(xué)習(xí)設(shè)計的高性能數(shù)值計算庫。其核心功能包括自動微分、即時編譯和跨多個加速器的無縫擴(kuò)展，使其成為高效構(gòu)建和訓(xùn)練復(fù)雜模型的理想平臺。JAX已成為推動AI邊界發(fā)展的研究人員和工程師的基石。其周圍的生態(tài)系統(tǒng)包括一套強(qiáng)大的領(lǐng)域特定庫，包括簡化神經(jīng)網(wǎng)絡(luò)架構(gòu)實現(xiàn)的Flax，以及實現(xiàn)最先進(jìn)優(yōu)化器的Optax。

基于JAX構(gòu)建的JAX-Privacy是一個強(qiáng)大的工具包，用于構(gòu)建和審計差分隱私模型。它使研究人員和開發(fā)者能夠快速高效地實現(xiàn)差分隱私算法，用于在大型數(shù)據(jù)集上訓(xùn)練深度學(xué)習(xí)模型，并提供將隱私訓(xùn)練集成到現(xiàn)代分布式訓(xùn)練工作流程所需的核心工具。JAX-Privacy的原始版本于2022年推出，旨在讓外部研究人員能夠重現(xiàn)和驗證我們在隱私訓(xùn)練方面的一些進(jìn)展。此后，它已發(fā)展成為Google內(nèi)部研究團(tuán)隊整合其在差分隱私訓(xùn)練和審計算法方面新穎研究見解的中心。

今天，我們自豪地宣布發(fā)布JAX-Privacy 1.0。這個新版本整合了我們最新的研究進(jìn)展，重新設(shè)計以實現(xiàn)模塊化，使研究人員和開發(fā)者比以往任何時候都更容易構(gòu)建差分隱私訓(xùn)練流水線，將最先進(jìn)的差分隱私算法與JAX提供的可擴(kuò)展性相結(jié)合。

多年來，研究人員一直將差分隱私作為量化和限制隱私泄露的黃金標(biāo)準(zhǔn)。差分隱私保證算法的輸出幾乎相同，無論數(shù)據(jù)集中是否包含單個個體或示例。

雖然差分隱私的理論已經(jīng)成熟，但在大規(guī)模機(jī)器學(xué)習(xí)中的實際實現(xiàn)仍然是一個挑戰(zhàn)。最常見的方法是差分隱私隨機(jī)梯度下降，需要定制的批處理程序、逐例梯度裁剪以及添加精心校準(zhǔn)的噪聲。這個過程計算量大，很難正確高效地實現(xiàn)，特別是在現(xiàn)代基礎(chǔ)模型的規(guī)模下。

JAX-Privacy使研究人員和開發(fā)者能夠使用最先進(jìn)的差分隱私算法在私有數(shù)據(jù)上以可擴(kuò)展和高效的方式訓(xùn)練和微調(diào)基礎(chǔ)模型，這要?dú)w功于其用于梯度裁剪和相關(guān)噪聲生成的原始構(gòu)建塊，這兩者在分布式環(huán)境中都能有效工作。

現(xiàn)有框架已經(jīng)取得了進(jìn)展，但在可擴(kuò)展性或靈活性方面往往不足。我們的工作一直在推動隱私機(jī)器學(xué)習(xí)的邊界，從開創(chuàng)新的差分隱私算法到開發(fā)復(fù)雜的審計技術(shù)。我們需要一個能夠跟上我們研究步伐的工具——一個不僅正確高效，而且從頭開始設(shè)計來處理最先進(jìn)模型的并行性和復(fù)雜性的庫。

JAX的函數(shù)式范式和強(qiáng)大的變換，如vmap（用于自動向量化）和shard_map（用于單程序多數(shù)據(jù)并行化），提供了堅實的基礎(chǔ)。通過基于JAX構(gòu)建，我們可以創(chuàng)建一個開箱即用的并行就緒庫，支持跨多個加速器和超級計算機(jī)訓(xùn)練大規(guī)模模型。JAX-Privacy是這一努力的成果，一個經(jīng)過時間考驗的庫，已為內(nèi)部生產(chǎn)集成提供支持，現(xiàn)在正與更廣泛的社區(qū)分享。

JAX-Privacy通過提供一套精心設(shè)計的組件簡化了差分隱私的復(fù)雜性：

JAX-Privacy實現(xiàn)了各種基礎(chǔ)工具，用于裁剪、噪聲添加、批選擇、記賬和審計，這些工具可以以各種方式組合來構(gòu)建端到端的差分隱私訓(xùn)練計劃。

JAX-Privacy最令人興奮的方面之一是其實際應(yīng)用。該庫設(shè)計用于支持用于預(yù)訓(xùn)練和微調(diào)大語言模型的現(xiàn)代機(jī)器學(xué)習(xí)框架。一個值得注意的例子是我們最近在訓(xùn)練VaultGemma時使用JAX-Privacy構(gòu)建塊，這是世界上最強(qiáng)大的差分隱私大語言模型。

通過這次開源發(fā)布，我們希望使開發(fā)者能夠通過流行的Keras框架僅用幾行代碼輕松微調(diào)大型模型。特別是，我們包含了微調(diào)Gemma系列模型的完整功能示例，這是由Google DeepMind基于Gemini構(gòu)建的開放模型集合。這些示例展示了如何將JAX-Privacy應(yīng)用于對話摘要和合成數(shù)據(jù)生成等任務(wù)，表明即使在使用最先進(jìn)的模型時，該庫也能提供最先進(jìn)的結(jié)果。

通過簡化差分隱私的集成，JAX-Privacy使開發(fā)者能夠從頭開始構(gòu)建保護(hù)隱私的應(yīng)用程序，無論他們是在為醫(yī)療應(yīng)用微調(diào)聊天機(jī)器人還是為個性化財務(wù)建議開發(fā)模型。它降低了隱私保護(hù)機(jī)器學(xué)習(xí)的入門門檻，使強(qiáng)大、負(fù)責(zé)任的AI更易獲得。

我們很興奮與研究社區(qū)分享JAX-Privacy。這次發(fā)布是多年專注努力的結(jié)果，代表了對隱私保護(hù)機(jī)器學(xué)習(xí)領(lǐng)域的重大貢獻(xiàn)。我們希望通過提供這些工具，能夠推動新一波造福所有人的研究和創(chuàng)新。

我們將繼續(xù)支持和開發(fā)該庫，整合新的研究進(jìn)展并響應(yīng)社區(qū)的需求。我們期待看到您使用JAX-Privacy構(gòu)建的成果。查看GitHub上的存儲庫或PIP包，立即開始訓(xùn)練保護(hù)隱私的機(jī)器學(xué)習(xí)模型。

JAX-Privacy包括以下貢獻(xiàn)者的貢獻(xiàn)：Leonard Berrada、Robert Stanforth、Brendan McMahan、Christopher A. Choquette-Choo、Galen Andrew、Mikhail Pravilov、Sahra Ghalebikesabi、Aneesh Pappu、Michael Reneer、Jamie Hayes、Vadym Doroshenko、Keith Rush、Dj Dvijotham、Zachary Charles、Peter Kairouz、Soham De、Samuel L. Smith、Judy Hanwen Shen。

Q&A

Q1：JAX-Privacy 1.0是什么？它解決了什么問題？

A：JAX-Privacy 1.0是谷歌基于JAX構(gòu)建的差分隱私機(jī)器學(xué)習(xí)工具庫。它解決了在大規(guī)模機(jī)器學(xué)習(xí)中實現(xiàn)差分隱私算法的挑戰(zhàn)，讓研究人員和開發(fā)者能夠在保護(hù)個人隱私的前提下，使用大型數(shù)據(jù)集訓(xùn)練AI模型，特別是在訓(xùn)練和微調(diào)基礎(chǔ)模型時。

Q2：差分隱私為什么重要？它如何保護(hù)隱私？

A：差分隱私是量化和限制隱私泄露的黃金標(biāo)準(zhǔn)。它保證算法的輸出幾乎相同，無論數(shù)據(jù)集中是否包含某個特定個體的數(shù)據(jù)。這意味著即使模型使用了你的數(shù)據(jù)進(jìn)行訓(xùn)練，也無法從模型輸出中推斷出你的個人信息，從而在獲得AI模型益處的同時保護(hù)個人隱私。

Q3：JAX-Privacy與現(xiàn)有框架相比有什么優(yōu)勢？

A：JAX-Privacy基于JAX的高性能計算能力，支持自動微分、即時編譯和多加速器擴(kuò)展，在可擴(kuò)展性和靈活性方面超越現(xiàn)有框架。它提供開箱即用的并行處理能力，支持大規(guī)模模型的分布式訓(xùn)練，并已經(jīng)在VaultGemma等世界級差分隱私大語言模型的訓(xùn)練中得到驗證。