FedJAX的庫結(jié)構(gòu)和內(nèi)容

聯(lián)邦學(xué)習(xí)是一種機(jī)器學(xué)習(xí)設(shè)置，允許多個(gè)客戶端（即移動(dòng)設(shè)備或者整個(gè)組織，取決于正在參與的任務(wù)）在一個(gè)中央服務(wù)器的編排下，協(xié)同訓(xùn)練同一個(gè)模型，同時(shí)還能保持訓(xùn)練數(shù)據(jù)的離散性。例如，通過聯(lián)邦學(xué)習(xí)，可以基于永遠(yuǎn)不會(huì)從移動(dòng)設(shè)備中消失的用戶數(shù)據(jù)訓(xùn)練虛擬鍵盤語言模型。

要實(shí)現(xiàn)這點(diǎn)，聯(lián)邦學(xué)習(xí)算法首先需要初始化服務(wù)器中的模型，然后完成以下對于每一輪訓(xùn)練而言都非常關(guān)鍵的三步：

1. 服務(wù)器將模型發(fā)送到一組采樣客戶端。

2. 這些采樣客戶端在本地?cái)?shù)據(jù)中訓(xùn)練模型。

3. 訓(xùn)練完成之后，客戶端將更新后的模型發(fā)送到服務(wù)器，然后服務(wù)器將所有這些模型匯總在一起。

隨著人們對隱私和安全的日益注重，聯(lián)邦學(xué)習(xí)已成為一個(gè)尤為活躍的研究領(lǐng)域。對于這個(gè)日新月異的領(lǐng)域，能夠輕松將想法轉(zhuǎn)換為代碼、快速迭代，以及比較和復(fù)制現(xiàn)有基線的重要性不言而喻。

日新月異的領(lǐng)域

https://research.google/pubs/pub49232/

因此，我們很高興為大家介紹 FedJAX。FedJAX 是一個(gè)基于 JAX 的開源庫，適用于注重研究易用性的聯(lián)邦學(xué)習(xí)模擬。FedJAX 擁有適用于執(zhí)行聯(lián)邦算法、預(yù)打包的數(shù)據(jù)集、模型和算法以及高模擬速度的簡單基本模塊，旨在讓研究員能夠更快速、更容易地開發(fā)和評(píng)估聯(lián)邦算法。

FedJAX

https://github.com/google/fedjax

JAX

https://github.com/google/jax

在這篇文章中，我們將討論 FedJAX 的庫結(jié)構(gòu)和內(nèi)容。我們會(huì)證明，在 TPU 中，F(xiàn)edJAX 可通過 EMNIST 數(shù)據(jù)集的聯(lián)合平均，在幾分鐘內(nèi)就能訓(xùn)練完模型。而通過 Stack Overflow 數(shù)據(jù)集的標(biāo)準(zhǔn)超參數(shù) （Hyperparameter），則需要將近 1 小時(shí)。

EMNIST

https://github.com/google/fedjax/blob/main/fedjax/datasets/emnist.py

聯(lián)合平均

https://fedjax.readthedocs.io/en/latest/fedjax.algorithms.html#module-fedjax.algorithms.fed_avg

Stack Overflow

https://github.com/google/fedjax/blob/main/fedjax/datasets/stackoverflow.py

庫結(jié)構(gòu)

FedJAX 注重易用性，因此僅引進(jìn)了少量新概念。使用 FedJAX 編寫的代碼與學(xué)術(shù)論文用于描述新穎算法的偽代碼類似，因此極易上手。除此之外，雖然 FedJAX 提供了聯(lián)邦學(xué)習(xí)的基本模塊，但用戶可以將其替換為最基本的實(shí)現(xiàn)（僅使用 NumPy 和 JAX），并且仍然可以將整體訓(xùn)練速度保持在一個(gè)合理的區(qū)間。

與學(xué)術(shù)論文用于描述新穎算法的偽代碼類似

https://github.com/google/fedjax/blob/main/README.md#quickstart

NumPy

https://numpy.org/

包含的數(shù)據(jù)集和模型

在當(dāng)前聯(lián)邦學(xué)習(xí)研究領(lǐng)域，存在各種各樣常用的數(shù)據(jù)集和模型，例如圖像識(shí)別 （Image recognition）、語言建模 （Language modeling） 等。越來越多這樣的數(shù)據(jù)集和模型無需安裝即可直接用于 FedJAX，因此用戶無需從頭開始編寫預(yù)處理數(shù)據(jù)集和模型。這不僅有利于對不同的聯(lián)邦算法進(jìn)行有效比較，還加速了新算法的開發(fā)。

目前，F(xiàn)edJAX 與以下數(shù)據(jù)集和示例模型一起打包：

EMNIST-62，一項(xiàng)字符識(shí)別任務(wù)

https://github.com/google/fedjax/blob/main/fedjax/datasets/emnist.py

Shakespeare，一項(xiàng)下一字符預(yù)測任務(wù)

https://github.com/google/fedjax/blob/main/fedjax/datasets/shakespeare.py

Stack Overflow，一項(xiàng)下一字詞預(yù)測任務(wù)

https://github.com/google/fedjax/blob/main/fedjax/datasets/stackoverflow.py

除了以上標(biāo)準(zhǔn)設(shè)置，F(xiàn)edJAX 還提供用于創(chuàng)建新數(shù)據(jù)集和模型的工具，這些新數(shù)據(jù)集和模型可以與庫的其余內(nèi)容共同使用。

