Adjusting Label Predictions for Class Imbalance

Why Class Frequencies Affect Predictions

A label-transfer classifier is trained using reference cells with known cell types. If one type is common in that reference and another is rare, the classifier’s predicted probabilities can reflect that imbalance. A cell with the same expression evidence for either type may receive a higher probability for the common type because that type had a larger prior probability in the reference data.

This is not necessarily a flaw. If the reference frequencies represent the population being annotated, they may be appropriate. But reference datasets are often assembled for particular purposes: some types may be deliberately oversampled, rare populations may be enriched, or the reference may come from a different tissue composition. In those cases, its class frequencies need not represent the cells to which predictions will be applied.

The relevant adjustment is a change in prior class probabilities, not a correction for every kind of model error. It is justified only when the class-conditional distribution of the features is stable between reference and target cells and the classifier probabilities represent the reference posterior reasonably well.

Definition: Let \(Y\) be a cell-type label and \(X\) the features used by a classifier. Write \(\pi_k^{(R)}=P_R(Y=k)\) for the reference class prior of type \(k\), and \(\pi_k^{(T)}=P_T(Y=k)\) for its target prior. The empirical reference frequency is the observed fraction \(n_k/n\) of reference cells labelled \(k\); it estimates the reference prior only when the sampling design supports that interpretation. A chosen target prior is a specified or externally estimated distribution for the cells being annotated. These quantities need not be equal.

The Prior-Adjustment Formula

Suppose the classifier gives \(q_R(k\mid x)\), its estimated probability that a cell with features \(x\) has type \(k\) under the reference distribution. The adjustment changes the prior odds while retaining the relative evidence from the features. For each class, multiply the reference probability by the target-to-reference prior ratio, then normalize the resulting values to sum to one.

q_T(k\mid x)= \frac{q_R(k\mid x)\,\dfrac{\pi_k^{(T)}}{\pi_k^{(R)}}} {\displaystyle\sum_j q_R(j\mid x)\,\dfrac{\pi_j^{(T)}}{\pi_j^{(R)}}}

Here \(q_T(k\mid x)\) is the adjusted probability for type \(k\). The denominator is essential: it turns the reweighted values into probabilities that sum to one. In the binary case, the same operation multiplies the reference posterior odds by the target-to-reference prior odds ratio.

\frac{q_T(k\mid x)}{q_T(\ell\mid x)} = \frac{q_R(k\mid x)}{q_R(\ell\mid x)} \frac{\pi_k^{(T)}/\pi_\ell^{(T)}} {\pi_k^{(R)}/\pi_\ell^{(R)}}

The formula follows from Bayes’ rule. If the distribution of \(X\) conditional on \(Y=k\) is unchanged, then the target likelihood for class \(k\) is the same as the reference likelihood. The reference posterior is proportional to that likelihood times the reference prior. Replacing the reference prior with the target prior therefore multiplies the reference posterior by \(\pi_k^{(T)}/\pi_k^{(R)}\); normalizing across classes gives the displayed formula.

Assumptions for prior adjustment: The label definitions and feature representation are the same in the reference and target data; \(P_T(X\mid Y=k)=P_R(X\mid Y=k)\) for each represented class; the reference probabilities \(q_R(k\mid x)\) are adequate estimates of reference posteriors; and every target class to be adjusted has positive reference prior. If any condition fails, the formula may not produce valid target probabilities.

Worked Example: A Binary Label Changes After Adjustment

Suppose a reference contains 80% type A and 20% type B, so \(\pi_A^{(R)}=0.80\) and \(\pi_B^{(R)}=0.20\). For one hypothetical cell, the classifier reports \(q_R(A\mid x)=0.60\) and \(q_R(B\mid x)=0.40\). Assume independent biological information supports target priors of 40% type A and 60% type B. The class-conditional feature distributions are assumed unchanged, and both classes occur in the reference.

For type A, the unnormalized adjusted value is \(0.60(0.40/0.80)=0.30\). For type B, it is \(0.40(0.60/0.20)=1.20\). The values sum to 1.50, so the adjusted probabilities are \(0.30/1.50=0.200\) for A and \(1.20/1.50=0.800\) for B. The predicted label changes from A under the reference probabilities to B under the target priors. This change follows from the stipulated prior shift; it does not show that the cell’s expression features changed.

Worked Multiclass Adjustment

For several cell types, each class has its own prior ratio. A rare reference class can receive a large multiplier when it is more common in the target, but its adjusted probability still depends on the original evidence and on the multipliers for all other classes.

Worked Example: Three Cell Types and a Chosen Target Prior

A hypothetical reference has 70% type A, 20% type B, and 10% type C. For a cell, the reference classifier probabilities are 0.60, 0.30, and 0.10, respectively. Suppose a separate, relevant source justifies target priors of 40% A, 40% B, and 20% C. Assume stable class-conditional feature distributions and adequate reference probabilities.

The class-specific multipliers are \(0.40/0.70=0.5714\) for A, \(0.40/0.20=2\) for B, and \(0.20/0.10=2\) for C. The unnormalized adjusted values are therefore \(0.60(0.5714)=0.3429\), \(0.30(2)=0.6000\), and \(0.10(2)=0.2000\). Their sum is 1.1429.

After normalization, the probabilities are approximately \(0.3429/1.1429=0.3000\) for A, \(0.6000/1.1429=0.5250\) for B, and \(0.2000/1.1429=0.1750\) for C. The largest adjusted probability is for B, whereas the largest reference probability was for A. Under the specified target prior, B is the predicted class. The calculation does not imply that B is certain; it gives the model’s adjusted probability under the stated assumptions.

