SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\varepsilon (\boldsymbol {\alpha}) := \mathbb {E} _ {P \mathcal {S} _ {i} \sim \mathcal {Q}} \left[ \varepsilon_ {i} (\boldsymbol {\alpha}) \mid P ^ {\mathcal {S} _ {i}} \right] = \mathbb {E} _ {P \mathcal {S} _ {i} \sim \mathcal {Q}} \mathbb {E} _ {(\boldsymbol {X} _ {k} ^ {\mathcal {S} _ {i}}, \boldsymbol {Y} _ {k} ^ {\mathcal {S} _ {i}}) \sim P \mathcal {S} _ {i}} \left[ \hat {\varepsilon} _ {i k} (\boldsymbol {\alpha}) \right]
$$
Remark B.19.  $\hat{\varepsilon} (\alpha)$  quantifies the averaged error rate across source clients' finite samples.  $\varepsilon (\alpha)$  quantifies the expected error rate on a new data batch from a new client (target client). Noting that both error rates are defined with respect to the personalized model after adaptation.
B.3.2 Generalization bound for one hypothesis
Next, we derive generalization bounds for one fixed adaptation rate  $\alpha$ . Since we consider fixed  $\alpha$ , for clarity, we denote
$$
Z _ {i k} := \hat {\varepsilon} _ {i k} (\boldsymbol {\alpha})
$$
$$
\bar {Z} _ {i}. := \frac {1}{K} \sum_ {k = 1} ^ {K} Z _ {i k} = \hat {\varepsilon} _ {i} (\boldsymbol {\alpha})
$$
$$
\mu_ {i} := \mathbb {E} _ {\left(\boldsymbol {X} _ {k} ^ {\mathcal {S} _ {i}}, \boldsymbol {Y} _ {k} ^ {\mathcal {S} _ {i}}\right) \sim P ^ {\mathcal {S} _ {i}}} \left[ Z _ {i k} \right] = \varepsilon_ {i} (\boldsymbol {\alpha})
$$
$$
\bar {Z} _ {\cdot \cdot} := \frac {1}{N} \sum_ {i = 1} ^ {N} \bar {Z} _ {i}. = \hat {\epsilon} (\boldsymbol {\alpha})
$$
$$
\bar {\mu}. := \frac {1}{N} \sum_ {i = 1} ^ {N} \mu_ {i}
$$
$$
\mu := \mathbb {E} _ {P ^ {s} _ {i} \sim \mathcal {Q}} \mu_ {i} = \epsilon (\boldsymbol {\alpha})
$$
Intuitively, with enough number of source clients and number of batches, we have  $\bar{Z}_{..} \approx \bar{\mu}_{..} \approx \mu_{..}$
Lemma B.20 (Hoeffding's inequality). Let  $X_{1}, \dots, X_{n}$  be independent random variables such that  $a_{i} \leq X_{i} \leq b_{i}$  almost surely. Consider the sum of these random variables  $S_{n} = X_{1} + \dots + X_{n}$ . For all  $\epsilon > 0$ ,
$$
\Pr (S _ {n} - \mathbb {E} [ S _ {n} ] \geq \epsilon) \leq \exp \left(- \frac {2 \epsilon^ {2}}{\sum_ {i = 1} ^ {n} (b _ {i} - a _ {i}) ^ {2}}\right)
$$
Proof. Please refer to [13]
Lemma B.21 (Concentration of averaged client expected error rates). For any  $\epsilon > 0$ , we have
$$
\Pr (\bar {\mu}. - \mu \geq \epsilon) \leq \exp (- 2 N \epsilon^ {2})
$$
Proof. Notice that  $\mu_1, \dots, \mu_N$  are independent given  $\mathcal{Q}$ . For all  $i = 1, \dots, N$ ,  $\mathbb{E} \mu_i = \mu$  and  $0 \leq \mu_i \leq 1$ . Therefore,
$$
\begin{array}{l} \Pr (\bar {\mu}. - \mu \geq \epsilon) = \Pr \left(\sum_ {i = 1} ^ {N} \mu_ {i} - N \mu \geq N \epsilon\right) \\ = \Pr \left(\sum_ {i = 1} ^ {N} \mu_ {i} - \mathbb {E} \left[ \sum_ {i = 1} ^ {N} \mu_ {i} \right] \geq N \epsilon\right) \\ \leq \exp \left(- \frac {2 \cdot (N \epsilon) ^ {2}}{N \cdot (1 - 0) ^ {2}}\right) \quad (\text {H o e f f d i n g} \\ = \exp (- 2 N \epsilon^ {2}) \\ \end{array}
$$
Lemma B.22 (Concentration of client empirical error rate). For any  $\epsilon > 0$ ,
$$
\Pr \left(\bar {Z}.. - \bar {\mu}. \geq \epsilon\right) \leq \exp \left(- 2 N K \epsilon^ {2}\right)
$$
Proof. Given distributions  $P^{S_1},\dots ,P^{S_N}$ , we have  $Z_{11},\dots ,Z_{1K},Z_{21},\dots ,Z_{NK}$  are independent. For any  $i = 1,\dots ,N$  and  $k = 1,\dots ,K$ , we have  $\mathbb{E}_{(\pmb {X}_k^{S_i},\pmb {Y}_k^{S_i})\sim P^{S_i}}Z_{ik} = \mu_i$  and  $0\leq Z_{ik}\leq 1$ . Therefore,
$$
\Pr \left(\bar {Z} _ {\cdot \cdot} - \bar {\mu} _ {\cdot} \geq \epsilon \mid P ^ {\mathcal {S} _ {1}}, \dots , P ^ {\mathcal {S} _ {N}}\right) = \Pr \left(\sum_ {i = 1} ^ {N} \sum_ {k = 1} ^ {K} Z _ {i k} - \sum_ {i = 1} ^ {N} K \mu_ {i} \geq N K \epsilon \Bigg | P ^ {\mathcal {S} _ {1}}, \dots , P ^ {\mathcal {S} _ {N}}\right)
$$
$$
\begin{array}{l} = \Pr \left(\sum_ {i = 1} ^ {N} \sum_ {k = 1} ^ {K} Z _ {i k} - \mathbb {E} \left[ \sum_ {i = 1} ^ {N} \sum_ {k = 1} ^ {K} Z _ {i k} \right] \geq N K \epsilon \mid P ^ {\mathcal {S} _ {1}}, \dots , P ^ {\mathcal {S} _ {N}}\right) \\ \leq \exp \left(- \frac {2 \cdot (N K \epsilon) ^ {2}}{N K \cdot (1 - 0) ^ {2}}\right) \quad (\text {H o e f f d i n g} s \text {i n e q u a l i t y}) \\ = \exp (- 2 N K \epsilon^ {2}) \\ \end{array}
$$
Then, we use the tower property,
$$
\begin{array}{l} \Pr (\bar {Z} _ {\cdot \cdot} - \bar {\mu} _ {\cdot \cdot} \geq \epsilon) = \mathbb {E} _ {P ^ {\mathcal {S} _ {1}, \dots , P ^ {\mathcal {S} _ {N}} \stackrel {\text {i i d .}} {\sim}} Q} \Pr (\bar {Z} _ {\cdot \cdot} - \bar {\mu} _ {\cdot \cdot} \geq \epsilon | P ^ {\mathcal {S} _ {1}}, \dots , P ^ {\mathcal {S} _ {N}}) \\ \leq \sup  _ {P ^ {\mathcal {S} _ {1}}, \dots , P ^ {\mathcal {S} _ {N}}} \Pr (\bar {Z}. - \bar {\mu}. \geq \epsilon | P ^ {\mathcal {S} _ {1}}, \dots , P ^ {\mathcal {S} _ {N}}) \\ \leq \exp (- 2 N K \epsilon^ {2}) \\ \end{array}
$$
Proposition B.23 (Generalization for one hypothesis). For any fixed global model  $\mathbf{w}_G$  and adaptation rate  $\alpha$ , for any  $\epsilon > 0$ , we have
$$
\operatorname * {P r} \left(\left| \hat {\varepsilon} (\boldsymbol {\alpha}) - \varepsilon (\boldsymbol {\alpha}) \right| \geq \epsilon\right) \leq 4 \exp \left(- \frac {2 N K \epsilon^ {2}}{(\sqrt {K} + 1) ^ {2}}\right)
$$
Proof. For any  $\epsilon > 0$ , we have
$$
\begin{array}{l} \Pr (\bar {Z} _ {\cdot \cdot} - \mu \geq \epsilon) = \Pr ((\bar {Z} _ {\cdot \cdot} - \bar {\mu} _ {\cdot}) + (\bar {\mu} _ {\cdot} - \mu) \geq \epsilon) \\ \leq \inf  _ {\epsilon^ {\prime}} \left[ \Pr \left(\left(\bar {Z}.. - \bar {\mu}. \geq \epsilon^ {\prime}\right) \vee \left(\bar {\mu}. - \mu \geq \epsilon - \epsilon^ {\prime}\right)\right) \right] \\ \leq \inf  _ {\epsilon^ {\prime}} \left[ \Pr \left(\bar {Z}.. - \bar {\mu}. \geq \epsilon^ {\prime}\right) + \Pr \left(\bar {\mu}. - \mu \geq \epsilon - \epsilon^ {\prime}\right) \right] \\ \leq \inf  _ {\epsilon^ {\prime}} \left[ \exp \left(- 2 N K \left(\epsilon^ {\prime}\right) ^ {2}\right) + \exp \left(- 2 N \left(\epsilon - \epsilon^ {\prime}\right) ^ {2}\right) \right] \quad (\text {L e m m a B . 2 1 a n d B . 2 2}) \\ \end{array}
$$
To make the bound clear (although not optimal), we choose  $\epsilon' = \frac{1}{\sqrt{K} + 1}\epsilon$  and thus  $\epsilon - \epsilon' = \frac{\sqrt{K}}{\sqrt{K} + 1}\epsilon$ . Then the bound becomes,
$$
\Pr (\bar {Z} _ {\cdot \cdot} - \mu \geq \epsilon) \leq 2 \exp \left(- \frac {2 N K \epsilon^ {2}}{(\sqrt {K} + 1) ^ {2}}\right)
$$
Similarly we can show that
$$
\operatorname * {P r} (\bar {Z}. -. \mu \leq - \epsilon) \leq 2 \exp \left(- \frac {2 N K \epsilon^ {2}}{(\sqrt {K} + 1) ^ {2}}\right)
$$
$$
\Pr \left(| \hat {\varepsilon} (\boldsymbol {\alpha}) - \varepsilon (\boldsymbol {\alpha}) | \geq \epsilon\right) = \Pr \left(\left| \bar {Z} _ {\cdot .} - \mu \right| \geq \epsilon\right) \leq 4 \exp \left(- \frac {2 N K \epsilon^ {2}}{(\sqrt {K} + 1) ^ {2}}\right)
$$
Remark B.24. The RHS is function of both (1)  $N$ , the number of source clients and (2)  $K$ , the number of data batches on each client.
- When  $N \to \infty$ , given any fixed  $K \geq 1$ , the RHS  $\rightarrow 0$ , indicating that  $\hat{\epsilon}(\alpha) \stackrel{p}{\to} \epsilon(\alpha)$  (convergence in probability).
- However, given a fixed finite  $N$ , when  $K \to \infty$ , the RHS does not limit to zero. Intuitively, sampling more batches on finite source clients only helps the algorithm learn finite distribution  $P^{S_1}, \dots, P^{S_N}$ . However, more data batches on existing clients does not help further exploration of the meta-distribution  $\mathcal{Q}$  and generalization to novel target clients. Actually, given a fixed finite  $N$ , when  $K \to \infty$ ,  $\hat{\epsilon}(\alpha) \stackrel{p}{\to} \frac{1}{N} \sum_{i=1}^{N} \epsilon_i(\alpha) \neq \epsilon(\alpha)$ .
- If we put data from  $N$  sources (each with  $K$  batches) together as one source with  $NK$  batches. The generalization bound is looser.
B.3.3 Generalization bound for hypothesis space (proof of Theorem 5.1)
Finally, we derive the generalization bound for the hypothesis space. We first show in Lemma B.27 that  $\hat{\varepsilon}_{ij}(\alpha)$  is  $(LH)$ -Lipschitz to  $\alpha$ , then we apply standard generalization analysis in Theorem 5.1 based on covering number [32].
Definition B.25 ( $L$ -Lipschitz). The neural network  $f(\pmb{x}; \pmb{w})$  is  $L$ -Lipschitz w.r.t.  $\pmb{w}$ , if  $\forall \pmb{x}$  and  $\pmb{w}_1, \pmb{w}_2$ .
$$
\left\| f (\boldsymbol {x}; \boldsymbol {w} _ {1}) - f (\boldsymbol {x}; \boldsymbol {w} _ {2}) \right\| _ {2} \leq L \cdot \left\| \boldsymbol {w} _ {1} - \boldsymbol {w} _ {2} \right\| _ {2}
$$
Definition B.26 (H-module-wise-bounded update direction). The update direction is  $H$ -module-wise-bounded for a data batch  $X_{k}^{S_{i}}$  if
Lemma B.27 (Lipschitz error rate). Given a data batch  $(\mathbf{X}_k^{\mathcal{S}_i},\mathbf{Y}_k^{\mathcal{S}_i})$ , if the update direction is  $H$ -module-wise-bounded, given any two adaptation rates  $\alpha_{1},\alpha_{2}$  and the global model  $\mathbf{w}_{G}$ , we have
$$
\left| \hat {\varepsilon} _ {i j} (\boldsymbol {\alpha} _ {1}) - \hat {\varepsilon} _ {i j} (\boldsymbol {\alpha} _ {2}) \right| \leq L H \cdot \| \boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2} \| _ {2}
$$
$$
\begin{array}{l} \left| \hat {\varepsilon} _ {i j} (\boldsymbol {\alpha} _ {1}) - \hat {\varepsilon} _ {i j} (\boldsymbol {\alpha} _ {2}) \right| \\ = \left| \left(\frac {1}{B} \sum_ {b = 1} ^ {B} \left(1 - \left(\boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) ^ {\top} f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\boldsymbol {\alpha} _ {1}\right)\right)\right)\right) - \left(\frac {1}{B} \sum_ {b = 1} ^ {B} \left(1 - \left(\boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) ^ {\top} f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathscr {S} _ {i}} \left(\boldsymbol {\alpha} _ {2}\right)\right)\right)\right) \right| \\ = \left| \frac {1}{B} \sum_ {b = 1} ^ {B} \left(\boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) ^ {\top} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\boldsymbol {\alpha} _ {1}\right)\right) - f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\boldsymbol {\alpha} _ {2}\right)\right)\right) \right| \\ \leq \frac {1}{B} \sum_ {b = 1} ^ {B} \| \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}} \| _ {2} \cdot \left\| f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\boldsymbol {\alpha} _ {1}\right)\right) - f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\boldsymbol {\alpha} _ {2}\right)\right) \right\| _ {2} \\ \leq \frac {1}{B} \sum_ {b = 1} ^ {B} \| \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}} \| _ {2} \cdot L \cdot \left\| \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} (\boldsymbol {\alpha} _ {1}) - \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} (\boldsymbol {\alpha} _ {2}) \right\| _ {2} \quad (L \text {- L i p s c h i t z m o d e l}) \\ \leq \frac {1}{B} \sum_ {b = 1} ^ {B} \| \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}} \| _ {2} \cdot L \cdot H \cdot \| \boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2} \| _ {2} \tag {Lemma B.8} \\ = L H \cdot \left\| \alpha_ {1} - \alpha_ {2} \right\| _ {2} \\ \end{array}
$$
Remark B.28. Intuitively, Lemma B.27 shows that small change in  $\alpha$  will result in bounded change on  $\hat{\varepsilon}_{ij}(\alpha)$ .
Corollary B.29.  $\hat{\epsilon} (\alpha)$  and  $\epsilon (\alpha)$  are (LH)-Lipschitz w.r.t.  $\alpha$
Proof.  $\hat{\epsilon}(\alpha)$  and  $\epsilon(\alpha)$  are expectations of  $\hat{\varepsilon}_{ij}(\alpha)$  given the empirical and expected distribution of  $(\mathbf{X}_k^{S_i}, \mathbf{Y}_k^{S_i})$ . Lipschitz property is preserved.
Theorem 5.1 (Generalization for hypothesis space). Let  $\mathcal{H} = \{\alpha : \| \alpha \|_2 \leq R\}$  be the hypothesis space (space of adaptation rates),  $N$  be the number of source clients, and  $K$  be the number of data batches on each source client. Assuming (1)  $L$ -Lipschitz model, and (2)  $H$ -module-wise-bounded update direction. For any fixed global model  $\boldsymbol{w}_G$  and any  $\epsilon > 0$ , we have
$$
\Pr \left(\sup  _ {\boldsymbol {\alpha} \in \mathcal {H}} | \varepsilon (\boldsymbol {\alpha}) - \hat {\varepsilon} (\boldsymbol {\alpha}) | \geq \epsilon\right) \leq \left(\frac {1 2 L H R}{\epsilon}\right) ^ {d} \cdot 4 \exp \left(- \frac {N K \epsilon^ {2}}{2 (\sqrt {K} + 1) ^ {2}}\right) \tag {10}
$$
Proof. We use covering number to derive the generalization bound [32]. Define estimation error
$$
\Delta_ {\epsilon} (\boldsymbol {\alpha}) = \varepsilon (\boldsymbol {\alpha}) - \hat {\varepsilon} (\boldsymbol {\alpha})
$$
$$
\begin{array}{l} \left| \Delta_ {\epsilon} \left(\boldsymbol {\alpha} _ {1}\right) - \Delta_ {\epsilon} \left(\boldsymbol {\alpha} _ {2}\right) \right| = \left| \left[ \varepsilon \left(\boldsymbol {\alpha} _ {1}\right) - \hat {\varepsilon} \left(\boldsymbol {\alpha} _ {1}\right) \right] - \left[ \varepsilon \left(\boldsymbol {\alpha} _ {2}\right) - \hat {\varepsilon} \left(\boldsymbol {\alpha} _ {2}\right) \right] \right| \\ \leq \left| \varepsilon (\boldsymbol {\alpha} _ {1}) - \varepsilon (\boldsymbol {\alpha} _ {2}) \right| + \left| \hat {\varepsilon} (\boldsymbol {\alpha} _ {1}) - \hat {\varepsilon} (\boldsymbol {\alpha} _ {2}) \right| \\ \leq 2 L H \cdot \| \alpha_ {1} - \alpha_ {2} \| _ {2} \quad (\text {C o r o l l a r y B . 2 9}) \\ \end{array}
$$
$\mathcal{A} = \{\alpha : \| \alpha \|_2 \leq R, \alpha \in \mathbb{R}^d\}$  can be covered by  $K = \mathcal{N}_2(R,r)$  L2 balls with radius  $r = \frac{\epsilon}{4LH}$ . Lemma 6.27 in [32] shows that
$$
S = \mathcal {N} _ {2} (R, r) \leq \left(\frac {3 R}{r}\right) ^ {d} = \left(\frac {1 2 L H R}{\epsilon}\right) ^ {d}
$$
Denote these L2 balls to be  $B_{1},\dots ,B_{S}$
$$
\Pr \left(\sup  _ {\boldsymbol {\alpha} \in \mathcal {A}} | \varepsilon (\boldsymbol {\alpha}) - \hat {\varepsilon} (\boldsymbol {\alpha}) | \geq \epsilon\right) \leq \sum_ {s = 1} ^ {S} \Pr \left(\sup  _ {\boldsymbol {\alpha} \in B _ {s}} | \varepsilon (\boldsymbol {\alpha}) - \hat {\varepsilon} (\boldsymbol {\alpha}) | \geq \epsilon\right)
$$
For each ball  $B_{s}$ ,  $s = 1, \dots, S$ , denote the center to be  $\alpha_{s}$ . For any  $\alpha \in B_{s}$ , we have  $\| \alpha - \alpha_{s} \| \leq \frac{\epsilon}{4LH}$ , therefore
$$
\left| \Delta_ {\epsilon} (\boldsymbol {\alpha} _ {s}) - \Delta_ {\epsilon} (\boldsymbol {\alpha}) \right| \leq 2 L H \cdot \| \boldsymbol {\alpha} - \boldsymbol {\alpha} _ {s} \| \leq \frac {\epsilon}{2}
$$
Intuitively, every  $\alpha \in B_s$  has similar error rate. Therefore, the error rate for the whole ball is upper bounded, as long as the center  $\alpha_{s}$  has a small error rate
$$
\begin{array}{l} \Pr \left(\sup  _ {\boldsymbol {\alpha} \in B _ {s}} | \varepsilon (\boldsymbol {\alpha}) - \hat {\varepsilon} (\boldsymbol {\alpha}) | \geq \epsilon\right) = \Pr \left(\sup  _ {\boldsymbol {\alpha} \in B _ {s}} | \Delta_ {\epsilon} (\boldsymbol {\alpha}) | \geq \epsilon\right) \\ \leq \Pr \left(\sup  _ {\boldsymbol {\alpha} \in B _ {s}} \left[ \left| \Delta_ {\epsilon} (\boldsymbol {\alpha} _ {s}) \right| + \left| \Delta_ {\epsilon} (\boldsymbol {\alpha} _ {s}) - \Delta_ {\epsilon} (\boldsymbol {\alpha}) \right| \right] \geq \epsilon\right) \\ \leq \Pr \left(\sup  _ {\boldsymbol {\alpha} \in B _ {s}} | \Delta_ {\epsilon} (\boldsymbol {\alpha} _ {s}) | + \frac {\epsilon}{2} \geq \epsilon\right) \\ = \Pr \left(| \varepsilon (\boldsymbol {\alpha} _ {s}) - \hat {\varepsilon} (\boldsymbol {\alpha} _ {s}) | \geq \frac {\epsilon}{2}\right) \\ \end{array}
$$
Finally, by Proposition B.23, for each  $\alpha_{s}$
$$
\operatorname * {P r} \left(| \varepsilon (\pmb {\alpha} _ {s}) - \hat {\varepsilon} (\pmb {\alpha} _ {s}) | \geq \frac {\epsilon}{2}\right) \leq 4 \exp \left(- \frac {N K \epsilon^ {2}}{2 (\sqrt {K} + 1) ^ {2}}\right)
$$
Put all together
$$
\Pr \left(\sup  _ {\boldsymbol {\alpha} \in \mathcal {A}} | \varepsilon (\boldsymbol {\alpha}) - \hat {\varepsilon} (\boldsymbol {\alpha}) | \geq \epsilon\right) \leq \left(\frac {1 2 L H R}{\epsilon}\right) ^ {d} \cdot 4 \exp \left(- \frac {N K \epsilon^ {2}}{2 (\sqrt {K} + 1) ^ {2}}\right)
$$
C.1 Detailed experiment settings
C.1.1 CIFAR-10 experiments
Data preparation We use a benchmarking three-way split [50]: we randomly split the dataset to 300 clients, 240 of them are source clients and 60 are target clients. Each source client has 160 training samples and 40 validation samples, while each target client has 200 testing samples. We simulate three kinds of distribution shifts: feature shift, label shift, and hybrid shift. For feature shift, we follow [12, 17], randomly apply 15 different kinds of corruptions to the source clients (Figure 7a), and 4 new kinds of corruptions to the target clients (Figure 7b) to test the generalization of ATP. The corruption severity is randomly selected from  $\{1,2,3,4,5\}$ . For label shift, we use the step partition [5], where each client has 8 minor classes with 5 images per class, and 2 major classes with 80 images per class. For the hybrid shift, we apply both step partition and feature corruptions.
(a) 15 corruptions for training clients
Figure 7:  $15 + 4$  different corruptions we use to construct feature shift
Global model training We first train a global model with FedAvg [31] over the training sets of source clients.
- ResNet-18: The global model is ResNet-18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for  $T = 200$  communication rounds with full participation (cohort size  $C = 240$ ), local epochs  $E = 1$ , learning rate  $\eta = 0.01$  and batch size  $B = 20$ .
- Shallow CNN: The global model is a randomly initialized 5-layer CNN. We train the global model for  $T = 200$  communication rounds with full participation (cohort size  $C = 240$ ), local epochs  $E = 1$ , learning rate  $\eta = 0.1$  and batch size  $B = 20$ .
ATP training We initialize the adaptation rates as a all-zero vector, and optimize it over the validation sets of source clients. We optimize the adaptation rates for  $T = 200$  (for ResNet-18) or 400 (for Shallow CNN) communication rounds with partial participation (cohort size  $C = 60$ ), learning rate  $\eta = 0.1$  and batch size  $B = 20$ .
ATP testing We test the optimized adaptation rates on each target client. We use batch size  $B = 20$  by default, and test different batch size in Subsection 6.3.
C.1.2 CIFAR-100 experiments
Data preparation The data preparation is similar to CIFAR-10 experiments. The only difference is for label shift, each client has 98 minor classes with 1 image per class, and 2 major classes with 51 images per class. Same partition is applied to hybrid shift.
Global model training We first train a global model with FedAvg [31] over the training sets of source clients. The global model is ResNet-18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for  $T = 200$  communication rounds with full participation (cohort size  $C = 240$ ), local epochs  $E = 1$ , learning rate  $\eta = 0.01$  and batch size  $B = 20$ .
ATP training We initialize the adaptation rates as a all-zero vector, and optimize it over the validation sets of source clients. We optimize the adaptation rates for  $T = 200$  communication rounds with partial participation (cohort size  $C = 60$ ), learning rate  $\eta = 0.1$  and batch size  $B = 20$ .
ATP testing We test the optimized adaptation rates on each target client. We use batch size  $B = 20$ .
C.1.3 Digits-5 experiments
Data preparation Digits-5 dataset contains five domains: MNIST, SVHN, USPS, SynthDigits, and MNIST-M. We adopt the leave-one-domain-out evaluation protocol [10], i.e., one domain is chosen as the held-out testing domain, and the remaining domains are regarded as source training domains. We follow the data preprocessing in [25], while additionally applying step partition to inject label shift. Each domain is divided into 10 clients, leading to a total of 40 source clients and 10 target clients. Consequently, each client ends up with approximately 743 images spread across 10 classes. Each source client has  $80\%$  of its samples as training set and the remained  $20\%$  as testing set. Each client has 2 major classes and 8 minor class, where the ratio of images per class is approximately  $16:1$  (the same as our CIFAR-10 experiments). Since there is already domain shift, we do not add corruptions.
Global model training We first train a global model with FedAvg [31] over the training sets of source clients. The global model is ResNet-18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for  $T = 200$  communication rounds with full participation (cohort size  $C = 50$ ), local epochs  $E = 1$ , learning rate  $\eta = 0.01$  and batch size  $B = 20$ .
ATP training We initialize the adaptation rates as a all-zero vector, and optimize it over the validation sets of source clients. We optimize the adaptation rates for  $T = 200$  communication rounds with partial participation (cohort size  $C = 10$ ), learning rate  $\eta = 0.5$  and batch size  $B = 200$ .
ATP testing We test the optimized adaptation rates on each target client. We use batch size  $B = 200$ .
C.1.4 PACS experiments
Data preparation PACS dataset contains four domains: art, cartoon, photo, and sketch. We adopt the leave-one-domain-out evaluation protocol [10], i.e., one domain is chosen as the held-out testing domain, and the remaining domains are regarded as source training domains. We follow the data preprocessing in [10], while additionally applying step partition to inject label shift. Each domain is divided into 7 clients, leading to a total of 21 source clients and 7 target clients. Each source client has  $80\%$  of its samples as training set and the remained  $20\%$  as testing set. Each client has 2 major classes and 5 minor class, where the ratio of images per class is approximately  $16:1$  (the same as our CIFAR-10 experiments). Since there is already domain shift, we do not add corruptions.
Global model training We first train a global model with FedAvg [31] over the training sets of source clients. The global model is ResNet-18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for  $T = 200$  communication rounds with full participation (cohort size  $C = 21$ ), local epochs  $E = 1$ , learning rate  $\eta = 0.05$  and batch size  $B = 20$ .
ATP training We initialize the adaptation rates as a all-zero vector, and optimize it over the validation sets of source clients. We optimize the adaptation rates for  $T = 500$  communication rounds with full participation (cohort size  $C = 21$ ), learning rate  $\eta = 0.5$  and batch size  $B = 200$ .
C.1.5 Algorithm details
Assignment matrix  $\mathbf{A}$  In the main test, we mentioned that  $\mathbf{A} \in \mathbb{R}^{D \times d}$  is a  $0 - 1$  assignment matrix that maps each adaptation rate  $\alpha^{[l]}$  to the indices of the  $l$ -th module's parameters in  $\mathbf{w}$ . Mathematically,
$$
A _ {k l} = \left\{ \begin{array}{l l} 1, & \text {i f t h e k - t h p a r a m e t e r i n w b e l o n g s t o t h e l - t h m o d u l e} \\ 0, & \text {o t h e r w i s e} \end{array} \right.
$$
If there are  $d = 3$  modules, each with 1, 2, and 3 parameters, so  $D = 1 + 2 + 3 = 6$ , the corresponding assignment matrix will be
$$
\boldsymbol {A} = \left[ \begin{array}{c c c} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \\ 0 & 0 & 1 \\ 0 & 0 & 1 \end{array} \right]
$$
We did our experiments with single NVIDIA Tesla V100 GPU. However, our experiment should only require less than 2GB of GPU memory.
C.2 Compatibility to model architecture (RQ1)
In this part, we evaluate ATP with two more model architectures: a 5-layer Shallow CNN as a smaller model and ResNet-50 as a larger model.