SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbb {E} \| \mathbf {x} _ {t + 1} - \mathbf {x} _ {*} \| ^ {2} = \mathbb {E} \left\| \mathbf {x} _ {t} - \eta_ {t} \sum_ {s \in [ S ]} \lambda_ {s} ^ {t} \Delta_ {s} ^ {t} - x _ {*} \right\| ^ {2} \tag {70}
$$
where the first equality is due to strongly-convex objective functions, and we set  $\epsilon = \frac{\eta_t\mu}{2}$ .
$$
\sum_ {s \in [ S ]} \mathbb {E} \left[ \lambda_ {s} ^ {t} \right] \left(f _ {s} (\mathbf {x}) - f _ {s} \left(\mathbf {x} _ {*}\right)\right) \leq \frac {1}{2 \eta_ {t}} \left(1 - \frac {\eta_ {t} \mu}{2}\right) \| \mathbf {x} _ {t} - \mathbf {x} _ {*} \| ^ {2} - \frac {1}{2 \eta_ {t}} \| \mathbf {x} _ {t + 1} - \mathbf {x} _ {*} \| ^ {2} \tag {86}
$$
$$
+ \underbrace {\frac {1}{\mu} S ^ {2} \left(\alpha \eta_ {L} ^ {2} K ^ {2} D ^ {2} + \beta \sigma^ {2}\right) + \frac {\eta_ {t} S ^ {2} D ^ {2}}{2}} _ {\delta} \tag {87}
$$
Averaging using weight  $w_{t} = (1 - \frac{\mu\eta_{t}}{2})^{1 - t}$  and using such weight to pick output  $\mathbf{x}$ . By using Lemma 1 in [27] with constant learning rate  $\eta_{t} = \eta = \Omega \left(\frac{1}{\mu T}\right)$ , we have
$$
\mathbb {E} \left[ \Delta_ {Q} \right] \leq \left\| \mathbf {x} _ {0} - \mathbf {x} _ {*} \right\| ^ {2} \mu \exp \left(- \frac {\eta}{2} \mu T\right) + \mathcal {O} (\delta) \tag {88}
$$
where  $\delta = \frac{1}{\mu} S^2 (\alpha \eta_L^2 K^2 D^2 +\beta \sigma^2) + \frac{\eta S^2D^2}{2}$ .
By letting  $\beta = \eta$ ,  $\eta_L = \mathcal{O}\left(\frac{1}{\sqrt{T}}\right)$  and  $\eta = \Theta\left(\frac{\log(\max(1, \mu^2 T))}{\mu T}\right)$ ,
$$
\mathbb {E} \left[ \Delta_ {Q} \right] \leq \tilde {\mathcal {O}} \left(\frac {1}{T}\right). \tag {89}
$$
C.2 Non-convex Functions
Proof. Similar to Lemma 2 and taking expectation on the random data samples conditioning on  $\mathbf{x}_t$ , we have
$$
\begin{array}{l} \mathbb {E} f _ {s} \left(\mathbf {x} _ {t + 1}\right) \leq f _ {s} \left(\mathbf {x} _ {t}\right) + \mathbb {E} \left\langle \nabla f _ {s} \left(\mathbf {x} _ {t}\right), - \eta_ {t} \mathbf {d} _ {t} \right\rangle + \frac {1}{2} L \mathbb {E} \| \eta_ {t} \mathbf {d} _ {t} \| ^ {2} (90) \\ = f _ {s} \left(\mathbf {x} _ {t}\right) + \mathbb {E} \left\langle \nabla f _ {s} \left(\mathbf {x} _ {t}\right) - \Delta_ {s} ^ {t}, - \eta_ {t} \mathbf {d} _ {t} \right\rangle - \eta_ {t} \mathbb {E} \left\langle \Delta_ {s} ^ {t}, \mathbf {d} _ {t} \right\rangle + \frac {1}{2} L \mathbb {E} \| \eta_ {t} \mathbf {d} _ {t} \| ^ {2} (91) \\ \leq f _ {s} \left(\mathbf {x} _ {t}\right) + \mathbb {E} \left\langle \nabla f _ {s} \left(\mathbf {x} _ {t}\right) - \Delta_ {s} ^ {t}, - \eta_ {t} \mathbf {d} _ {t} \right\rangle - \eta_ {t} \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} + \frac {1}{2} L \mathbb {E} \| \eta_ {t} \mathbf {d} _ {t} \| ^ {2} (92) \\ \leq f _ {s} \left(\mathbf {x} _ {t}\right) + \frac {1}{2} \mathbb {E} \| \nabla f _ {s} \left(\mathbf {x} _ {t}\right) - \Delta_ {s} ^ {t} \| ^ {2} + \frac {1}{2} \eta_ {t} ^ {2} \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} - \eta_ {t} \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} + \frac {1}{2} L \mathbb {E} \eta_ {t} ^ {2} \| \mathbf {d} _ {t} \| ^ {2} (93) \\ = f _ {s} \left(\mathbf {x} _ {t}\right) + \frac {1}{2} \mathbb {E} \| \nabla f _ {s} \left(\mathbf {x} _ {t}\right) - \Delta_ {s} ^ {t} \| ^ {2} - \eta_ {t} \left(1 - \frac {1}{2} \eta_ {t} - \frac {1}{2} L \eta_ {t}\right) \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2}, (94) \\ \end{array}
$$
where  $\mathbf{d}_t = \sum_{s\in [S]}\lambda_s^{t,*}\Delta_s^t$  and  $\lambda_s^{t,*}$  is calculated by the accumulated stochastic gradients  $\Delta_s^t$ ,  $s\in [S]$ , but we drop the  $*$  of  $\lambda$  for simplicity.
With  $\eta_t \leq \frac{3}{2(1 + L)}$ , we have
$$
\begin{array}{l} \frac {\eta_ {t}}{4} \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} \leq - f _ {s} (\mathbf {x} _ {t + 1}) + f _ {s} (\mathbf {x} _ {t}) + \frac {1}{2} \mathbb {E} \| \nabla f _ {s} (\mathbf {x} _ {t}) - \Delta_ {s} ^ {t} \| ^ {2} (95) \\ \leq - f _ {s} \left(\mathbf {x} _ {t + 1}\right) + f _ {s} \left(\mathbf {x} _ {t}\right) + \frac {1}{2} \left(\alpha \eta_ {L} ^ {2} K ^ {2} D ^ {2} + \beta \sigma^ {2}\right) (96) \\ \end{array}
$$
$$
\frac {1}{T} \sum_ {t \in [ T ]} \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} \leq \frac {4 \left(f _ {s} (\mathbf {x} _ {1}) - \mathbb {E} f _ {s} (\mathbf {x} _ {T + 1})\right)}{\eta T} + 2 \left(\alpha \eta_ {L} ^ {2} K ^ {2} D ^ {2} + \beta \sigma^ {2}\right) \tag {97}
$$
Note that we want to use  $\left\| \bar{\mathbf{d}}_t \right\|^2$  are used as the metrics, where  $\bar{\mathbf{d}}_t = \boldsymbol{\lambda}_t^T \nabla (\mathrm{Diag}(\mathbf{FA}^\top))$  and  $\boldsymbol{\lambda}_t$  is calculated based on accumulated (stochastic) gradients  $\Delta_t$ . Then we have
$$
\left\| \bar {\mathbf {d}} _ {t} \right\| ^ {2} \leq 2 \left\| \sum_ {s \in [ S ]} \lambda_ {t} ^ {s} \nabla f _ {s} (\mathbf {x} _ {t}) - \mathbf {d} _ {t} \right\| ^ {2} + 2 \| \mathbf {d} _ {t} \| ^ {2}. \tag {98}
$$
With constant learning rate  $\eta_t = \eta$  and averaging from  $T$  communication rounds, we have
$$
\begin{array}{l} \frac {1}{T} \sum_ {t \in [ T ]} \mathbb {E} \left\| \bar {\mathbf {d}} _ {t} \right\| ^ {2} \leq \frac {1}{T} \sum_ {t \in [ T ]} 2 \mathbb {E} \left\| \sum_ {s \in [ S ]} \lambda_ {t} ^ {s} \nabla f _ {s} (\mathbf {x} _ {t}) - \mathbf {d} _ {t} \right\| ^ {2} + \frac {1}{T} \sum_ {t \in [ T ]} 2 \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} (99) \\ \leq \frac {1}{T} \sum_ {t \in [ T ]} 2 S \mathbb {E} \sum_ {s \in [ S ]} \left| \left| \lambda_ {t} ^ {s} (\nabla f _ {s} (\mathbf {x} _ {t}) - \Delta_ {s} ^ {t}) \right| \right| ^ {2} + \frac {1}{T} \sum_ {t \in [ T ]} 2 \mathbb {E} \| \mathbf {d} _ {t} \| ^ {2} (100) \\ \leq \frac {8 \left(f _ {s} \left(\mathbf {x} _ {1}\right) - \mathbb {E} f _ {s} \left(\mathbf {x} _ {T + 1}\right)\right)}{\eta T} + (2 S ^ {2} + 4) \left(\alpha \eta_ {L} ^ {2} K ^ {2} D ^ {2} + \beta \sigma^ {2}\right) (101) \\ \end{array}
$$
With constant learning rate  $\eta = \frac{1}{\sqrt{T}}$ , local learning rate  $\eta_{L} = \mathcal{O}\left(\frac{1}{T^{1/4}}\right)$  and  $\beta = \eta$ ,
$$
\frac {1}{T} \sum_ {t \in [ T ]} \mathbb {E} \left\| \bar {\mathbf {d}} _ {t} \right\| ^ {2} = \mathcal {O} \left(\frac {1}{\sqrt {T}}\right). \tag {102}
$$
D Further Experiments and Additional Results
In the following, we provide the detailed machine learning models for our experiments:
1) MultiMNIST Datasets and Learning Tasks: We test the convergence performance of our algorithms using the "MultiMNIST" dataset [53], which is a multi-task learning version of the MNIST dataset [57] from LIBSVM repository. Specifically, to convert the hand-written classification problem into a multi-task problem, we randomly chose 60000 images and divided them into  $M$  agents. Each agent has two tasks, where each task has  $n = 60000 / (2 * M)$  samples. Due to space limitations, we only present the convergence results for the case of non-i.i.d. data partition (i.e., data heterogeneity) and relegate the results of the i.i.d. data case to the appendix. For the non-i.i.d. data partition, we use the same data partition strategy as in [28], where each client can access data with at most two labels. In our experiments, a group of images is positioned in the top left corner, while another group of images is positioned in the bottom right. The two tasks are task "L" (to categorize the top-left digit) and task "R" (to classify the bottom-right digit). The overall problem is to classify the images of different tasks at different agents. All algorithms use the same randomly generated initial point. Here, we present experiments with  $M = 10$  agents, where each agent has two tasks (i.e.,  $\mathbf{A} \in \mathbb{R}^{M \times 2}$  is an all-one matrix). We set the local update rounds  $K = 10$ . Experiments with a larger number of agents ( $M = 5, 10, 30$ ) are provided here. The learning rates are chosen as  $\eta_L = 0.1$  and  $\eta_t = 0.1, \forall t$ .
2): River Flow Dataset and Learning Tasks: We further test our algorithms on FMOL problems of larger sizes. In this experiment, we use the River Flow dataset[54], which is for flow prediction flow at eight locations within the Mississippi River network. Thus, there are eight tasks in this problem. In this experiment, we set  $\eta_{L} = 0.001$ ,  $\eta_{t} = 0.1$ ,  $M = 10$ , and keep the batch size = 256 while comparing  $K$ , and keep  $K = 30$  while comparing the batch size. To better visualize 8 different tasks, we illustrate the normalized loss in radar charts in Fig. 2(b). We again verify that utilizing a larger training batch size and conducting additional local steps  $K$  results in accelerated convergence.
3): CelebA Dataset and Learning Tasks: We utilize the CelebA dataset [55], consisting of 200K facial images annotated with 40 attributes. We approach each attribute as a binary classification task, resulting in a 40-way multi-task learning (MTL) problem. To create a shared representation function, we implement ResNet-18 [58] without the final layer, attaching a linear layer to each attribute for classification. In this experiment, we set  $\eta_{L} = 0.0005$ ,  $\eta_{t} = 0.1$ ,  $M = 10$ , and  $K = 10$ . Figure 3 displays a radar chart depicting the loss value of each binary classification task. In Figure 3, we demonstrate the efficacy of our FMGDA and FSMGDA algorithms in both i.i.d. case and non-i.i.d. case.
Experiments on i.i.d. data: First, we compare the convergence results with the same experimental settings in our Section. 5 but tested on the i.i.d data. As shown in Fig. 4, both FMGDA and FSMGDA
(a) Non-i.i.d. case.
Figure 3: Experiments on CelebA dataset.
(b) i.i.d case.
(a) Loss v.s. Communication rounds.
Figure 4: Experiments on i.i.d. data.
(b) Loss v.s. Local update rounds.
successfully converged in i.i.d. data, and the algorithm with a larger training batch size and more local updates  $K$  may converge faster.
Impact of the number of clients: In this experiment, we choose the different number of clients from the discrete set  $\{5, 10, 30\}$  and fix learning rates at 0.1 and local update rounds at 10. As shown in Fig. 5, a larger number of workers leads to faster convergence rates of our proposed algorithms both in i.i.d. case and non-i.i.d. case; this is mainly because more samples have been used while training while having more workers.
(a) i.i.d. case.
(b) Non-i.i.d. case.
Figure 5: Loss value comparisons of algorithms on a different numbers of clients  $M$ .
(a) FMGDA.
Figure 6: Comparisons of different step-sizes.
(b) FSGMDA.
Impact of the Step-size: In this experiment, we choose the value of the learning rate  $\eta_L$  from the discrete set  $\{0.05, 0.01, 0.1\}$  and fix worker number at 5, local update rounds at 10. As shown in Fig. 6, larger local step-sizes lead to faster convergence rates on both FMGDA algorithm and FSMGDA algorithm.