Neural Prior Estimation: Learning Class Priors from Latent Representations

Masoud Yavari; Payman Moallem

cs.LGFeb 19, 2026

Neural Prior Estimation: Learning Class Priors from Latent Representations

Authors: Masoud Yavari, Payman Moallem

Organizations: Independent Researcher Isfahan, Iran · Department of Electrical Engineering Faculty of Engineering University of Isfahan, Iran

Abstract

Logit adjustment corrects class imbalance using the empirical class prior. We study whether a comparable class-frequency signal can instead be learned from the network representation, without explicitly supplying class counts to the correction rule. We introduce the Neural Prior Estimator (NPE), which attaches one or more lightweight Prior Estimation Modules (PEMs) to the latent representation. Each PEM is trained with a one-way logistic objective on the ground-truth coordinate. The resulting frequency-dependent outputs are combined into an NPE estimate and used as a learned logit correction, forming NPE-LA. In a simplified scalar model, the optimum of the PEM objective is monotone in the class count and grows asymptotically as $\log N_c$ , up to a slower $\log \log N_c$ term. Experiments on long-tailed CIFAR-10 and CIFAR-100 show that NPE-LA is competitive with standard logit adjustment and improves minority-class performance over CE and cRT in the reported settings. Experiments on STARE and ADE20K further show that the same idea can be used as a lightweight recalibration mechanism for dense prediction.

Figures & tables

Setting	LR	Weight Decay	Batch Size	Epochs	LR decay milestones
HP-1	0.1	$2\times 10^{-4}$	124	200	160, 180
HP-2	0.05	$10^{-3}$	64	120	100, 110

Table 1: Training hyperparameters for classification experiments. LR decay factor = 0.1 at each milestone.

Dataset	Backbone + Head	Optimizer	LR	WD	Iterations	LR decay milestones
STARE	UNet-FCN	SGD	1e-2	5e-4	600	200, 400
ADE20K	DeepLab-V3	SGD	1e-2	5e-4	2 500	1 250
ADE20K	Swin-T + UPerNet	AdamW	6e-5	1e-2	2 500	1 250

Table 2: Per-model training schedule and optimizer settings. LR decay factor = 0.1 at each milestone.

Figure 1: Class-wise accuracy on CIFAR-100 (

\rho=100

) under HP-2, illustrating the impact of different numbers of PEMs (

N_{\mathrm{PEM}}

) across Head, Medium, and Tail classes. NPE is utilized during training only; no logit adjustment is performed at inference time.

# of PEMs	CIFAR-100			CIFAR-10
# of PEMs	200	100	50	200	100	50
0	37.62	40.53	45.98	64.40	75.01	77.00
1	37.28	41.00	45.96	65.24	73.05	78.29
4	37.09	42.35	45.45	66.31	73.95	77.99
8	38.39	41.65	46.29	66.13	73.72	78.98
16	38.42	41.78	46.36	66.22	74.20	78.33

Table 3: Top-1 accuracy (%) on CIFAR-100 and CIFAR-10 for different imbalance ratios

\rho

and PEMs under hyperparameter setting HP-2. NPE is utilized during training only; no logit adjustment is performed at inference time.

Method	CIFAR-100			CIFAR-10
Method	200	100	50	200	100	50
CE	34.83 $\pm$ 0.26	39.14 $\pm$ 0.41	43.87 $\pm$ 0.25	63.54 $\pm$ 0.28	71.24 $\pm$ 0.29	76.77 $\pm$ 0.37
cRT	37.04 $\pm$ 0.24	41.04 $\pm$ 0.14	45.83 $\pm$ 0.32	67.19 $\pm$ 0.31	73.80 $\pm$ 0.42	79.07 $\pm$ 0.51
LA	38.59 $\pm$ 0.18	42.63 $\pm$ 0.46	47.07 $\pm$ 0.49	72.10 $\pm$ 0.27	77.37 $\pm$ 0.32	81.25 $\pm$ 0.38
NPE-LA (1 PEM)	38.85 $\pm$ 0.30	42.53 $\pm$ 0.40	46.86 $\pm$ 0.78	72.42 $\pm$ 0.21	77.89 $\pm$ 0.38	82.22 $\pm$ 0.46
NPE-LA (16 PEMs)	39.36 $\pm$ 0.12	42.73 $\pm$ 0.28	46.44 $\pm$ 0.24	73.18 $\pm$ 0.41	78.17 $\pm$ 0.73	82.31 $\pm$ 0.22

Table 4: Top-1 accuracy (%) on CIFAR-100 and CIFAR-10 for different imbalance ratios

\rho

and methods under hyperparameter setting HP-1. Bold indicates the highest value; underlined indicates the second highest

