SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Supervised refinement After unsupervised adaptation, we refine the adaptation rates on each source client with label information to minimize  $\ell_{CE}(f(\pmb{X}_k,\pmb{w}_k),\pmb{Y}_k)$ , where  $\ell_{CE}$  is the cross-entropy loss. We use gradient descent to optimize  $\alpha$ , i.e.,
$$
\boldsymbol {\alpha} \leftarrow \boldsymbol {\alpha} - \eta \nabla_ {\boldsymbol {\alpha}} \ell_ {C E} (f (\boldsymbol {X} _ {k}; \boldsymbol {w} _ {k}), \boldsymbol {Y} _ {k}) \tag {7}
$$
where  $\eta$  is the learning rate of adaptation rates. Notice that the gradient of  $\alpha$  can be computed as
$$
\nabla_ {\boldsymbol {\alpha}} \ell_ {C E} (f (\boldsymbol {X} _ {k}; \boldsymbol {w} _ {k}), \boldsymbol {Y} _ {k}) = \frac {\partial \ell_ {C E} (f (\boldsymbol {X} _ {k} ; \boldsymbol {w} _ {k}) , \boldsymbol {Y} _ {k})}{\partial \boldsymbol {w} _ {k}} \frac {\partial \boldsymbol {w} _ {k}}{\partial \boldsymbol {\alpha}} = \boldsymbol {A} ^ {\top} (\boldsymbol {h} _ {k} \odot \nabla_ {\boldsymbol {w} _ {k}} \ell_ {C E} (f (\boldsymbol {X} _ {k}; \boldsymbol {w} _ {k}), \boldsymbol {Y} _ {k}))
$$
To estimate the gradient of  $\alpha$ , each training client only needs to adjacently compute the unsupervised and supervised gradient, and compute their module-wise inner products. Different from many meta-learning algorithms [9, 26], ATP is computationally very efficient since it requires no second-order derivatives. In the practical implementation, since each module in the model has significantly different number of parameters, the raw gradient for each  $\alpha^{[l]}$  usually has different scales. Therefore we normalize the gradient with the square root of the number of parameters in the corresponding module.
Server aggregation To incorporate adaptation knowledge from multiple source clients and enhance generalization to the clients' population, ATP use standard federated aggregation [31] to periodically aggregates the local adaptation rates. In each communication rounds, after each client locally update  $\alpha$  for a few iterations, the local adaptation rates are uploaded to the server for averaging (as shown in line 6 of Algorithm 1), and then sent to source clients for the next round of training. With server aggregation, ATP learn the adaptation rates that enables successful adaptation to all source clients in average.
Communication cost Notice that ATP only optimizes the adaptation rates  $\alpha$  without changing the global model  $w_{G}$ . Therefore, only the adaptation rates are kept transmitted between the server and each client, while the global model parameter is only broadcasted once at the start of the ATP training. Such design significantly reduces the communication cost from 2TD (for standard FedAvg) to  $D + 2Td$ .
<sup>1</sup> Some literatures consider  $(1 - m)$  as the momentum. Here we follow the definition in PyTorch.
Algorithm 1 ATP Training
ServerTrain  $(\pmb{w}_G,\pmb{\alpha}_G^0 = \mathbf{0})$    
1: Broadcast  $\pmb{w}_{G}$  to all source clients   
2: for communication round  $t = 1$  to  $T$  do   
3:  $\mathbb{S}^t\gets$  (random set of  $C$  source clients)   
4: for source client  $S_{i}\in \mathbb{S}^{t}$  in parallel do   
5:  $\pmb{\alpha}_{i}^{t}\gets$  ClientTrain  $(S_{i},\alpha_{G}^{t - 1})$    
6:  $\pmb{\alpha}_{G}^{t} = \frac{1}{C}\sum_{S_{i}\in \mathbb{S}^{t}}\pmb{\alpha}_{i}^{t}$    
7: return  $\pmb{\alpha}_{G}^{T}$    
ClientTrain  $(S_i,\alpha)$  # Run on source client  $S_{i}$    
8: for local epoch  $e = 1$  to  $E$  do   
9:  $\mathbb{B}^{S_i}\gets$  (split  $\mathbb{D}^{S_i}$  into  $K^{S_i}$  batches of size  $B$    
10: for batch  $k = 1$  to  $K^{S_i}$  do   
11:  $(X_k^{S_i},Y_k^{S_i})\gets$  (k-th labeled batch in  $\mathbb{B}^{S_i}$ )   
12: Estimate update direction  $\pmb{h}_k^{S_i}$  with unlabeled  $X_{k}^{S_{i}}$  according to Eq. (4) and (5)   
13:  $\pmb{w}_k^{S_i}\gets \pmb {w}_G + (\pmb {A}\pmb {\alpha})\odot \pmb{h}_k^{S_i}$    
14:  $\pmb{\alpha}\gets \pmb{\alpha}-\eta\nabla_{\pmb{\alpha}}\ell_{CE}(f(\pmb{X}_j^{S_i};\pmb{w}_k^{S_i}),\pmb{Y}_k^{S_i})$    
15: return  $\pmb{\alpha}$
Algorithm 2 ATP Testing
ClientTest  $(\mathcal{T}_j,\boldsymbol{w}_G,\alpha)$  # Run on target client  $\mathcal{T}_j$   
1:  $\mathbb{B}^{\mathcal{T}_j}\gets$  (split  $\mathbb{X}^{\mathcal{T}_j}$  into  $K^{\mathcal{T}_j}$  batches of size  $B$   
2:  $\pmb{h}_{\mathrm{history}}\gets 0$  # Cumulative moving average  
3: for batch  $k = 1$  to  $K^{\mathcal{T}_j}$  do  
4: Estimate update direction  $\pmb{h}_k^{\mathcal{T}_j}$  with unlabeled  $\pmb{X}_k^{\mathcal{T}_j}$  according to Eq. (4) and (5)  
5: if TTBA then  
6:  $\pmb{w}_k^{\mathcal{T}_j}\gets \pmb{w}_G + (\pmb {A}\pmb {\alpha})\odot \pmb{h}_k^{\mathcal{T}_j}$   
7: else if OTTA then  
8:  $\pmb{h}_{\mathrm{history}}\gets \frac{k - 1}{k}\pmb{h}_{\mathrm{history}} + \frac{1}{k}\pmb{h}_k^{\mathcal{T}_j}$   
9:  $\pmb{w}_k^{\mathcal{T}_j}\gets \pmb{w}_G + (\pmb {A}\pmb {\alpha})\odot \pmb{h}_{\mathrm{history}}$   
10: Make prediction:  $\hat{\pmb{Y}}_k^{\mathcal{T}_j} = f(\pmb {X}_k^{\mathcal{T}_j};\pmb {w}_k^{\mathcal{T}_j})$
4.2 Testing phase: exploit adaptation rates on target clients
During testing, each target client downloads both the global model and the adaptation rates. We propose two versions of ATP: ATP-batch for test-time batch adaptation (TTBA) and ATP-online for online test-time adaptation (OTTA). We summarize the testing phase in Algorithm 2.
ATP-batch For TTBA, each target client makes independent predictions on each batch. For each batch of target data, ATP-batch first conducts the unsupervised adaptation identical to source clients, and then makes prediction.
ATP-online For OTTA, data comes in a stream of batches  $\left[\boldsymbol{X}_1^{\mathcal{T}_j},\boldsymbol{X}_2^{\mathcal{T}_j},\dots \right]$ . Previous works [43, 45] usually keep updating the model batch after batch. However, such accumulative adaptation can introduce severe batch dependency problem, i.e., each batch is evaluated when the model takes different number of update steps [54]. For the first few batches, the model has not adapted to the local distribution well; while for the last few batches, the model may over-minimize the entropy but increase the cross-entropy loss. To avoid batch dependency, we propose an averaged adaptation mechanism for online adaptation, whose scale of adaptation is stable during online adaptation.
For each batch  $\boldsymbol{X}_k^{\mathcal{T}_j}$  in the data stream, we always compute the update direction  $\boldsymbol{h}_k^{\mathcal{T}_j}$  starting with the fixed global model  $\boldsymbol{w}_G$  according to Eq. (4) and (5). Subsequently, instead of using only the current update direction to adapt the model, we average all the stored update direction to update the model, i.e.,
$$
\boldsymbol {w} _ {k} ^ {\mathcal {T} _ {j}} \leftarrow \boldsymbol {w} _ {G} + (\boldsymbol {A} \boldsymbol {\alpha}) \odot \left(\frac {1}{k} \sum_ {s = 1} ^ {k} \boldsymbol {h} _ {s} ^ {\mathcal {T} _ {j}}\right) \tag {8}
$$
By using the average of previous updates, we simulate updating with larger batch size to utilize historical data, while controlling the number of update steps to be one. In the practical implementation, we use cumulative moving average (as shown in line 8 of Algorithm 2), whose space complexity does not increase with the increment of step  $k$ .
5 Theoretical analysis
In this section, we show that ATP enjoys good generalization guarantees because of the low dimensionality of adaptation rates. Formal definitions, assumptions and full proofs are provided in Appendix B.3. We also show in Appendix B.2 that ATP has convergence guarantee similar to FedAvg [31, 46].
Theorem 5.1 (Generalization). 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$ -upper-bounded 2-norms for each module's update. For any fixed global model  $\mathbf{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 {9}
$$
where  $\hat{\varepsilon} (\alpha)$  is the average post-adaptation error rate on source clients, and  $\varepsilon (\alpha)$  is the expected post-adaptation error rate on clients' population.
Theorem 5.1 shows that, although ATP improves the model expressiveness by adapting the model to each client's distribution, ATP can still provably generalize well to the clients' population. Especially, this generalization benefit from low dimensionality of adaptation rates, since the bound gets looser when  $d$  increases. Moreover, this bound shows the importance of learning adaptation rates from multiple source clients: if we merge all  $N$  source domains with  $K$  batches into one domain with  $NK$  batches, then the bound will be much looser.
In this section, we design experiments to answer the following research questions:
- RQ1: Can ATP handle different distribution shift and outperform prior TTA methods?
- RQ2: Does ATP learn adaptation rates specific to distribution shift?
Setup We evaluate ATP on a variety of models, datasets and distribution shifts. We first evaluate on CIFAR-10(-C) with a standard three-way split [50]: we randomly split the dataset to 300 clients: 240 source clients and 60 target clients. Each source client has 160 training samples and 40 validation samples, while each target client has 200 unlabeled 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, and 4 new kinds of corruptions to the target clients to test the generalization of ATP. 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 perturbations. To test ATP under more challenging domain shifts, we then evaluate ATP on two domain generalization datasets: Digits-5 [25] and PACS [21]. We adopt the leave-one-domain-out evaluation protocol [10], i.e., one domain is chosen to construct target clients, and the remaining domains are used to construct source clients. We follow similar data preprocessing in [25], while additionally applying step partition to inject label shift. Each domain is divided into 10 clients, leading to 40/10 source/target clients for Digits-5 and 30/10 source/target clients for PACS. For the experiments above, we use ResNet-18 [11] as a common choice in FL experiments [42, 14, 33]. We also test ATP with two different architectures: a five-layer CNN on CIFAR-10(-C) and ResNet-50 on CIFAR-100(-C). Detailed experiment settings are given in Appendix C.1.
6.1 RQ1: Can ATP handle different distribution shift?
We compare ATP with three kinds of baseline TTA methods. For feature shift methods, we compare to BN-Adapt [38] and Tent [45] which adjusts the batch normalization layers, SHOT [28] which adjusts the feature extractor, T3A [16] which adjusts the final classifier, and MEMO [52] which uses augmentation to adjust the whole network. For label shift, we compare to EM [36] which adjusts the label priori unsupervisedly with expectation-maximization, and BBSE [29] which uses the validation data to construct a confusion matrix to estimate the label priori. Since re-training a model with different label weights for each
Table 1: Accuracy (mean ± s.d. %) on target clients under various distribution shifts on CIFAR-10
Method	Feature shift	Label shift	Hybrid shift	Avg. Rank
No adaptation	69.42 ± 0.13	72.98 ± 0.24	63.68 ± 0.24	7.7
BN-Adapt	73.52 ± 0.22	54.54 ± 0.10	50.42 ± 0.39	7.0
SHOT	71.76 ± 0.17	48.13 ± 0.18	44.68 ± 0.32	9.3
Tent	71.76 ± 0.09	50.13 ± 0.21	46.05 ± 0.26	8.3
T3A	69.53 ± 0.08	71.70 ± 0.32	62.17 ± 0.17	8.0
MEMO	72.43 ± 0.22	77.30 ± 0.15	68.07 ± 0.28	4.3
EM	65.18 ± 0.12	80.73 ± 0.18	69.85 ± 0.43	5.0
BBSE	63.98 ± 0.17	79.30 ± 0.17	67.96 ± 0.43	6.7
Surgical	69.85 ± 0.22	76.00 ± 0.17	66.94 ± 0.43	6.3
ATP-batch	73.68 ± 0.10	79.90 ± 0.22	73.05 ± 0.35	2.3
ATP-online	74.06 ± 0.18	81.96 ± 0.14	75.37 ± 0.22	1.0
client is not realistic in FL. We use the estimated label distribution to adjust the output of a classifier.
Table 2: Accuracy (mean  $\pm$  s.d. %) on target clients under hybrid shift on Digits-5 and PACS
Method	Digits-5	PACS
MNIST	SVHN	USPS	SynthDigits	MNIST-M	Art	Cartoon	Photo	Sketch
No adaptation	95.47 ± 0.22	52.28 ± 1.45	89.62 ± 0.44	79.75 ± 0.69	55.62 ± 0.80	71.57 ± 1.16	74.71 ± 0.70	90.25 ± 0.75	74.20 ± 0.72
BN-Adapt	94.90 ± 0.29	57.57 ± 0.53	89.51 ± 0.39	75.34 ± 0.48	59.68 ± 0.44	73.55 ± 0.51	71.54 ± 0.55	92.07 ± 0.26	70.92 ± 0.53
SHOT	94.69 ± 0.31	57.91 ± 0.23	89.55 ± 0.69	76.43 ± 0.34	60.19 ± 0.69	69.32 ± 0.67	67.77 ± 0.40	86.97 ± 0.60	59.40 ± 0.91
Tent	95.48 ± 0.29	60.67 ± 0.49	91.65 ± 0.61	78.56 ± 0.45	62.49 ± 0.73	71.59 ± 0.71	71.03 ± 0.97	88.06 ± 0.24	63.15 ± 1.10
T3A	94.63 ± 0.61	49.90 ± 1.10	88.46 ± 0.75	75.47 ± 1.14	51.25 ± 1.55	72.15 ± 0.72	75.02 ± 0.78	91.51 ± 0.62	70.14 ± 1.21
MEMO	95.92 ± 0.19	52.85 ± 1.09	89.84 ± 0.44	80.12 ± 0.90	55.48 ± 1.13	71.47 ± 1.29	75.57 ± 0.98	90.65 ± 0.90	76.30 ± 0.65
EM	96.64 ± 0.31	57.21 ± 1.65	92.29 ± 0.32	85.69 ± 0.46	62.08 ± 0.60	73.96 ± 1.85	78.91 ± 0.92	92.30 ± 0.92	80.82 ± 1.52
BBSE	94.47 ± 0.58	57.26 ± 1.47	91.34 ± 0.39	85.54 ± 0.46	61.59 ± 0.91	74.33 ± 1.78	78.69 ± 1.00	91.82 ± 0.68	80.15 ± 1.42
Surgical	97.35 ± 0.13	59.93 ± 2.01	94.19 ± 0.40	86.06 ± 0.44	65.87 ± 0.78	74.59 ± 2.69	77.48 ± 0.64	92.34 ± 0.78	80.90 ± 3.42
ATP-batch	97.81 ± 0.27	62.18 ± 1.71	95.41 ± 0.26	87.91 ± 0.45	69.98 ± 1.96	82.92 ± 0.96	79.64 ± 0.75	95.40 ± 0.41	82.28 ± 1.57
ATP-online	97.81 ± 0.23	62.64 ± 1.92	95.56 ± 0.23	88.33 ± 0.47	70.78 ± 2.36	83.51 ± 0.84	79.46 ± 0.77	95.52 ± 0.40	82.80 ± 1.69
Figure 4: Adaptation rates learned by ATP with different distribution shifts on CIFAR-10
We also compare to Surgical [20] which uses the validation data to decide which blocks to adapt. For all baselines we use the validation data to select hyperparameters.
ATP can handle different types of distribution shifts Table 1 shows the results on CIFAR-10. Under feature and label shifts, most TTA methods suffer from performance trade-off as they improve the performance on one distribution shift while harm the other. The only exception is MEMO, which utilizes data augmentation to robustify the model prediction. However, it also introduces significant computational cost during inference. As an adaptive framework simpler than ours, Surgical also introduces accuracy gain across all distribution shifts. However, its coarse-grained adaptation rule prevents further improvement on the accuracy. ATP reaches great performance comparable to the strongest baseline TTA method under both feature and label shifted. Under the more complex hybrid shift, ATP achieves the highest performance gain with a significant margin. Meanwhile, ATP-online can further improve the performance of ATP-batch by using information from previous batches.
ATP can handle more challenging domain shifts Table 2 shows the results on two domain generalization datasets with a hybrid of domain and label shifts. Compared to baselines, ATP consistently achieves higher accuracy across all domains.
ATP is compatible to multiple model architectures. Finally, we evaluate ATP on more model architectures: Shallow-CNN as smaller model and ResNet-50 as larger model. As shown in Table 5 in Appendix C.2, ATP has uniformly good performance on two new models.
6.2 RQ2: Does ATP learn adaptation rates specific to distribution shift?
Besides ATP's good performance, we are also interested in whether ATP successfully learns adaptation rates specific to the type distribution shift. To explore this, we group the adaptation rates by their corresponding block and module type under three kinds of distribution shifts. As shown in Figure 4, ATP learns significantly different adaptation rates under different distribution shifts. In Figure 4 (left), ATP learns to adapt the last linear layer under label shift, while mainly adapt the former layers under feature shift. More interestingly, we notice in Figure 4 (right) that the adaptation rates for batch norm running statistics are positive under feature shift, but negative under label shift. Negative adaptation rate is usually counter-intuitive, since it disaligns the training and testing distributions. However, it benefits the model under label shifts because it explicitly adapts the label prior distribution towards
Figure 5: Effect of cohort size and batch size
the prediction distribution. We use a toy example in Appendix C.5 to show why negative adaptation rate can improve the model performance under label distribution.
Moreover, we examine whether the learn adaptation rates are specific to the type of distribution shift by training on one distribution shift, but testing on another. We observe in Table 3 that, ATP performs the best when trained and tested with the same type of distribution shifts. However, the adaptation rates trained on feature/label shift fails to boost the performance on the other distribution shift. The adaptation rates trained on hybrid shift can generalize to feature shift and label shift, but still worse than the adaptation rates trained with the same type of distribution shifts. These results show that the learn adaptation rates are specific to the type of distribution shift.
Table 3: Train and test adaptation rates with different distribution shifts, accuracy (mean ± s.d. %)
Train	Test
Feature shift	Label shift	Hybrid shift
No adaptation	69.42 ± 0.13	72.98 ± 0.24	63.68 ± 0.24
Feature shift	73.68 ± 0.10	65.05 ± 1.82	60.64 ± 1.43
Label shift	67.99 ± 0.28	79.90 ± 0.22	69.50 ± 0.52
Hybrid shift	72.69 ± 0.14	78.92 ± 0.34	73.05 ± 0.35
6.3 Further discussion
Ablation study We present two variants of ATP to study how trainable parameters and running statistics contribute to the adaptability of ATP. ATP-params only learns to adapt the trainable parameters, while ATP-stats focuses solely on adapting the running statistics. As shown in Table 4, adapting trainable parameters and running statistics both play critical roles in achieving successful adaptation. More specifically, ATP-params primarily facilitate adaptation to label shift, whereas ATP-stats essentially aid in adapting to feature shift.
Table 4: Ablation study, accuracy (mean ± s.d. %)
Method	Feature shift	Label shift	Hybrid shift
No adaptation	69.42 ± 0.13	72.98 ± 0.24	63.68 ± 0.24
ATP-params	69.23 ± 0.27	78.29 ± 0.14	68.05 ± 0.54
ATP-stats	71.27 ± 0.17	74.03 ± 0.18	64.78 ± 0.27
ATP-batch	73.71 ± 0.14	79.90 ± 0.22	73.05 ± 0.35
Hyperparameter sensitivity Figure 5 shows the effects of cohort size and batch size with CIFAR-10 under the hybrid shift, where cohort size refers to the number of clients sampled at each round. ATP demonstrates remarkable consistency in accuracy across different cohort sizes, indicating its robustness. For batch size, we optimize the adaptation rates with  $B = 20$  and subsequently evaluate the algorithm with different batch sizes. We find that ATP consistently improves the model's accuracy across different batch sizes, with larger batch sizes yielding greater benefits for the model. ATP-online is more robust to batch size than ATP-batch since it can utilizes information from previous batches.
In this paper, we propose ATP that unsupervisedly learns the adaptation rate for each module to handle various types of distribution shifts encountered in test-time personalized federated learning. As a potential future direction, incorporating the training of the global model could offer advantages in terms of facilitating easier and better personalization.
This work is supported by National Science Foundation under Award No. IIS-1947203, IIS-2117902, IIS-2137468, IIS-2002540, Agriculture and Food Research Initiative (AFRI) grant no. 2020-67021-32799/project accession no.1024178 from the USDA National Institute of Food and Agriculture, the U.S. Department of Homeland Security under Grant Award Number, 17STQAC00001-06-00, and IBM-Illinois Discovery Accelerator Institute - a new model of an academic-industry partnership designed to increase access to technology education and skill development to spur breakthroughs in emerging areas of technology. The views and conclusions are those of the authors and should not be interpreted as representing the official policies of the funding agencies or the government.
A More discussions
A.1 More related works
Partial fine-tuning, i.e., updating a subset of modules of a pretrained network on a new dataset, has been studied in supervised settings [26, 35, 40]. In FL, PartialFed [42] adaptively decides whether each parameter is shared or personalized. However, it cannot generalize to testing clients that do not participate in the training. Recently, surgical fine-tuning [20] selectively fine-tunes a subset of blocks with a similar intuition that the type of distribution shift influences which part of the network to be adapted. Different from their method, we focus on the unsupervised setting and propose to refine the adaptation rate for each module.
Hyperparameter optimization is also related to our algorithm if considering adaptation rates as a set of hyperparameters. [19] first investigates the problem of federated hyperparameter tuning and proposed FedEX that leverages weight-sharing from neural architecture search to efficiently tune hyperparameters. [55] introduces FloRA that addresses use cases of tabular data and enables single-shot federated hyperparameter tuning. While these methods focus on improving the efficiency of hyperparameter optimization, our paper focuses on finding the optimal adaptation rates that benefit test-time personalization.
A.2 Broader impacts and limitations
Broader impacts We are not aware of any potential negative societal impacts regarding our work to the best of our knowledge. For all the used data sets, there is no private personally identifiable information or offensive content.
Limitations One possible limitation is that we consider a fix global model for lower communication cost and better generalization, while it might be beneficial to also train a global model for easier personalization, which could be a promising future direction.
B Theoretical analysis
In this section, we give theoretical proofs of convergence, and generalization of ATP.
B.1 Approximation analysis
In this subsection, we give detailed proofs of Proposition 3.1 and 3.2 in Section 3 of the main text. These propositions show why certain types of distribution shifts can be handled by adapting certain layers in a neural network.
B.1.1 Proof of Proposition 3.1
Proof. W.l.o.g., assuming the last layer of the neural network is a linear layer. Denoting  $g(\pmb{x}; \pmb{w}_g)$  as the input of the last layer, where  $\pmb{x}$  is the input and  $\pmb{w}_g$  is the model parameters for the feature extractor (i.e., all layers except for the last classification layer). Denote  $\pmb{w}_1, \dots, \pmb{w}_K$  as the weights of the last layer and  $b_1, \dots, b_K$  as the bias terms of the last layer, assuming  $K$  classes. Then we have
$$
f (\pmb {x}; \pmb {w}) _ {c} = \frac {\exp (\pmb {w} _ {c} ^ {\top} g (\pmb {x} ; \pmb {w} _ {g}) + b _ {c})}{\sum_ {c ^ {\prime} = 1} ^ {K} \exp (\pmb {w} _ {c ^ {\prime}} ^ {\top} g (\pmb {x} ; \pmb {w} _ {g}) + b _ {c ^ {\prime}})}
$$
Since the neural network is calibrated on  $p$ , for all class index  $c = 1, \dots, K$ , we have
$$
f (\boldsymbol {x}, \boldsymbol {w}) _ {c} = p (\boldsymbol {y} = \boldsymbol {e} _ {c} | \boldsymbol {x})
$$
where  $e_{c}$  is an one-hot vector with its  $c$ -th element as one. For distribution  $q$  with the same conditional distribution and different priors, by Bayes' theorem,  $\forall x, y$