SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
q (\boldsymbol {y} | \boldsymbol {x}) = \frac {q (\boldsymbol {x} | \boldsymbol {y}) q (\boldsymbol {y})}{\sum_ {\boldsymbol {y}} q (\boldsymbol {x} | \boldsymbol {y}) q (\boldsymbol {y})} = \frac {p (\boldsymbol {x} | \boldsymbol {y}) q (\boldsymbol {y})}{\sum_ {\boldsymbol {y}} p (\boldsymbol {x} | \boldsymbol {y}) q (\boldsymbol {y})} = \frac {p (\boldsymbol {y} | \boldsymbol {x}) \cdot \frac {q (\boldsymbol {y})}{p (\boldsymbol {y})}}{\sum_ {\boldsymbol {y}} p (\boldsymbol {y} | \boldsymbol {x}) \cdot \frac {q (\boldsymbol {y})}{p (\boldsymbol {y})}}
$$
Therefore, we can calibrate the neural network on distribution  $q$  simply by adding  $\log \frac{q(\pmb{y})}{p(\pmb{y})}$  to the bias terms, i.e.,
$$
\begin{array}{l} f _ {c a l} (\pmb {x}; \pmb {w} _ {c a l}) _ {c} = \frac {\exp (\pmb {w} _ {c} ^ {\top} g (\pmb {x} ; \pmb {w} _ {g}) + b _ {c} + \log \frac {q (\pmb {e} _ {c})}{p (\pmb {e} _ {c})})}{\sum_ {c ^ {\prime} = 1} ^ {K} \exp (\pmb {w} _ {c ^ {\prime}} ^ {\top} g (\pmb {x} ; \pmb {w} _ {g}) + b _ {c ^ {\prime}} + \log \frac {q (\pmb {e} _ {c ^ {\prime}})}{p (\pmb {e} _ {c ^ {\prime}})}} \\ = \frac {\exp (\boldsymbol {w} _ {c} ^ {\top} g (\boldsymbol {x} ; \boldsymbol {w} _ {g}) + b _ {c}) \cdot \frac {q (\boldsymbol {e} _ {c})}{p (\boldsymbol {e} _ {c})}}{\sum_ {c ^ {\prime} = 1} ^ {K} \exp (\boldsymbol {w} _ {c ^ {\prime}} ^ {\top} g (\boldsymbol {x} ; \boldsymbol {w} _ {g}) + b _ {c ^ {\prime}}) \cdot \frac {q (\boldsymbol {e} _ {c ^ {\prime}})}{p (\boldsymbol {e} _ {c ^ {\prime}})}} \\ = \frac {p (\boldsymbol {e} _ {c} | \boldsymbol {x}) \cdot \frac {q (\boldsymbol {e} _ {c})}{p (\boldsymbol {e} _ {c})}}{\sum_ {c ^ {\prime} = 1} ^ {K} p (\boldsymbol {e} _ {c ^ {\prime}} | \boldsymbol {x}) \cdot \frac {q (\boldsymbol {e} _ {c ^ {\prime}})}{p (\boldsymbol {e} _ {c ^ {\prime}})}} \\ = q (\boldsymbol {y} = \boldsymbol {e} _ {c} | \boldsymbol {x}) \\ \end{array}
$$
B.1.2 Proof of Proposition 3.2
Proposition 3.2 (Adapting the BN layer to handle feature shift [38]). When the feature shift only causes differences in the first and second order moments of the feature activations  $\mathbf{z} = g(\mathbf{x})$  where  $g$  is the combination of layers before the BN layer, assuming independent activations, the feature shift can be removed by adapting running mean and variance of the BN layer.
Proof. Denote the source and target feature (marginal) distributions to be  $p(\pmb{x})$  and  $q(\pmb{x})$ . Given independent, activations, we only need to test the marginal distribution of each  $z \in \mathbf{z} = g(\pmb{x})$ . For each  $z$ , since the feature shift only introduces differences in the first and second order moments, there exists  $\Delta$  and  $r > 0$ , s.t.,  $\forall z_t \in \mathbb{R}$
$$
\Pr_ {\boldsymbol {x} \sim q} (z \geq z _ {t}) = \Pr_ {\boldsymbol {x} \sim p} \left(z \geq \frac {z _ {t} - \Delta}{r}\right)
$$
which indicates that the distribution of  $z$  is first shifted by  $\Delta$  and then scaled by  $r$ . Such distribution shift in the feature activation can be removed by adapting the running mean  $\mu_p$  and variance  $\sigma_p^2$
$$
\mu_ {q} = r \cdot \mu_ {p} + \Delta
$$
$$
\sigma_ {q} = \sigma_ {p} \cdot r
$$
As a result, for all  $t\in \mathbb{R}$
$$
\begin{array}{l} \Pr_ {\boldsymbol {x} \sim q} \left(\frac {z - \mu_ {q}}{\sigma_ {q}} \geq t\right) = \Pr_ {\boldsymbol {x} \sim q} (z \geq \mu_ {q} + \sigma_ {q} \cdot t) \\ = \Pr_ {\boldsymbol {x} \sim p} (z \geq \frac {\mu_ {q} + \sigma_ {q} \cdot t - \Delta}{r}) \\ = \operatorname * {P r} _ {\boldsymbol {x} \sim p} (z \geq \mu_ {p} + \sigma_ {p} \cdot t) \\ = \Pr_ {\boldsymbol {x} \sim p} \left(\frac {z - \mu_ {p}}{\sigma_ {p}} \geq t\right) \\ \end{array}
$$
which indicates that the feature shift is removed after normalization with running statistics  $\mu_q, \sigma_q$ .
B.2 Convergence analysis
In this part, we show that ATP has the same convergence guarantee as FedAvg [31]. We first show in Lemma B.5 and B.10 that ATP preserves convexity and smoothness, which are two important conditions in the analysis of convergence. Then we formally prove the convergence of ATP in Theorem B.11.
B.2.1 Definitions: local and global objective
For clarity, we first formally define the data generation process, and local/global objectives for optimization.
Figure 6: Data generation process
Data generation We consider a two-stage sampling process as illustrated in Figure 6.
- There are  $N$  source clients' distributions  $P^{\mathcal{S}_1}, P^{\mathcal{S}_2}, \dots, P^{\mathcal{S}_N}$  and  $M$  target clients' distribution  $P^{\mathcal{T}_1}, P^{\mathcal{T}_2}, \dots, P^{\mathcal{T}_M}$  i.i.d. drawn from a meta-distribution  $\mathcal{Q}$ .
- For each source client  $i$ 's distribution  $P^{\mathcal{S}_i}$ , there are  $K$  data batches  $(X_1^{\mathcal{S}_i}, Y_1^{\mathcal{S}_i})$ ,  $(X_2^{\mathcal{S}_i}, Y_2^{\mathcal{S}_i})$ , ...,  $(X_K^{\mathcal{S}_i}, Y_K^{\mathcal{S}_i})$  drawn i.i.d. from  $P^{\mathcal{S}_i}$ .
Each batch consists of  $B$  samples,  $(\pmb{X}_k^{S_i},\pmb{Y}_k^{S_i}) = \{(\pmb{x}_{k,b}^{S_i},\pmb{y}_{k,b}^{S_i})\}_{b = 1}^B$  where  $B$  is the batch size.
- For simplicity, we assume that all source client has the same number of batches  $K$  and batch size  $B$ .
Definition B.1 (Batch objective). Define the batch objective of the  $k$ -th batch on client  $\mathcal{S}_i$  to be
$$
F _ {i k} (\boldsymbol {\alpha}) = \frac {1}{B} \sum_ {b = 1} ^ {B} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}, \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}}, \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) \right.
$$
where  $\pmb{w}_k^{S_i} = \pmb{w}_G + (\pmb{A}\pmb{\alpha})\odot \pmb{h}_k^{S_i}$  and  $\pmb{h}_k^{S_i}$  is the update direction computed with  $\pmb{X}_k^{S_i} = \{\pmb{x}_{k,b}^{S_i}\}_{b = 1}^B$  with Eq. (4) and (5).
Definition B.2 (Local objective). Define the local objective of client  $i$  to be
$$
F _ {i} (\boldsymbol {\alpha}) = \frac {1}{K} \sum_ {k = 1} ^ {K} F _ {i k} (\boldsymbol {\alpha})
$$
Definition B.3 (Global objective). Define the global objective to be
$$
F (\boldsymbol {\alpha}) = \frac {1}{N} \sum_ {i = 1} ^ {N} F _ {i} (\boldsymbol {\alpha})
$$
B.2.2 ATP preserves convexity and smoothness
In this part, we show that ATP preserves convexity and smoothness, which are two important conditions in the analysis of convergence.
Definition B.4 (Convexity). A function  $f: \mathbb{R}^D \to \mathbb{R}$  is convex if for all  $x_1, x_2 \in \mathbb{R}^D$  and  $\lambda \in [0,1]$
$$
f \left(\lambda \boldsymbol {x} _ {1} + (1 - \lambda) \boldsymbol {x} _ {2}\right) \leq \lambda f (\boldsymbol {x} _ {1}) + (1 - \lambda) f (\boldsymbol {x} _ {2})
$$
Lemma B.5 (Convexity preserving). If  $\ell_{CE}(f(\pmb{x};\pmb{w}),\pmb{y})$  is convex w.r.t.  $\pmb{w}$  given any data sample  $(\pmb{x},\pmb{y})$ , then  $F_{i}(\alpha)$  is convex w.r.t.  $\pmb{\alpha}$ .
Proof. Noticing that  $\boldsymbol{w}_k^{S_i} = \boldsymbol{w}_G + (\boldsymbol{A}\alpha) \odot \boldsymbol{h}_k^{S_i}$  is linear to  $\alpha$ , linear transformation preserves convexity. For any update direction  $\boldsymbol{h}_k^{S_i}$  and data sample  $(\boldsymbol{x}_{k,b}^{S_i},\boldsymbol{y}_{k,b}^{S_i})$ , we find that
$$
\begin{array}{l} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {G} + \left(\boldsymbol {A} (\lambda \boldsymbol {\alpha} _ {1} + (1 - \lambda) \boldsymbol {\alpha} _ {2})\right) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}, \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) \right. \\ = \ell_ {C E} \left( \right.f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \lambda \left[ \boldsymbol {w} _ {G} + (\boldsymbol {A} \boldsymbol {\alpha} _ {1}) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}} \right] + (1 - \lambda) \left[ \boldsymbol {w} _ {G} + (\boldsymbol {A} \boldsymbol {\alpha} _ {2}) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}} \right], \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) \\ \leq \lambda \ell_ {C E} \left( \right.f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {G} + \left(\boldsymbol {A} \boldsymbol {\alpha} _ {1}\right) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}, \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) + (1 - \lambda) \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {G} + \left(\boldsymbol {A} \boldsymbol {\alpha} _ {2}\right) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}, \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right)\right. \\ \end{array}
$$
i.e.,  $\ell_{CE}(f(\pmb{x}_{k,b}^{S_i};\pmb{w}_G + (\pmb{A}\pmb{\alpha})\odot \pmb{h}_k^{S_i}),\pmb{y}_{k,b}^{S_i})$  is convex w.r.t.  $\pmb{\alpha}$
Finally, since
$$
F _ {i} (\boldsymbol {\alpha}) = \frac {1}{K B} \sum_ {k = 1} ^ {K} \sum_ {b = 1} ^ {B} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {G} + (\boldsymbol {A} \boldsymbol {\alpha}) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right), \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right)
$$
which is the average of  $KB$  convex functions, we have that  $F_{i}(\alpha)$  is also convex to  $\alpha$ .
Definition B.6 ( $\beta$ -smoothness). A function  $f: \mathbb{R}^D \to \mathbb{R}$  is  $L$ -smoothness with  $\beta > 0$  if for all  $\pmb{x}_1, \pmb{x}_2 \in \mathbb{R}^D$ ,
$$
\| \nabla f (\boldsymbol {x} _ {1}) - \nabla f (\boldsymbol {x} _ {2}) \| _ {2} \leq \beta \| \boldsymbol {x} _ {1} - \boldsymbol {x} _ {2} \| _ {2}
$$
Definition B.7 (H-module-wise-bounded update direction). The update direction is  $H$ -module-wise-bounded for a data batch  $X_{k}^{S_{i}}$  if
$$
\left\| \left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {[ l ]} \right\| _ {2} \leq H, \quad \forall l = 1, \dots , d
$$
where  $(\pmb{h}_k^{\mathcal{S}_i})^{[l]}$  is the update direction corresponding to the  $l$ -th module and  $d$  is the number of modules in the neural network.
Lemma B.8 (Lipschitz parameter). If the update direction is  $H$ -module-wise-bounded for a data batch  $X_{k}^{S_{i}}$ . Given two adaptation rates  $\alpha_{1},\alpha_{2}$  and the global model  $w_{G}$ , we have
$$
\left\| \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} (\boldsymbol {\alpha} _ {1}) - \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} (\boldsymbol {\alpha} _ {2}) \right\| _ {2} \leq H \cdot \left\| \boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2} \right\| _ {2}
$$
where  $\pmb{w}_k^{S_i}(\pmb{\alpha}_1) = \pmb{w}_G + (\pmb{A}\pmb{\alpha}_1) \odot \pmb{h}_k^{S_i}$  is the personalized model updated with  $\pmb{\alpha}_1$  as the adaptation rate.
$$
\begin{array}{l} \left\| \boldsymbol {w} _ {k} ^ {S _ {i}} \left(\boldsymbol {\alpha} _ {1}\right) - \boldsymbol {w} _ {k} ^ {S _ {i}} \left(\boldsymbol {\alpha} _ {2}\right) \right\| _ {2} = \left\| \left(\boldsymbol {w} _ {G} + \left(\boldsymbol {A} \boldsymbol {\alpha} _ {1}\right) \odot \boldsymbol {h} _ {k} ^ {S _ {i}}\right) - \left(\boldsymbol {w} _ {G} + \left(\boldsymbol {A} \boldsymbol {\alpha} _ {2}\right) \odot \boldsymbol {h} _ {k} ^ {S _ {i}}\right) \right\| _ {2} \\ = \left\| \left(\boldsymbol {A} \left(\boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2}\right)\right) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}} \right\| _ {2} \\ = \left\| h _ {k} ^ {\mathcal {S} _ {i}} \odot \left(\boldsymbol {A} \left(\boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2}\right)\right) \right\| _ {2} \\ = \sqrt {\sum_ {l = 1} ^ {d} \left\| \left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {[ l ]} \right\| _ {2} ^ {2} \left(\alpha_ {1} ^ {[ l ]} - \alpha_ {2} ^ {[ l ]}\right) ^ {2}} \\ \leq \sqrt {\sum_ {l = 1} ^ {d} H ^ {2} \left(\alpha_ {1} ^ {[ l ]} - \alpha_ {2} ^ {[ l ]}\right) ^ {2}} \\ = H \cdot \| \boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2} \| _ {2} \\ \end{array}
$$
Remark B.9. Lemma B.8 indicates that when the adaptation rate is perturbed by a little, the personalized model parameter  $\boldsymbol{w}_k^{S_i}$  is also only perturbed by a little.
Lemma B.10 (Smoothness preserving). If (1)  $\ell_{CE}(f(\pmb{x};\pmb{w}),\pmb{y})$  is  $\beta$ -smooth w.r.t.  $\pmb{w}$  given any data sample  $(\pmb{x},\pmb{y})$ , and (2) the update direction  $\pmb{h}_k^{\mathcal{S}_i}$  is  $H$ -module-wise-bounded for all data batches  $\pmb{X}_k^{\mathcal{S}_i}$ , then  $F_i(\alpha)$  is  $(H^{2}\beta)$ -smoothness w.r.t.  $\pmb{\alpha}$ .
Proof. We first give an upper bound of  $\| A^\top \mathrm{diag}(h_k^{S_i})\|_2$  when  $h_k^{S_i}$  is  $H$ -module-wise-bounded. The update direction  $h_k^{S_i} \in \mathbb{R}^D$  is the concatenation of update directions for each module  $\{(h_k^{S_i})^{[l]}\}_{l=1}^d$ , i.e.,
$$
\left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {\top} = \left[ \left(\left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {[ 1 ]}\right) ^ {\top}, \dots , \left(\left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {[ d ]}\right) ^ {\top} \right]
$$
where  $(\pmb{h}_k^{\mathcal{S}_i})^{[l]}$  is a column vector representing the update direction of the  $l$ -th module in the model. Similarly, any other vector  $\pmb{v} \in \mathbb{R}^D$  can be correspondingly expressed as
$$
\boldsymbol {v} ^ {\top} = \left[ \left(\boldsymbol {v} ^ {[ 1 ]}\right) ^ {\top}, \dots , \left(\boldsymbol {v} ^ {[ d ]}\right) ^ {\top} \right]
$$
$$
\begin{array}{l} \| \boldsymbol {A} ^ {\top} \operatorname {d i a g} \left(\boldsymbol {h} _ {k} ^ {S _ {i}}\right) \| _ {2} = \sup  _ {\boldsymbol {v} \in \mathbb {R} ^ {D}} \frac {\| \boldsymbol {A} ^ {\top} \operatorname {d i a g} \left(\boldsymbol {h} _ {k} ^ {S _ {i}}\right) \boldsymbol {v} \| _ {2}}{\| \boldsymbol {v} \| _ {2}} \\ = \sup  _ {\boldsymbol {v} \in \mathbb {R} ^ {D}} \frac {\left\| \boldsymbol {A} ^ {\top} \left(\boldsymbol {h} _ {k} ^ {\boldsymbol {S} _ {i}} \odot \boldsymbol {v}\right) \right\| _ {2}}{\left\| \boldsymbol {v} \right\| _ {2}} \\ = \sup  _ {\boldsymbol {v} \in \mathbb {R} ^ {D}} \sqrt {\frac {\sum_ {l = 1} ^ {d} \left[ \left(\left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {[ l ]}\right) ^ {\top} \boldsymbol {v} ^ {[ l ]} \right] ^ {2}}{\sum_ {l = 1} ^ {d} \left\| \boldsymbol {v} ^ {[ l ]} \right\| _ {2} ^ {2}}} \\ \leq \sup  _ {\boldsymbol {v} \in \mathbb {R} ^ {D}} \sqrt {\frac {\sum_ {l = 1} ^ {d} \left[ \left\| \left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) ^ {[ l ]} \right\| _ {2} \cdot \left\| \boldsymbol {v} ^ {[ l ]} \right\| _ {2} \right] ^ {2}}{\sum_ {l = 1} ^ {d} \left\| \boldsymbol {v} ^ {[ l ]} \right\| _ {2} ^ {2}}} \\ \leq \sup  _ {\boldsymbol {v} \in \mathbb {R} ^ {D}} \sqrt {\frac {\sum_ {l = 1} ^ {d} \left[ H \cdot \left\| \boldsymbol {v} ^ {[ l ]} \right\| _ {2} \right] ^ {2}}{\sum_ {l = 1} ^ {d} \left\| \boldsymbol {v} ^ {[ l ]} \right\| _ {2} ^ {2}}} \tag {Definition B.7} \\ = H \\ \end{array}
$$
We then prove that for any  $H$ -module-wise-bounded update direction  $\pmb{h}_k^{\mathcal{S}_i}$  and data sample  $(\pmb{x}_{k,b}^{S_i},\pmb{y}_{k,b}^{S_i})$ , we have  $\ell_{CE}(f(\pmb{x}_{k,b}^{S_i};\pmb{w}_k^{S_i}(\pmb{\alpha})),\pmb{y}_{k,b}^{S_i})$  is  $H^2\beta$ -smoothness w.r.t.  $\pmb{\alpha}$ .
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {\alpha} _ {1}} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} (\boldsymbol {\alpha} _ {1})\right), \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) - \nabla_ {\boldsymbol {\alpha} _ {2}} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {\mathcal {S} _ {i}}; \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} (\boldsymbol {\alpha} _ {2})\right), \boldsymbol {y} _ {k, b} ^ {\mathcal {S} _ {i}}\right) \right\| _ {2} \\ = \| \boldsymbol {A} ^ {\top} \left( \right.\boldsymbol {h} _ {k} ^ {S _ {i}} \odot \nabla_ {\boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {1})} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {S _ {i}}; \boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {1})\right), \boldsymbol {y} _ {k, b} ^ {S _ {i}}\right) - \boldsymbol {A} ^ {\top} \left( \right.\boldsymbol {h} _ {k} ^ {S _ {i}} \odot \nabla_ {\boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {2})} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {S _ {i}}; \boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {2})\right), \boldsymbol {y} _ {k, b} ^ {S _ {i}}\right)\left. \right\| _ {2} \\ \leq \| \boldsymbol {A} ^ {\top} \operatorname {d i a g} \left(\boldsymbol {h} _ {k} ^ {S _ {i}}\right) \| _ {2} \cdot \| \nabla_ {\boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {1})} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {S _ {i}}; \boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {1})\right), \boldsymbol {y} _ {k, b} ^ {S _ {i}}\right) - \nabla_ {\boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {2})} \ell_ {C E} \left(f \left(\boldsymbol {x} _ {k, b} ^ {S _ {i}}; \boldsymbol {w} _ {k} ^ {S _ {i}} (\boldsymbol {\alpha} _ {2})\right), \boldsymbol {y} _ {k, b} ^ {S _ {i}}\right) \| _ {2} \\ \leq \left\| \boldsymbol {A} ^ {\top} \operatorname {d i a g} \left(\boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}\right) \right\| _ {2} \cdot \beta \cdot \left\| \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\alpha_ {1}\right) - \boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} \left(\alpha_ {2}\right) \right\| _ {2} (Definition B.6) \\ \leq \left\| \boldsymbol {A} ^ {\top} \operatorname {d i a g} \left(\boldsymbol {h} _ {k} ^ {S _ {i}}\right) \right\| _ {2} \cdot \beta \cdot H \cdot \left\| \boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2} \right\| _ {2} (Lemma B.8) \\ \leq H ^ {2} \cdot \beta \cdot \| \boldsymbol {\alpha} _ {1} - \boldsymbol {\alpha} _ {2} \| _ {2} \\ \end{array}
$$
B.2.3 Convergence of ATP under FedAvg framework
Finally, we show that with preservation of convexity and smoothness, ATP shares the same convergence guarantee as FedAvg [31]. We apply the proof in [46].
Theorem B.11 (Convergence of ATP). Assume that
1. At any round  $t$ , each client takes  $\tau$  SGD steps with learning rate  $\eta$ .
2. Full participation, i.e., each source client participates every round
3.  $\ell_{CE}(f(\pmb{x};\pmb{w}),\pmb{y})$  is convex and  $\beta$ -smooth w.r.t.  $\pmb{w}$  given any data sample  $(\pmb{x},\pmb{y})$ .
4. The update direction  $\pmb{h}_k^{\mathcal{S}_i}$  is  $H$ -module-wise-bounded for all  $i, j$
5. Bounded inner variance: for any  $\alpha$  and client  $i$ ,
$$
\mathbb {E} _ {j} \nabla_ {\boldsymbol {\alpha}} F _ {i j} (\boldsymbol {\alpha}) = \nabla_ {\boldsymbol {\alpha}} F _ {i} (\boldsymbol {\alpha}), \quad \mathbb {E} _ {j} \| \nabla_ {\boldsymbol {\alpha}} F _ {i j} (\boldsymbol {\alpha}) - \nabla_ {\boldsymbol {\alpha}} F _ {i} (\boldsymbol {\alpha}) \| _ {2} ^ {2} \leq \sigma^ {2}
$$
6. Bounded outer variance: for any  $\alpha$  and client  $i$ ,
$$
\left\| \nabla_ {\boldsymbol {\alpha}} F _ {i} (\boldsymbol {\alpha}) - \nabla_ {\boldsymbol {\alpha}} F (\boldsymbol {\alpha}) \right\| _ {2} ^ {2} \leq \zeta^ {2}
$$
If the client learning rate satisfies  $\eta \geq \frac{1}{4H^2\beta}$ , then one has
$$
\mathbb {E} \left[ \frac {1}{\tau T} \sum_ {t = 0} ^ {T - 1} \sum_ {k = 1} ^ {\tau} F (\bar {\boldsymbol {\alpha}} ^ {t, k}) - F (\boldsymbol {\alpha} ^ {*}) \right] \leq \frac {\| \boldsymbol {\alpha} _ {G} ^ {0} - \boldsymbol {\alpha} ^ {*} \| _ {2} ^ {2}}{2 \eta \tau T} + \frac {\eta \sigma^ {2}}{N} + 4 \tau \eta^ {2} H ^ {2} \beta \sigma^ {2} + 1 8 \tau^ {2} \eta^ {2} H ^ {2} \beta \zeta^ {2}
$$
where  $\alpha^{*} = \arg \min_{\alpha} F(\alpha)$  and  $\bar{\alpha}^{t,k} = \frac{1}{N}\sum_{i=1}^{N}\alpha_{i}^{t,k}$ .  $\alpha_{i}^{t,k}$  is the local adaptation rates after the communication rounds and  $k$  local epochs.
Proof. The optimization process of ATP is similar as FedAvg [31], where the difference is that ATP adapts the adaptation rates instead of model parameter. Lemma B.5 and B.10 that ATP preserves convexity and smoothness, i.e., for each client  $i$ ,  $F_{i}(\alpha)$  is convex and  $(H^{2}\beta)$ -smoothness. Therefore, we can apply Theorem 1 in [46] to complete the proof.
Remark B.12. The convergence rate of ATP is  $\mathcal{O}\left(\frac{1}{\tau T}\right)$ .
B.3 Generalization analysis
In this part, we studied how an adaptation rate  $\alpha$  learned by ATP that performs well on source clients can generalize to target clients. More specifically, we are interested in how many different source clients are required to ensure a certain generalization error. Similar to most of the other generalization analysis, we (1) derive generalization bound for any fixed hypothesis  $(\alpha)$ , and (2) quantify the size of hypothesis space.
B.3.1 Definitions: data generation and error rates
We first formally define the error rates.
Definition B.13 (Error rate for one data sample). Let  $f(\cdot; \boldsymbol{w}_k^{S_i}) : \mathcal{X} \to \Delta^{|\mathcal{Y}| - 1}$  be the neural network with model parameters  $\boldsymbol{w}_k^{S_i}$  that takes one data sample  $\boldsymbol{x}_{k,b}^{S_i}$  as input and outputs a probability distribution over the label space, i.e.,  $f(\boldsymbol{x}_{k,b}^{S_i}; \boldsymbol{w}_k^{S_i}) \geq \mathbf{0}$  and  $\mathbf{1}^\top f(\boldsymbol{x}_{k,b}^{S_i}; \boldsymbol{w}_k^{S_i}) = 1$ . Given adapted model parameters  $\boldsymbol{w}_k^{S_i}$ , define the error rate on one data sample  $(\boldsymbol{x}_{k,b}^{S_i}, \boldsymbol{y}_{k,b}^{S_i})$  to be
$$
\hat {e} _ {i k b} \left(\boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}}\right) := 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}}\right)
$$
Remark B.14. Definition B.13 is equivalent to the expected misclassification rate if when making random decision based on the output probability  $f(\pmb{x}_{k,b}^{\mathcal{S}_i};\pmb{w}_k^{\mathcal{S}_i})$ .
Definition B.15 (Error rate for one data batch). Given global model parameter  $\boldsymbol{w}_{G}$ , adaptation rate  $\alpha$ , and a batch of data  $\boldsymbol{X}_k^{\mathcal{S}_i} = \{\boldsymbol{x}_{k,b}^{\mathcal{S}_i}\}_{b=1}^B$ ,  $\boldsymbol{Y}_k^{\mathcal{S}_i} = \{\boldsymbol{y}_{k,b}^{\mathcal{S}_i}\}_{b=1}^B$ , define the error rate for one data batch
$$
\hat {\varepsilon} _ {i k} (\boldsymbol {\alpha}) := \hat {e} _ {i k} \left(\boldsymbol {w} _ {k} ^ {S _ {i}}\right) := \frac {1}{B} \sum_ {b = 1} ^ {B} \hat {e} _ {i k b} \left(\boldsymbol {w} _ {k} ^ {S _ {i}}\right)
$$
$$
\boldsymbol {w} _ {k} ^ {\mathcal {S} _ {i}} = \boldsymbol {w} _ {G} + (\boldsymbol {A} \boldsymbol {\alpha}) \odot \boldsymbol {h} _ {k} ^ {\mathcal {S} _ {i}}
$$
and  $h_k^{S_i}$  is the update direction computed with  $X_k^{S_i}$ .
Definition B.16 (Error rates for one client). Given global model parameter  $\mathbf{w}_G$ , adaptation rate  $\alpha$  and a source client  $S_i$  with  $K$  data batches  $\{(X_k^S, Y_k^S)\}_{k=1}^K$ , define the empirical error rate for source client  $S_i$
$$
\hat {\varepsilon} _ {i} (\boldsymbol {\alpha}) := \frac {1}{K} \sum_ {k = 1} ^ {K} \hat {\varepsilon} _ {i k} (\boldsymbol {\alpha})
$$
Also, define the expected error rate for source client  $S_{i}$
$$
\varepsilon_ {i} (\boldsymbol {\alpha}) := \mathbb {E} _ {(\boldsymbol {X} _ {k} ^ {S _ {i}}, \boldsymbol {Y} _ {k} ^ {S _ {i}}) \sim P S _ {i}} \left[ \hat {\varepsilon} _ {i k} (\boldsymbol {\alpha}) \right]
$$
Remark B.17.  $\hat{\varepsilon}_i(\pmb {\alpha})$  quantifies the error rate on client  $S_{i}$  's finite dataset  $\mathbb{D}^{S_i} = \{(X_k^{S_i},Y_k^{S_i})\}_{k = 1}^K.$ $\varepsilon_{i}(\pmb {\alpha})$  quantifies the expected error rate on a new data batch from client  $S_{i}$ . Notice that the same definition applies to target clients.
Definition B.18 (Source error rate and expected target error rate). Given global model parameter  $\pmb{w}_{G}$ , adaptation rate  $\alpha$ , and  $N$  source client  $S_{1}, \dots, S_{N}$ , each with  $K$  data batches  $\{(X_k^{\mathcal{S}_i}, Y_k^{\mathcal{S}_i})\}_{k=1}^{K}$ , define the training error rate
$$
\hat {\varepsilon} (\boldsymbol {\alpha}) := \frac {1}{K} \sum_ {i = 1} ^ {N} \hat {\varepsilon} _ {i} (\boldsymbol {\alpha})
$$
Also, define the expected testing error rate