Method	CIFAR-100			CIFAR-10
Method	200	100	50	200	100	50
CE	37.04 $\pm$ 0.23	41.00 $\pm$ 0.12	45.46 $\pm$ 0.11	65.59 $\pm$ 0.20	73.18 $\pm$ 0.77	77.76 $\pm$ 0.44
cRT	41.58 $\pm$ 0.24	45.43 $\pm$ 0.11	50.50 $\pm$ 0.12	72.34 $\pm$ 0.71	77.85 $\pm$ 0.25	81.01 $\pm$ 0.50
LA	43.35 $\pm$ 0.33	46.85 $\pm$ 0.53	51.30 $\pm$ 0.22	74.91 $\pm$ 0.51	80.13 $\pm$ 0.13	82.95 $\pm$ 0.36
NPE-LA (1 PEM)	42.99 $\pm$ 0.30	46.72 $\pm$ 0.23	51.25 $\pm$ 0.36	74.37 $\pm$ 0.25	79.35 $\pm$ 0.37	82.32 $\pm$ 0.19
NPE-LA (16 PEMs)	43.96 $\pm$ 0.40	47.47 $\pm$ 0.36	51.82 $\pm$ 0.17	74.83 $\pm$ 0.55	80.12 $\pm$ 0.31	83.16 $\pm$ 0.08

Table 5: Top-1 accuracy (%) on CIFAR-100 and CIFAR-10 for different imbalance ratios

\rho

and methods under hyperparameter setting HP-2. Bold indicates the highest value; underlined indicates the second highest

Figure 2: Class-wise accuracy on CIFAR-100 (

\rho=100

) under HP-2, comparing different methods across head, medium, and tail classes.

Backbone	Decode Head	PEM	Scale	Dice (%)			Accuracy (%)
		Type	Factor	BG	FG	Mean	BG	FG	Mean
UNet	FCN	-	-	98.55	81.20	89.78	99.13	75.21	87.17
		FCN	1	97.34	75.02	86.27	95.22	94.93	95.08
		FCN	0.2	98.52	82.39	90.46	98.44	83.21	90.83

Table 6: STARE segmentation with UNet–FCN: Dice and pixel accuracy for background (BG), foreground (FG), and mean, comparing the baseline and NPE-LA across scale factors

\alpha

Backbone	Decode Head	PEM Type	Scale Factor	mIoU (%)	mAcc (%)
DeepLab-V3	ASPPHead	-	-	42.42	53.60
		ASPPHead	1	24.55	65.59
		ASPPHead	0.1	42.82	57.99
		FCN	1	29.28	67.05
		FCN	0.1	42.92	56.92
Swin-T	UPerNet	-	-	44.41	55.96

Table 7: ADE20K segmentation with DeepLabv3 and Swin-T: mIoU and pixel accuracy for baseline and NPE-LA with different PEM types and scale factors

\alpha

Figure 3: Class-wise IoU on ADE20K for baseline DeepLab-V3 and NPE-LA variant with a single FCN PEM and scale factor 0.1. Baseline IoUs are ordered in ascending order for improved visual clarity.

Figure 4: Class-wise accuracy on ADE20K for baseline DeepLab-V3 and NPE-LA variant using a single FCN PEM and scale factor 0.1. Baseline accuracies are ordered in ascending order for improved visual clarity.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 5: NPE estimated prior

\hat{p}^{\mathrm{val}}(y)

and empirical class prior

p(y)

(left axis), and representation confidence

\mathrm{Conf}(y)

(right axis), on CIFAR-100-LT (

\rho=100

Logit Adjustment

See all 8 papers

Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification
May 20, 2026Samuel McDowell, Nathan Stromberg, Lalitha SankarClass-Imbalanced Learning Tabular Classification
School of Electrical, Computer and Energy Engineering · Arizona State University, Tempe, AZ
Selective Posterior Margin Regularization for Forward-Corrected Classification
Sep 5, 2026Zexing Zhang, Jichao Li, Tianyang Lei +2Learning with Noisy Labels
College of Systems Engineering, National University of Defense Technology Changsha 410073, China
Cross-Fitting for Neural Posterior Estimation
Oct 19, 2026Jeffrey RegierNeural Posterior Estimation Simulation-Based Inference
Department of Statistics, University of Michigan, Ann Arbor, MI 48109

May 20, 2026J/K move · Enter open · S save

Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification

Samuel McDowell, Nathan Stromberg, Lalitha Sankar

School of Electrical, Computer and Energy Engineering · Arizona State University, Tempe, AZ

Class-Imbalanced Learning Tabular Classification Prior-Data Fitted Networks

Open paper AskPDF arXiv

Prior-data fitted networks (PFNs) have achieved exceptional performance on tabular classification tasks. However, like other classifiers, their performance can suffer under the effect of class imbalance, resulting in poor performance for rare classes. Several techniques exist which attempt to mitigate the deleterious effect of class imbalance on classification performance, but the in-context learning (ICL) dynamic of PFNs means that loss-based strategies are impossible, and other techniques are unproven. We have adapted several classical techniques addressing class imbalance and analyzed their performance on PFN classification. We observe that thresholding performs exceptionally well because of the calibration characteristics of PFNs, and downsampling performs comparably because of PFNs exceptional limited-data performance, with the additional benefit of reduced computation cost for inference.

Abstract

Logit Adjustment

Explore similar work

Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification

Selective Posterior Margin Regularization for Forward-Corrected Classification

Cross-Fitting for Neural Posterior Estimation