Choosing a Prior and Checking the Adjustment

The empirical reference frequency is easy to calculate, but its meaning depends on how reference cells were selected. If cells were sampled to represent the target population, \(n_k/n\) may be a reasonable estimate of \(\pi_k^{(R)}\), and the same frequencies may be suitable target priors. If the reference was deliberately balanced across types, or assembled from distinct enriched samples, its observed frequencies describe that dataset’s sampling composition rather than necessarily describing the target.

A chosen target prior should come from a defensible source, such as a representative sample of the intended tissue or a clearly stated application-specific assumption. It should sum to one, use the same set of labels, and refer to the population actually being annotated. A prior chosen merely to increase the number of rare labels is not evidence that those labels are more prevalent.

Worked Example: Empirical Frequency Is Not Automatically the Target Prior

A hypothetical reference contains 900 type A cells and 100 type B cells. Thus its empirical frequencies are \(\pi_A^{(R)}=0.90\) and \(\pi_B^{(R)}=0.10\). The reference was assembled by selecting equal numbers from two pre-existing groups, so those frequencies are not intended to estimate the target tissue composition. Suppose independent information instead supports target priors of 0.50 for each type.

For a cell with reference probabilities \(q_R(A\mid x)=0.70\) and \(q_R(B\mid x)=0.30\), the unnormalized adjusted values are \(0.70(0.50/0.90)=0.3889\) for A and \(0.30(0.50/0.10)=1.500\) for B. They sum to 1.8889. The adjusted probabilities are \(0.3889/1.8889=0.2059\) for A and \(1.500/1.8889=0.7941\) for B. Under these stated priors and assumptions, the predicted type is B.

If instead the target population were known to have the reference composition, the target priors would be 0.90 and 0.10, both multipliers would equal 1, and the original probabilities 0.70 and 0.30 would remain unchanged. The example shows why the observed reference count alone does not determine the right adjustment.

Diagnostics, Alternatives, and Limitations

Before applying the formula, check the source of the reference frequencies, the rationale for the target priors, and whether the two datasets plausibly share class-conditional feature distributions. If representative target cells with known labels are available, examine whether the adjusted predictions behave plausibly on them. Also recalculate predictions under a range of defensible target priors. If small, plausible changes in priors cause many labels to switch, the assigned labels depend strongly on prior information and should be reported as such.

A practical diagnostic is to compare the unadjusted and adjusted probability vectors for representative cells. A large change is expected when a class prior ratio is large; it is a warning to inspect the prior source and model assumptions, not proof that the correction is wrong. If target priors are estimated from the same model predictions being adjusted, the estimate and adjustment can depend on each other. Treating such a result as independent confirmation of the chosen composition would be misleading.

Prior adjustment is not the only response to imbalance. One alternative is to train with class weights or a deliberately balanced sampling scheme, changing how the model is fitted rather than correcting its output afterward. Such training choices can affect the fitted decision rule and do not automatically yield probabilities for the target population; their effective prior and probability interpretation must be understood. Another option is to retain the reference probabilities when the reference composition is the desired target composition. No adjustment is preferable to an unjustified adjustment.

Several failure modes matter in single-cell annotation. A tissue or batch difference can change \(P(X\mid Y)\), violating the prior-shift assumption; changing only class priors cannot repair that. If a target type is absent from the reference, its reference prior is zero and the ratio is undefined: the formula cannot create evidence for a class the model did not represent. Likewise, an inaccurate or uncalibrated reference probability cannot be made reliable merely by multiplying it by prior ratios. The adjustment also does not resolve uncertainty about label granularity or establish that a predicted cell type is biologically correct.

Common mistake: “Type B is rare in the reference, so every B probability should be increased.” The correction is not to increase rare-class probabilities automatically. First decide whether the reference frequency differs from the intended target prior, then use the ratio \(\pi_k^{(T)}/\pi_k^{(R)}\) under the prior-shift assumptions. If reference and target priors are equal, the correct adjustment is no adjustment.

Define the target.
Specify the cell population and the label set to which predictions will apply.

Specify both priors.
Document the reference prior used by the classifier and the evidence supporting the target prior.

Check assumptions and support.
Assess whether class-conditional features are plausibly stable, probabilities are usable, and each target class appears in the reference.

Reweight and normalize.
Apply the prior ratio to each class probability, normalize across classes, and assess sensitivity to plausible priors.

Check Your Understanding

Answer each question, then compare with the brief solutions below.

A classifier has reference priors 0.75 for A and 0.25 for B, and reference probabilities 0.80 for A and 0.20 for B. If target priors are also 0.75 and 0.25, what are the adjusted probabilities?
For a cell, the reference probabilities are 0.50 for A and 0.50 for B. Reference priors are 0.80 and 0.20, while target priors are 0.40 and 0.60. Calculate the adjusted probabilities and predicted class.
Why can a balanced reference’s observed class frequencies be a poor choice for target priors?
What prevents the formula from assigning a positive adjusted probability to a cell type absent from the reference?

Solutions: (1) Both prior ratios equal 1, so the probabilities remain 0.80 for A and 0.20 for B. (2) The unnormalized values are \(0.50(0.40/0.80)=0.25\) for A and \(0.50(0.60/0.20)=1.50\) for B. The sum is 1.75, giving adjusted probabilities \(0.25/1.75=0.1429\) for A and \(1.50/1.75=0.8571\) for B; the predicted class is B. (3) Balancing may be a sampling or training choice rather than a representation of the target population’s composition. (4) Its reference prior is zero, so the target-to-reference ratio is undefined; the class is unsupported by the reference classifier.