聯邦學習：按混合分佈劃分Non-IID樣本

我們在博文《聯邦學習：按病態獨立同分佈劃分Non-IID樣本》中學習了聯邦學習開山論文^[1]中按照病態獨立同分布（Pathological Non-IID）劃分樣本。 在上一篇博文《聯邦學習：按Dirichlet分佈劃分Non-IID樣本》中我們也已經提到了按照Dirichlet分佈劃分聯邦學習Non-IID資料集的一種演算法。下面讓我們來看按Dirichlet分佈劃分資料集的另外一種變種，即按混合分佈劃分Non-IID樣本，該方法為論文^[2]中首次提出。

該論文采取了一個重要的假設，那就是雖然聯邦學習每個client的資料是Non-IID的，但我們假設每個client的資料都來自於某個混合分佈(混合成分個數\(K\)為超引數可調)。

\[p(x|\theta_t) = \sum_{k=1}^Kz_{tk} p(x|\theta_{k}) \]

其中\(t\)意思為第\(t\)個client，\(z_{tk}\)為（不可觀測的）隱變數(latent variable)，意為第\(t\)個client中的資料來自成分\(k\)的概率。第\(t\)個client的某個樣本點\(x\)進行生成時，會從\(K\)個成分中選擇一個成分\(p(x|\theta_{k})\)進行取樣，選擇該成分的概率為\(\alpha_{tk}\)。

形象化的展示圖片如下：

有了這個假設， 那麼每個client的資料都可以視為來自這三個分佈的資料的混合（每個client的Non-IID區別只是混合比例係數各不相同而已，下面我們提到混合比例係數由Dirichlet分佈隨機生成），那我們相當於假定了每個client資料間的一種"相似性"，即在各節點資料表面的Non-IID（\(p(x|\theta_t)\)）中其實潛藏IID的成分(\(p(x|\theta_{k}),k=1,2,..K\))。經過我的實驗，一旦這樣劃分資料，那麼對於基準的個性化聯邦學習演算法都會提升精度， 但是^[2]作者提出了一種基於子模型整合的演算法來更加充分地利用這種相似性。比如，假設一個client一共有A、B、C這3個子成分， 那麼我們就設計三個子模型分別對這些成分進行學習，每個模型的引數可以作為成分資料分佈引數的一種體現。對於隱變數\(z_{tk}\)(做為子模型加權使用)，作者設計了EM演算法來進行推斷。

注意，這裡作者的思想讓我們聯想到高斯混合分佈。高斯混合分佈就假設每個節點的資料取樣自高斯混合分佈中的一個成分(對應一個聚類簇)，而經典的高斯混合聚類就是要確定每個節點和簇的的對應關係（並推斷出隱變數係數）, 可以參見我的部落格《統計學習：EM演算法及其在高斯混合模型(GMM)中的應用》。

接下來我們來看這個劃分演算法的函式如何設計。除了常規Dirichlet劃分演算法所要求的n_clients、n_classes、alpha等， 它還有一個專門的n_clusters引數，表示混合成分個數。我們來看函式原型：

def mixture_distribution_split_noniid(dataset, n_classes, n_clients, n_clusters, alpha):

我們解釋一下函式的引數，這裡dataset是torch.utils.Dataset型別的資料集，n_classes表示資料集裡樣本分類數，n_clusters是簇的個數（後面會解釋其含義，如果設定為-1，則就預設n_clusters=n_classes，即每個簇對應一個標籤類別），alpha 為Dirichlet分佈引數，用於控制clients之間的資料diversity（Non-IID多樣性）。該函式返回一個由n_client個client所需的樣本索引組成的列表組成的列表client_idcs。

接下來我們看這個函式的內容。這個函式的內容可以概括為：先將所有類別不重疊地劃分為n_clusters個簇（每個簇對應一個不同的標籤分佈，標籤不重疊）；再對每個簇c，將樣本按照Non-IID劃分給不同的clients（每個client的樣本數量按照dirichlet分佈來確定）。

首先，我們判斷n_clusters的數量，如果為-1，則預設每一個cluster對應一個資料class：

    if n_clusters == -1:
        n_clusters = n_classes

然後將打亂後的標籤集合\(\{0,1,...,n\_classes-1\}\)分為n_clusters個簇。注意，這就意為著每個簇對應的標籤集合沒有重疊，也就是說各個簇之間的樣本資料是Non-IID的。

    all_labels = list(range(n_classes))
    np.random.shuffle(all_labels)
    def avg_divide(l, g):
        """
        將列表`l`分為`g`個獨立同分布的group（其實就是直接劃分）
        每個group都有 `int(len(l)/g)` 或者 `int(len(l)/g)+1` 個元素
        返回由不同的groups組成的列表
        """
        num_elems = len(l)
        group_size = int(len(l) / g)
        num_big_groups = num_elems - g * group_size
        num_small_groups = g - num_big_groups
        glist = []
        for i in range(num_small_groups):
            glist.append(l[group_size * i: group_size * (i + 1)])
        bi = group_size * num_small_groups
        group_size += 1
        for i in range(num_big_groups):
            glist.append(l[bi + group_size * i:bi + group_size * (i + 1)])
        return glist
    clusters_labels = avg_divide(all_labels, n_clusters)

然後再根據上面劃分好的label集合建立key為label, value為簇id(group_idx)的字典，

    label2cluster = dict()  # maps label to its cluster
    for group_idx, labels in enumerate(clusters_labels):
        for label in labels:
            label2cluster[label] = group_idx