工具

https://fedjax.readthedocs.io/en/latest/fedjax.html#federated-data

此外，F(xiàn)edJAX 支持聯(lián)合平均的標(biāo)準(zhǔn)實(shí)現(xiàn)，也支持用于在分散式示例上訓(xùn)練共享模型的其他聯(lián)邦算法，例如自適應(yīng)聯(lián)邦優(yōu)化器、不可知聯(lián)合平均以及 Mime，從而讓比較和評(píng)估現(xiàn)有算法變得更加簡單。

自適應(yīng)聯(lián)邦優(yōu)化器

https://fedjax.readthedocs.io/en/latest/fedjax.algorithms.html#module-fedjax.algorithms.fed_avg

不可知聯(lián)合平均

https://fedjax.readthedocs.io/en/latest/fedjax.algorithms.html#module-fedjax.algorithms.agnostic_fed_avg

Mime

https://fedjax.readthedocs.io/en/latest/fedjax.algorithms.html#module-fedjax.algorithms.mime

性能評(píng)估

我們在兩項(xiàng)任務(wù)上對自適應(yīng)聯(lián)合平均的標(biāo)準(zhǔn) FedJAX 實(shí)現(xiàn)進(jìn)行了基準(zhǔn)測試：圖像識(shí)別任務(wù)（測試聯(lián)邦 EMNIST-62 數(shù)據(jù)集）和下一字詞預(yù)測任務(wù)（測試 Stack Overflow 數(shù)據(jù)集）。聯(lián)邦 EMNIST-62 數(shù)據(jù)集較小，由 3400 名用戶和他們創(chuàng)建的示例（共 62 個(gè)拉丁字母數(shù)字字符）構(gòu)成；而 Stack Overflow 數(shù)據(jù)集較大，由數(shù)百萬問題和答案構(gòu)成（這些問題和答案來自于擁有成千上萬名用戶的 Stack Overflow 論壇）。

自適應(yīng)聯(lián)合平均

https://openreview.net/pdf？id=LkFG3lB13U5

聯(lián)邦 EMNIST-62 數(shù)據(jù)集

https://github.com/google/fedjax/blob/main/fedjax/datasets/emnist.py

Stack Overflow 數(shù)據(jù)集

https://github.com/google/fedjax/blob/main/fedjax/datasets/stackoverflow.py

我們在專門用于機(jī)器學(xué)習(xí)的各種硬件上測量性能。對于聯(lián)邦 EMNIST-62，我們在 GPU （NVIDIA V100） 和 TPU（Google TPU v2 上的 1 個(gè) TensorCore）加速器上對單一模型進(jìn)行了 1500 輪訓(xùn)練（每輪 10 個(gè)客戶端）。

對于 Stack Overflow，我們在 GPU （NVIDIA V100）、單核 TPU（Google TPU v2 上 1 個(gè) TensorCore）及多核 TPU（Google TPU v2 上 8 個(gè) TensorCore）上對單一模型進(jìn)行了 1500 輪訓(xùn)練（每輪 50 個(gè)客戶端）。其中，在 GPU 上使用 jax.jit，在單核 TPU 上僅使用 jax.jit，而在多核 TPU 上使用 jax.pmap。在下方圖表中，我們記錄了每輪訓(xùn)練的平均完成時(shí)間、完整評(píng)估測試數(shù)據(jù)所需時(shí)間以及整體執(zhí)行時(shí)間（整體執(zhí)行包含訓(xùn)練和完整評(píng)估）。

通過標(biāo)準(zhǔn)超參數(shù)和 TPUs，聯(lián)邦 EMNIST-62 的整個(gè)實(shí)驗(yàn)可以在幾分鐘之內(nèi)完成，而 Stack Overflow.的實(shí)驗(yàn)需要 1 小時(shí)左右的時(shí)間。

我們還評(píng)估了隨著每輪客戶端數(shù)量增加之后的 Stack Overflow 平均每輪訓(xùn)練時(shí)長。通過比較圖表上 8 核 TPU 與單核 TPU 的平均每輪訓(xùn)練時(shí)長，我們很容易就能發(fā)現(xiàn)，如果每輪參與的客戶端數(shù)量較多，則使用多核 TPU 能極大縮短運(yùn)行時(shí)間（對微分化的不公開學(xué)習(xí)等應(yīng)用來說非常有幫助）。

微分化的不公開學(xué)習(xí)

https://openreview.net/forum？id=BJ0hF1Z0b

結(jié)論和未來研究方向

在這篇文章中，我們介紹了 FedJAX 這種適用于研究、速度較快且簡單易用的聯(lián)邦學(xué)習(xí)模擬庫。我們希望 FedJAX 能推動(dòng)聯(lián)邦學(xué)習(xí)的深入研究，同時(shí)引起人們對于該領(lǐng)域的更多關(guān)注。未來，我們計(jì)劃繼續(xù)發(fā)展現(xiàn)有算法集、聚合機(jī)制、數(shù)據(jù)集和模型。

歡迎各位隨時(shí)查閱我們的教程筆記本，或者親自體驗(yàn) FedJAX！

教程筆記本

https://fedjax.readthedocs.io/en/latest/

親自體驗(yàn) FedJAX

https://github.com/google/fedjax/blob/main/examples

若想進(jìn)一步了解 FedJAX 及其與 Tensorflow Federated 等平臺(tái)的關(guān)系，請參閱我們的論文、README 或常見問題解答。

責(zé)任編輯：haq