SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 3: A qualitative comparison between task balancing techniques.  $T$  representing the computation cost when no additional task-balancing techniques are employed.
Method	Balance Magnitude	Balance Learning	Grads Required	No Extra Tuning	FLOPs	Motivation
None	✓	✓	T	/
Uncertainty	✓	✓	2T	Homoscedastic uncertainty
Gradnorm	✓	✓	✓	✓	4T	Balance learning and magnitudes
DWA	✓	3T	Balance learning
Minimax (Ours)	✓	✓	2T	Task robust
4.2 CV Scenario: Multi-Modal Mask MAE
4.2.1 Model and Settings
Architectures An overview of MM-MAE is shown in Figure 4.1.2. Multi-Modal Mask MAE contains three encoders, each of which processes different modalities of one image. During pretraining, we try to recover each modality from its masked tokens. Each modality is divided into  $16 \times 16$  patches and then tokenize the patches with modality-dependent linear projections. Projected patches are concatenated into a sequence of tokens and given as input to the same transformer encoder. We also add a global token with 2D sine-cosine positional embeddings. Each task owns a specialized decoder, and the computational cost of decoders scales linearly with the number of tasks.
Tasks and Datasets We select two datasets with different scales, ImageNet1K [53] and ImageNetS50 [25], to conduct unsupervised training upstream to see whether the minimax pre-training method can help the downstream tasks with poor performance. The classification task is evaluated on the validation part of the original dataset, while the semantic segmentation and depth estimation tasks are validated on the NYUv2 dataset [56] by fine-tuning. Due to the absence of a sizeable multi-task dataset with aligned task images [19, 2] we generate pseudo-labels on ImageNet and ImageNetS50 with GPT-3 and Mask2Former.
4.2.2 Quantitative Result
Classification tasks The quantitative results are presented in Table 2. We evaluate our models and baseline by fine-tuning them on the supervised ImageNetS50 and ImageNet1K. We fine-tune our models for 100 epochs and report the top-1 validation accuracy. The result shows a tiny gap between our method and the average method in the classification task. Classification tasks are regarded as the least challenging task category of the three. Cause different downstream tasks have different optimal parameter requirements, this gap is unavoidable. After the pre-training of the model reaches a specific step, the training weight of the classification tasks will continue to decrease.
Semantic segmentation tasks We further evaluate our models on semantic segmentation tasks on the NYUv2 dataset. We report the mean intersection over the union (mIoU) metric. Notice that semantic segmentation is the most challenging task of these downstream transfers. Our method
(a) Input
Figure 3: Comparative intermediate results. The first and final columns represent the image input and ground truth, respectively, while the intermediary images depict the intermediate results yielded by various task-balancing methodologies. Distinct colors correspond to the prediction of different objects. Our approach ensures robustness in downstream tasks.
(b) None
(c) Uncertainty
(d) Gradnorm
(e) DWA
(f) Minimax(Ours)
(g) Ground truth
benefits more than the average loss model from pseudo-labeled modalities as input. In particular, the correctness is improved by  $9.6\%$  on ImageNetS50 pre-training. With the progress of model training, our task-robust loss forces the model to improve poorly trained semantic segmentation tasks by increasing the training weight. The following section 5 will explain why a simple strategy can significantly help worst-case downstream tasks.
Depth evaluate tasks For depth estimation, we use NYUv2. We report  $\delta_{1}$  on the NUYv2 test set, showing the percentage of pixels  $p$  with error max  $\{\frac{\hat{y}_p}{y_p},\frac{y_p}{\hat{y}_p}\}$  less than 1.25 [20]. According to Table 2, the accuracy is improved by  $3.9\%$  on ImageNetS50 pre-training with the help of the downstream-task robustness loss function. The depth estimation task in the same data volume has a higher tolerance for prediction errors per pixel than semantic segmentation. However, it is still more complicated than the classification task, which only predicts the image once. After the classification task is well-trained, the depth estimation task will benefit from our strategy in the subsequent training.
Table 2 delineates several strategies designed to equilibrate the contribution of each task during the training of a multi-task network. For a qualitative comparison of these methods, refer to Table 3. We appraise these strategies based on several criteria [61]. An overview of our examination suggests that our approach achieves a synergistic blend of simplicity, efficiency, and effectiveness.
To facilitate a more intuitive comparison of the differences in results throughout the training process, we undertook downstream tasks in semantic segmentation, comparing the performance of various approaches midway through pre-training (400 epoch, ImageNetS50). As depicted in Figure 3, our methodology exhibits superior performance, even under conditions of insufficient training.
5 Explanation
We show why the proposed minimax pre-training method can be more effective than the average expected risk minimization in some cases. We consider a simplification of the model and the task relationship. Such a simplification makes our analysis convenient and intuitive. We assume that for all  $t \in [T]$ , the function  $\ell_t(\cdot, z)$  is  $\mu$ -strongly-convex,  $L$ -smooth, and  $L'$ -Lipschitz continuous for any fixed  $z \in \mathbb{Z}$  and the function  $\ell_t(\cdot, \cdot) \leq B$  for all  $\theta \in \Theta$  and  $z \in \mathbb{Z}^4$ . Note that pre-training on
4In fact, it is sufficient to assume that the expectation  $\mathbb{E}_z[\ell_t(\cdot ,z)]$  is  $\mu$ -strongly-convex,  $L$ -smooth, and  $L^{\prime}$ -Lipschitz continuous.
irrelevant upstream tasks does little help to the downstream tasks in general. Here, we only discuss the case where the upstream tasks and the downstream tasks are closed related. We ideally assume that the loss functions of the downstream tasks are convex combinations of the loss functions of the upstream tasks, i.e.,
$$
\ell_ {\lambda} = \sum_ {t = 1} ^ {T} \lambda_ {t} \ell_ {t}, \text {f o r a l l} \lambda \in \Lambda = \Delta_ {T}, \tag {9}
$$
where  $\Delta_T$  is the  $(T - 1)$ -dimensional probability simplex. We further assume that for each task there exists a parameter such that the expected risk of the task is zero, i.e.,  $\min_{\theta \in \Theta} \mathbb{E}_{z \sim P}[\ell_t(\theta, z)] = 0$  for all  $t \in [T]$ .
We first show that in the above setting, the proposed minimax optimization pre-training method can guarantee a better worst-case initial expected risk than the minimization method.
Proposition 5.1. Let  $\theta_{\text{max}}^*$  and  $\theta_{\text{average}}^*$  be the pre-trained parameters obtained by minimizing the maximal expected risk and the average expected risk, respectively. Then for the worst-case expected risks of the downstream tasks, we have
$$
\max  _ {\lambda \in \Lambda} \mathbb {E} _ {z \sim P} \left[ \ell_ {\lambda} \left(\theta_ {\text {m a x}} ^ {*}, z\right) \right] \leq \max  _ {\lambda \in \Lambda} \mathbb {E} _ {z \sim P} \left[ \ell_ {\lambda} \left(\theta_ {\text {a v e r a g e}} ^ {*}, z\right) \right]. \tag {10}
$$
Remark 5.2. The gap between  $\max_{\lambda \in \Lambda} \mathbb{E}_{z \sim P}\left[\ell_{\lambda}\left(\theta_{\mathrm{max}}^{*}, z\right)\right]$  and  $\max_{\lambda \in \Lambda} \mathbb{E}_{z \sim P}\left[\ell_{\lambda}\left(\theta_{\mathrm{average}}^{*}, z\right)\right]$  can be large. We provide an example in the appendix, where the ratio between  $\max_{\lambda \in \Lambda} \mathbb{E}_{z \sim P}\left[\ell_{\lambda}\left(\theta_{\mathrm{average}}^{*}, z\right)\right]$  and  $\max_{\lambda \in \Lambda} \mathbb{E}_{z \sim P}\left[\ell_{\lambda}\left(\theta_{\mathrm{max}}^{*}, z\right)\right]$  is as large as  $O(T)$ .
We then illustrate that a good initialization can serve as an implicit regularization. We suppose that the downstream tasks are trained with gradient descent. (For stochastic gradients with bounded variances, the analysis below also holds for sufficiently small step sizes, within neglectable approximation errors.) For the downstream task  $\lambda$ , with certain step sizes, the parameters will always be in a subset
$$
\Theta_ {\lambda} \left(\theta_ {0}\right) = \left\{\theta \in \Theta \mid \| \theta - \theta_ {\lambda} ^ {*} \| ^ {2} \leq \frac {2}{\mu} \mathbb {E} _ {z \sim P} \left[ \ell_ {\lambda} \left(\theta_ {0}, z\right) \right] \right\}, \tag {11}
$$
where  $\theta_0$  is the initial parameter and  $\theta_{\lambda}^{*} = \arg \min_{\theta \in \Theta}\mathbb{E}_{z\sim P}\left[\ell_{\lambda}\left(\theta ,z\right)\right]$
Proposition 5.3. Suppose that a function  $f: \mathbb{R}^d \mapsto \mathbb{R}$  is  $\mu_f$ -strongly-convex and  $L_f$ -smooth for all  $x \in \mathbb{R}^d$  and  $x^* \in \arg \min_{x \in \mathbb{R}^d} f(x)$ . Let  $\{x_k\}_{k=0}^{K-1}$  be the sequence generated by gradient descent with a step size  $\eta > 0$ , i.e.,  $x_k = x_{k-1} - \eta \nabla f(x_{k-1})$  for all  $k \in [K-1]$ . If the step size  $\eta \leq \frac{1}{L_f}$ , then we have
$$
\left\| x _ {k} - x ^ {*} \right\| ^ {2} \leq \frac {2}{\mu_ {f}} \left(f \left(x _ {0}\right) - f \left(x ^ {*}\right)\right), \tag {12}
$$
for all  $k = 0,1,\ldots ,K - 1$
By Proposition 5.3, we can deem that the downstream task's parameter space is the subset  $\Theta_{\lambda}(\theta_0)$ .
Consider the worst sample complexity to find an  $\epsilon$ -approximately optimal parameter by ERM within the parameter space  $\Theta_{\lambda}(\theta_0)$  for a downstream task  $\lambda \in \Lambda$ .
Theorem 5.4. The worst-case sample complexity  $\max_{\lambda \in \Lambda}N_{\lambda}\left(\theta_0,\epsilon ,\delta\right)$  with initialization  $\theta_0$  satisfies
$$
\max  _ {\lambda \in \Lambda} N _ {\lambda} \left(\theta_ {0}, \epsilon , \delta\right) \leq \frac {8 d B ^ {2}}{\epsilon^ {2}} \log \left(1 + \frac {1 6 L ^ {\prime}}{\epsilon} \sqrt {\frac {2}{\mu} \max  _ {\lambda \in \Lambda} \mathbb {E} _ {z \sim P} \left[ \ell_ {\lambda} \left(\theta_ {0} , z\right) \right]}\right) + \frac {8 B ^ {2}}{\epsilon^ {2}} \log \frac {2}{\delta}. \tag {13}
$$
Theorem 5.4 characterizes the upper bound of the worst-case sample complexity of downstream tasks. If we regard  $\epsilon$  and  $\delta$  as constants, the worst-case sample complexity with respect to the initialization  $\theta_0$  is  $O\left(\log \max_{\lambda \in \Lambda} \mathbb{E}_{z \sim P}\left[\ell_{\lambda}\left(\theta_0,z\right)\right]\right)$ . Combined with (10), Theorem 5.4 demonstrates that the proposed minimax pre-training procedure implies tighter sample complexity than the average minimization pre-training procedure in the worst case. Though the dependency on the worst-case initial expected risk is logarithmic in the upper bound analysis, we find that the initialization can have an evident effect on the generalization of the downstream tasks in practice. We claim that the upper bound for general cases may not be tight for our deep learning applications. Special structures in applications might lead to tighter bounds for generalization errors, which remains for further study.
This paper introduces the concept of downstream-task robustness for pre-training, aiming to improve the performance of foundation models across various downstream tasks. As models such as ChatGPT become more prevalent, safety and consistent performance are increasingly important. Our proposed minimax loss for pre-training, validated through extensive experiments, offers a potential solution to enhance the robustness and safety of such models. In the future, we will explore grouping the upstream tasks adaptively. We would say though still in its early stages, the study of downstream-task robustness holds significant promise for the reliable and safe deployment of AI infrastructure.
C. Fang and Z. Lin were supported by National Key R&D Program of China (2022ZD0160301). C. Fang was also supported by the NSF China (No. 62376008) and Wudao Foundation. Z. Lin was also supported by the NSF China (No. 62276004), the major key project of PCL, China (No. PCL2021A12) an Qualcomm.
[56] Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus. Indoor segmentation and support inference from rgbd images. In European Conference on Computer Vision (ECCV), 2012.
[57] Wilson L Taylor. “cloze procedure”: A new tool for measuring readability. Journalism quarterly, 1953.
[58] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
[59] Paul Tseng. On linear convergence of iterative methods for the variational inequality problem. Journal of Computational and Applied Mathematics, 1995.
[60] Paul Tseng. Accelerated proximal gradient methods for convex optimization. Technical report, Technical report, University of Washington, Seattle, 2008., 2008.
[61] Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc Proesmans, Dengxin Dai, and Luc Van Gool. Multi-task learning for dense prediction tasks: A survey. IEEE transactions on pattern analysis and machine intelligence, 2021.
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[63] Roman Vershynin. High-dimensional probability: An introduction with applications in data science. Cambridge university press, 2018.
[64] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Annual Conference on Empirical Methods in Natural Language Processing (EMNLP), 2018.
[65] Chen Wei, Haoqi Fan, Saining Xie, Chao-Yuan Wu, Alan Yuille, and Christoph Feichtenhofer. Masked feature prediction for self-supervised visual pre-training. In Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[66] Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In Conference on Computer Vision and Pattern Recognition (CVPR), pages 9653-9663, 2022.
[67] Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. Image bert pre-training with online tokenizer. In International Conference on Learning Representations (ICLR), 2021.
[68] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2008.
Minimax Optimization in Deep Learning. Minimax optimization has wide application in deep learning, e.g., adversarial robustness [11, 17], distributional robustness [22, 36, 54], adversarial generative models [27], imitation learning [68, 30]. In the field of pre-training, previous work has employed minimax optimization for adversarial robustness [15, 33]. Our work proposes a minimax optimization procedure for pre-training but with a different formulation. While the previous works aims at adversarial robustness, our work focuses on the worst-case generalization of downstream tasks.
Masked language modeling. Masked language modeling (MLM) was first proposed by [57] as a Cloze task. Adapted as a novel pre-training task, MLM and its autoregressive counterpart, e.g., BERT [18], GPT [49, 50, 9], and T5 [51], are highly successful methods to deal with NLP problems. MLM first masks out some tokens from the input sentences and then trains the model to retrieve the missing context from the rest of the tokens. These methods have been demonstrated to scale excellently so that various downstream tasks can utilize the pre-trained representations. In particular, BERT is constructed based on the transformer [62] model. After preparing the input samples, an embedding layer and a stack of Transformer layers are followed to conduct the bi-directional semantic modeling. We exploit BERT, the most typical masked language model, as the backbone model to process our ablation experiments.
Masked image modeling. Encouraged by transformers, which have gradually become a primary architecture for generic language understanding, ViT [21] later illusion the potential of adopting a pure transformer in image tasks. To generalize better for vision tasks and motivated by the success of BERT [18] in NLP, many recent works propose various masked image prediction methods for pre-training vision models in a self-supervised way. These methods reconstruct the target such as pixels [1, 14, 21, 24, 29, 66], discrete tokens [5, 67], and (deep) features [3, 65]. Notably, the masked autoencoder (MAE) [29] adopts an asymmetric design to allow the large encoder to operate only on unmasked patches and is followed by a lightweight decoder to reconstruct the complete signal from the latent representation along with mask tokens. MultiMAE [2] leverages the efficiency of the MAE approach and extends it to multi-modal and multitask settings. Based on MultiMAE, we apply our approach to increase the transfer capability for downstream tasks.
B.1 Convergence Rate of Algorithm 1
B.1.1 Proof of Theorem 3.1
For notation simplicity, we define that
$$
f _ {t} (\theta) := \mathbb {E} _ {z \sim P} \left[ \ell_ {t} (\theta , z) \right].
$$
$$
F _ {k} (\theta) := \sum_ {t = 1} ^ {T} w _ {\alpha , t} \left(\theta_ {k}\right) \mathbb {E} _ {z \sim P} \left[ \ell_ {t} (\theta , z) \right]
$$
$$
F (\theta) := \max  _ {t \in [ T ]} f _ {t} (\theta),
$$
where  $w_{\alpha ,t}(\theta) = \frac{\exp(\alpha\mathbb{E}_{z\sim P}[\ell_t(\theta,z)])}{\sum_{t' = 1}^{T}\exp(\alpha\mathbb{E}_{z\sim P}[\ell_{t'}(\theta,z)])}.$
Our proof consists of two parts. The first part is to show how well  $F_{k}(\theta_{k})$  approximates  $F(\theta_{k})$  with our choice of the softmax hyperparameter  $\alpha$ . The second part is to analyze the dynamics of the algorithm and the total convergence rate.
We first show that  $F_{k}(\theta_{k}) \geq F(\theta_{k}) + \frac{R_{0}L^{\prime}}{2\sqrt{k + 1}}$  for  $\alpha_{k} \geq \frac{4\sqrt{k + 1}}{R_{0}L^{\prime}}\log \frac{4TB\sqrt{k + 1}}{R_{0}L^{\prime}}$ . Define  $T_{k,\epsilon}(\theta) = \{t \in [T] \mid f_{t}(\theta_{k}) \geq F(\theta_{k}) - \epsilon\}$ . If  $\alpha_{k} \geq \frac{1}{\epsilon_{k}}\log \frac{TB}{\epsilon_{k}}$  for an  $\epsilon_{k} > 0$ , we have
$$
\begin{array}{l} F _ {k} (\theta_ {k}) = \sum_ {t = 1} ^ {T} w _ {\alpha_ {k}, t} (\theta_ {k}) f _ {t} (\theta_ {k}) \\ \geq \sum_ {t \in T _ {k, \epsilon_ {k}} (\theta_ {k})} w _ {\alpha_ {k}, t} \left(\theta_ {k}\right) f _ {t} \left(\theta_ {k}\right) \\ \overset {(A)} {\geq} \frac {\sum_ {t \in T _ {k , \epsilon_ {k}} (\theta_ {k})} \exp \left(\alpha_ {k} f _ {t} (\theta_ {k})\right)}{\sum_ {t ^ {\prime} = 1} ^ {T} \exp \left(\alpha_ {k} f _ {t ^ {\prime}} (\theta_ {k})\right)} \left(F \left(\theta_ {k}\right) - \epsilon_ {k}\right) \\ = \left[ 1 + \frac {\sum_ {t \notin T _ {k} , \epsilon_ {k} (\theta_ {k})} \exp \left(\alpha_ {k} f _ {t} \left(\theta_ {k}\right)\right)}{\sum_ {t ^ {\prime} \in T _ {k} , \epsilon_ {k} (\theta_ {k})} \exp \left(\alpha_ {k} f _ {t ^ {\prime}} \left(\theta_ {k}\right)\right)} \right] ^ {- 1} (F (\theta_ {k}) - \epsilon_ {k}) \\ \stackrel {(B)} {\geq} \left[ 1 + \frac {T \exp \left(\alpha_ {k} \left(F \left(\theta_ {k}\right) - \epsilon_ {k}\right)\right)}{\exp \left(\alpha_ {k} F \left(\theta_ {k}\right)\right)} \right] ^ {- 1} \left(F \left(\theta_ {k}\right) - \epsilon_ {k}\right) \\ = \left[ 1 + T \exp \left(- \alpha_ {k} \epsilon_ {k}\right) \right] ^ {- 1} \left(F \left(\theta_ {k}\right) - \epsilon_ {k}\right) \\ \begin{array}{l} \stackrel {(C)} {\geq} F (\theta_ {k}) - 2 \epsilon_ {k}. \end{array} \\ \end{array}
$$
The inequality A is due to the definition of  $T_{k,\epsilon}(\theta)$ . The inequality B is because  $\sum_{t \notin T_{k,\epsilon_k}(\theta_k)} \exp(\alpha_k f_t(\theta_k)) \leq T \exp(\alpha_k (F(\theta_k) - \epsilon_k))$  by the definition of  $T_{k,\epsilon}(\theta)$  and there exists  $t_k^* \in T_{k,\epsilon_k}(\theta)$  such that  $f_{t_k^*} = F(\theta_k)$ , which further implies  $\sum_{t' \in T_{k,\epsilon_k}(\theta_k)} \exp(\alpha_k f_{t'}(\theta_k)) \geq \exp(\alpha_k F(\theta_k))$ . The inequality C is because the hyperparameter  $\alpha_k \geq \frac{1}{\epsilon_k} \log \frac{T_B}{\epsilon_k}$ . Specifically, let  $\epsilon_k = \frac{R_0 L'}{4\sqrt{k + 1}}$  and  $\alpha_k \geq \frac{4\sqrt{k + 1}}{R_0 L'} \log \frac{4TB\sqrt{k + 1}}{R_0 L'}$  correspondingly, we have  $F_k(\theta_k) \geq F(\theta_k) + \frac{R_0 L'}{2\sqrt{k + 1}}$  for all  $k = 0, \ldots, K - 1$ .
We then analyze the dynamics of the algorithm and give the total convergence rate. For  $k = 0, \ldots, K - 1$ , it holds that
$$
\begin{array}{l} F _ {k} \left(\theta_ {k}\right) \stackrel {(A)} {\leq} F _ {k} \left(\theta^ {*}\right) + \left\langle \nabla F _ {k} \left(\theta_ {k}\right), \theta_ {k} - \theta^ {*} \right\rangle \tag {14} \\ \stackrel {(B)} {=} F _ {k} \left(\theta^ {*}\right) + \frac {1}{2 \eta} \left(\left\| \theta_ {k + 1} - \theta_ {k} \right\| ^ {2} + \left\| \theta_ {k} - \theta^ {*} \right\| ^ {2} - \left\| \theta_ {k + 1} - \theta^ {*} \right\| ^ {2}\right), \\ \end{array}
$$
where  $\theta^{*}\in \arg \max_{\theta \in \Theta}F(\theta)$ . The inequality A is due to the convexity of  $F_{k}(\theta)$ . The equality B is due to the update steps  $\theta_{k + 1} = \theta_k - \eta \nabla F_k(\theta_k)$  in Algorithm 1 and the fact  $\langle a,b\rangle = \frac{1}{2}\left(\| a\| ^2 +\| b\| ^2 -\| a - b\| ^2\right)$ .
Plugging the approximation error  $F_{k}(\theta_{k}) \geq F(\theta_{k}) + \frac{R_{0}L^{\prime}}{2\sqrt{k + 1}}$  and the inequality  $F_{k}(\theta) \leq F(\theta)$  for  $k = 0,\dots ,K - 1$  and all  $\theta \in \Theta$  into (14), we have
$$
F \left(\theta_ {k}\right) \leq F \left(\theta^ {*}\right) + \frac {1}{2 \eta} \left(\left\| \theta_ {k + 1} - \theta_ {k} \right\| ^ {2} + \left\| \theta_ {k} - \theta^ {*} \right\| ^ {2} - \left\| \theta_ {k + 1} - \theta^ {*} \right\| ^ {2}\right) + \frac {R _ {0} L ^ {\prime}}{2 \sqrt {k + 1}}. \tag {15}
$$
Taking the average over  $k = 0,\dots ,K - 1$  , we further have
$$
\begin{array}{l} \frac {1}{K} \sum_ {k = 0} ^ {K - 1} F (\theta_ {k}) \leq F (\theta^ {*}) + \frac {1}{2 \eta} \left(\sum_ {k = 0} ^ {K - 1} \| \theta_ {k + 1} - \theta_ {k} \| ^ {2} + \| \theta_ {0} - \theta^ {*} \| ^ {2} - \| \theta_ {K} - \theta^ {*} \| ^ {2}\right) + \frac {1}{K} \sum_ {k = 0} ^ {K - 1} \frac {R _ {0} L ^ {\prime}}{2 \sqrt {k + 1}} \\ \leq F \left(\theta^ {*}\right) + \frac {1}{2 \eta} \left(\sum_ {k = 0} ^ {K - 1} \left\| \theta_ {k + 1} - \theta_ {k} \right\| ^ {2} + \left\| \theta_ {0} - \theta^ {*} \right\| ^ {2}\right) + \frac {1}{K} \sum_ {k = 0} ^ {K - 1} \frac {R _ {0} L ^ {\prime}}{2 \sqrt {k + 1}} \\ \stackrel {(A)} {\leq} F \left(\theta^ {*}\right) + \frac {1}{2 \eta} \left(\sum_ {k = 0} ^ {K - 1} \left\| \theta_ {k + 1} - \theta_ {k} \right\| ^ {2} + \left\| \theta_ {0} - \theta^ {*} \right\| ^ {2}\right) + \frac {R _ {0} L ^ {\prime}}{\sqrt {K}} \\ \stackrel {(B)} {\leq} F \left(\theta^ {*}\right) + \frac {K L ^ {\prime 2} \eta}{2} + \frac {R _ {0} ^ {2}}{2 \eta} + \frac {R _ {0} L ^ {\prime}}{\sqrt {K}} \\ \stackrel {(C)} {=} F \left(\theta^ {*}\right) + \frac {2 R _ {0} L ^ {\prime}}{\sqrt {K}}. \tag {16} \\ \end{array}
$$
The inequality A is due to the fact  $\sum_{k=1}^{K} \frac{1}{\sqrt{k}} < 2\sqrt{K}$ . The inequality B is because for  $k = 0, \ldots, K-1$ , the function  $F_k(\theta)$  is  $L'$ -Lipschitz continuous, which implies  $\|\theta_{k+1} - \theta_k\|^2 = \eta^2 \|\nabla F_k(\theta_k)\|^2 \leq \eta^2 L'^2$ . The equality C is due to our choice for the step sizes, i.e.,  $\eta_k = \eta = \frac{R_0}{L' \sqrt{K}}$  for all  $k = 0, \ldots, K-1$ .