接著獲取資料集的索引

    data_idcs = list(range(len(dataset)))

之後，我們將根據樣本的label和前面建立的label->cluster對映，再將樣本劃分到對應簇裡。

    # 記錄每個cluster大小的向量
    clusters_sizes = np.zeros(n_clusters, dtype=int)
    # 儲存每個cluster對應的資料索引
    clusters = {k: [] for k in range(n_clusters)}
    for idx in data_idcs:
        _, label = dataset[idx]
        # 由樣本資料的label先找到其cluster的id
        group_id = label2cluster[label]
        # 再將對應cluster的大小+1
        clusters_sizes[group_id] += 1
        # 將樣本索引加入其cluster對應的列表中
        clusters[group_id].append(idx)

    # 將每個cluster對應的樣本索引列表打亂
    for _, cluster in clusters.items():
        rng.shuffle(cluster)

我們已經得到了屬於每個cluster的樣本索引，接著我們按照Dirichlet分佈再將每個cluster中的樣本Non-IID地劃分到各client上去。

    # 記錄某個cluster的樣本分到某個client上的數量
    clients_counts = np.zeros((n_clusters, n_clients), dtype=np.int64) 

    # 遍歷每一個cluster
    for cluster_id in range(n_clusters):
        # 對每個client賦予一個滿足dirichlet分佈的權重，用於該cluster樣本的分配
        weights = np.random.dirichlet(alpha=alpha * np.ones(n_clients))
        # np.random.multinomial 表示投擲骰子clusters_sizes[cluster_id](該cluster中的樣本數)次，落在各client上的權重依次是weights
        # 該函式返回落在各client上各多少次，也就對應著各client應該分得來自該cluster的樣本數
        clients_counts[cluster_id] = np.random.multinomial(clusters_sizes[cluster_id], weights)

    # 對每一個cluster上的每一個client的計數次數進行字首（累加）求和，
    # 相當於最終返回的是每一個cluster中按照client進行劃分的樣本分界點下標
    clients_counts = np.cumsum(clients_counts, axis=1)

然後，我們根據上面已經得到的屬於各cluster的樣本集合，和各cluster中樣本分到各client中的情況（我們已經得到了每一個cluster中按照client進行劃分的樣本分界點下標），合併歸納得到每一個client中分得的樣本情況。

    def split_list_by_idcs(l, idcs):
        """
        將列表`l` 劃分為長度為 `len(idcs)` 的子列表
        第`i`個子列表從下標 `idcs[i]` 到下標`idcs[i+1]`
        （從下標0到下標`idcs[0]`的子列表另算）
        返回一個由多個子列表組成的列表
        """
        res = []
        current_index = 0
        for index in idcs: 
            res.append(l[current_index: index])
            current_index = index

        return res
    
    clients_idcs = [[] for _ in range(n_clients)]
    for cluster_id in range(n_clusters):
        # cluster_split為一個cluster中按照client劃分好的樣本
        cluster_split = split_list_by_idcs(clusters[cluster_id], clients_counts[cluster_id])

        # 將每一個client的樣本累加上去
        for client_id, idcs in enumerate(cluster_split):
            clients_idcs[client_id] += idcs

最後，我們返回每個client對應的樣本索引：

    return clients_idcs

接下來我們在EMNIST資料集上呼叫該函式進行測試，並進行視覺化呈現。我們設client數量\(N=10\)，Dirichlet概率分佈的引數向量\(\bm{\alpha}\)滿足\(\alpha_i=0.4,\space i=1,2,...N\), 混合成分個數為3：

import torch
from torchvision import datasets
import numpy as np
import matplotlib.pyplot as plt

torch.manual_seed(42)

if __name__ == "__main__":

    N_CLIENTS = 10
    DIRICHLET_ALPHA = 1
    N_COMPONENTS = 3

    train_data = datasets.EMNIST(root=".", split="byclass", download=True, train=True)
    test_data = datasets.EMNIST(root=".", split="byclass", download=True, train=False)
    n_channels = 1


    input_sz, num_cls = train_data.data[0].shape[0],  len(train_data.classes)


    train_labels = np.array(train_data.targets)

    # 注意每個client不同label的樣本數量不同，以此做到Non-IID劃分
    client_idcs = mixture_distribution_split_noniid(train_data, num_cls, N_CLIENTS, N_COMPONENTS, DIRICHLET_ALPHA)


    # 展示不同client的不同label的資料分佈
    plt.figure(figsize=(20,3))
    plt.hist([train_labels[idc]for idc in client_idcs], stacked=True, 
            bins=np.arange(min(train_labels)-0.5, max(train_labels) + 1.5, 1),
            label=["Client {}".format(i) for i in range(N_CLIENTS)], rwidth=0.5)
    plt.xticks(np.arange(num_cls), train_data.classes)
    plt.legend()
    plt.show()

最終的視覺化結果如下：

可以看到，62個類別標籤在不同client上的分佈雖然不同，但相對下面的完全基於Dirichlet的樣本劃分演算法，每個client之間的資料分佈顯得"更加相似"，即看得出來都來自於一個混合分佈，這證明我們的混合分佈樣本劃分演算法是有效的。

參考

• [1] McMahan B, Moore E, Ramage D, et al. Communication-efficient learning of deep networks from decentralized data[C]//Artificial intelligence and statistics. PMLR, 2017: 1273-1282.

• [2] Marfoq O, Neglia G, Bellet A, et al. Federated multi-task learning under a mixture of distributions[J]. Advances in Neural Information Processing Systems, 2021, 34.