Exploring shared representations for personalized federated learning paper notes + code interpretation

The address of the paper is here

One . Introduce

In federated learning, due to the heterogeneous data on each client , As a result, the global training model cannot meet the requirements of each client . The author solves this problem by using the common representatives between clients . say concretely , The Federative learning problem with heterogeneous data is regarded as a parallel learning task , There may be some common structure between these tasks , The author's goal is to learn and utilize this common representation to improve the model quality of each client , Based on this, this paper puts forward FedRep( Federation represents learning ).
FedRep： Federation means learning to leverage all data stored across clients , Use gradient based updates to learn global low dimensional representations . Besides , Enables each client to compute a personalized 、 Low dimensional classifier , Responsible for the unique identification of local data for each client .

Two . Problem definition

Traditional federalism learns from n Optimize the following targets on clients ：

$$\underset{(q_1,...,q_n)\in {\mathcal{Q}}_n}{\min }\frac{1}{n}\sum _{i=1}^n{f}_i(q_i)$$

among $f_i$ It means the first one i Loss functions on clients ,$q_i$ It means the first one i Models on clients . However, there is less data on the client , At the same time, the number of clients is huge , The client cannot learn a model with a small loss , Therefore, federated learning allows parameter interaction between clients . The traditional way is to let the client learn a common model , That is to say $q_1=q_2=...=q_n$, But when the client data is obviously heterogeneous , The client model should be closer to the local data . So we need to learn a group $\{q_i\}$ Make it satisfied with its own data .
Learn a common expression (Learning a Common Representation). We consider a global representation $\varphi :{\mathbb{R}}^d\to {\mathbb{R}}^k$, Map data to a lower dimension k; Special presentation header of the client :${\mathbb{R}}^k\to \mathcal{Y}$. According to this , The first i A model on a client is a combination of local parameters and global representations on the client ：$q_i(x)=(h_i\circ \varphi )(x)$. It is worth noting that ,k Far less than d, That is to say, the number of parameters that each client must learn locally is very small . We rewrite our global optimization objectives according to the new content ：

$$\underset{\varphi \in \Phi }{\min }\frac{1}{n}\sum _{i=1}^n\underset{h_i\in \mathcal{H}}{\min }{f}_i(h_i\circ \varphi )$$

among $\Phi$ Is a feasible representation class , and $\mathcal{H}$ Is a feasible header . The client uses the data of all customers to learn the global model , At the same time, use your own local information to learn personalized headers .

3、 ... and . FedRep Algorithm

The algorithm idea is shown in the figure ：

Server and client learn together $\varphi$, The client learns its own parameter headers $h$.
Client update ： In every round , The selected client is trained . These clients come from the server ${\varphi }_i$ Update your own $h_i$, as follows ：

$$h_i^{t,s}=GRD(f_i(h_i^{t,s-1},{\varphi }^t),h_i^{t,s-1},\alpha )$$

GRD Is an optimal representation of gradient descent , It means that we have a set of parameters h stay f Use a gradient descent to $\alpha$ Update for step size . After training ${\tau }_h$ Step update h after , We are the same $\varphi$ Conduct ${\tau }_{\varphi }$ Secondary update , as follows ：

$${\varphi }_i^{t,s}=GRD(f_i(h_i^{t,{\tau }_h},{\varphi }_i^{t,s-1}),{\varphi }_i^{t,s-1},\alpha )$$

Server update ： After the client completes the update, it returns to the server ${\varphi }_i^{t,{\tau }_{\varphi }}$, After the service end is aggregated, it is averaged .
The algorithm is shown in the figure below ：

Four . Code details

The author's code points here
I believe that this article should not be difficult to understand , That is, it can be processed once in layers .
First , The main concern is how to layer . According to the idea , We need to divide into rep Layer and the head layer ,head For your own parameters .rep Is to participate in sharing , Before layering , Let's take a look at the Internet ：

class CNNCifar100(nn.Module):
    def __init__(self, args):
        super(CNNCifar100, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.drop = nn.Dropout(0.6)
        self.conv2 = nn.Conv2d(64, 128, 5)
        self.fc1 = nn.Linear(128 * 5 * 5, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, args.num_classes)
        self.cls = args.num_classes

        self.weight_keys = [['fc1.weight', 'fc1.bias'],
                            ['fc2.weight', 'fc2.bias'],
                            ['fc3.weight', 'fc3.bias'],
                            ['conv2.weight', 'conv2.bias'],
                            ['conv1.weight', 'conv1.bias'],
                            ]

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 128 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = self.drop((F.relu(self.fc2(x))))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

A very simple one CNN The Internet , We store the names of each layer , Facilitate layering .

if args.alg == 'fedrep' or args.alg == 'fedper':
    if 'cifar' in  args.dataset:
        w_glob_keys = [net_glob.weight_keys[i] for i in [0,1,3,4]]
    elif 'mnist' in args.dataset:
        w_glob_keys = [net_glob.weight_keys[i] for i in [0,1,2]]
    elif 'sent140' in args.dataset:
        w_glob_keys = [net_keys[i] for i in [0,1,2,3,4,5]]
    else:
        w_glob_keys = net_keys[:-2]

Here is a brief hierarchical operation , You can see that for us to deal with cifar100 Words ,rep Layer get yes 0 1 3 4, The corresponding is except fc3 The last floor of . So the last layer is head, Others are rep.
Then start training , Training is to obtain the parameters of the server for the client rep Add your own parameters head, The code is ：

if args.alg != 'fedavg' and args.alg != 'prox':
    for k in w_locals[idx].keys():
        if k not in w_glob_keys:
            w_local[k] = w_locals[idx][k]

among w_glob_keys Namely rep Parameters of ,w_local For all parameters .
The last is training ：

for iter in range(local_eps):
    done = False

    # for FedRep,  First we train head Fix rep, A few rounds of training 
    if (iter < head_eps and self.args.alg == 'fedrep') or last:
        for name, param in net.named_parameters():
            if name in w_glob_keys:
                param.requires_grad = False
            else:
                param.requires_grad = True
    
    #  Then train rep Fix head
    elif iter == head_eps and self.args.alg == 'fedrep' and not last:
        for name, param in net.named_parameters():
            if name in w_glob_keys:
                param.requires_grad = True
            else:
                param.requires_grad = False