🤖 machine learning

Rethinking Loss Reweighting for Imbalance Learning as an Inverse Problem: A Neural Collapse Point of View

This paper proposes a novel inverse-view loss reweighting strategy for long-tailed classification that dynamically infers class weights to achieve an ideal equal per-class loss objective inspired by Neural Collapse, thereby reducing loss imbalance and outperforming existing baselines.

Original authors: Jinping Wang, Zixin Tong, Zhiwu Xie, Zhiqiang Gao

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jinping Wang, Zixin Tong, Zhiwu Xie, Zhiqiang Gao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a teacher trying to grade a class of students for a final exam. However, this class has a weird problem: 90% of the students are "Head Students" who have studied the same 10 chapters over and over again, while the remaining 10% are "Tail Students" who have only studied 1 or 2 chapters.

In a standard classroom (a balanced dataset), the teacher treats everyone equally. But in this long-tailed classroom, if the teacher just grades based on the total score, the "Head Students" will dominate the conversation. The teacher will naturally focus on the 90% because they show up more often, and the "Tail Students" will get ignored, even though they are just as important to the class's overall success.

In the world of Artificial Intelligence, this is called Long-Tailed Learning. The AI model gets really good at recognizing common things (like "dogs" or "cars") but terrible at recognizing rare things (like "sloths" or "specific types of beetles").

The Old Way: Guessing the Weights

To fix this, researchers have tried Loss Reweighting. Think of "Loss" as the grade the AI gets when it makes a mistake.

The Old Strategy: Researchers would try to guess how much to "boost" the grades of the rare students. They might say, "Okay, since the sloth students are rare, let's multiply their mistakes by 10 so the teacher pays more attention to them."
The Problem: These methods are mostly guesses (heuristics). They don't really know what the perfect final grade distribution should look like. They are just trying different multipliers until something works, without a clear target.

The New Idea: The "Perfect Classroom" Blueprint

This paper introduces a new way of thinking, inspired by a phenomenon called Neural Collapse.

Imagine that when a student is perfectly prepared, their knowledge becomes super organized. In the AI world, Neural Collapse is a state where the AI's brain organizes itself perfectly:

All examples of "dogs" cluster tightly together.
All examples of "cats" cluster tightly together.
The distance between the "dog" cluster and the "cat" cluster is perfectly symmetrical.

The authors discovered a beautiful rule: In this perfectly organized state, every single class (dog, cat, sloth, beetle) should have the exact same average "grade" (loss). If the AI is truly balanced, it shouldn't be struggling more with one group than another.

The Solution: Working Backwards (The Inverse Problem)

Instead of guessing how much to boost the rare students, the authors decided to work backwards.

Set the Goal: "We want the AI to reach a state where every class has the exact same average loss."
The Inverse Question: "Given that we want equal losses, what specific multipliers (weights) do we need to apply right now to make that happen?"

They treated this like a math puzzle (an inverse problem). Instead of saying, "I think we need a weight of 5," they solved the equation: "To get an equal loss, the weight must be X."

They found a simple, closed-form formula (a direct mathematical recipe) to calculate exactly what weight each class needs at every step of the training.

The Two-Part Strategy

The paper proposes a two-step "teacher" strategy to make this work:

The Micro-Adjustment (Batch-wise): Every time the AI looks at a small group of students (a "mini-batch"), it calculates exactly how much to boost the rare ones in that specific group to balance the immediate grades. It's like a teacher instantly adjusting the difficulty of a quiz based on who is in the room today.
The Macro-Adjustment (Frequency): Sometimes, the "Tail Students" (rare classes) just don't show up often enough in these small groups. Even if you boost them when they do show up, they still get less practice time overall. So, the authors add a second layer: they track how often each class appears. If a class is rare, they give it a permanent "bonus" to compensate for the fact that it appears less often.

The Results

When they tested this "Inverse Reweighting" method:

Better Balance: The AI stopped ignoring the rare classes. The "grades" (losses) for common and rare classes became much more equal.
Perfect Geometry: The AI's internal brain structure looked more like the "perfectly organized" Neural Collapse blueprint.
Higher Scores: On standard tests (like CIFAR-100 and ImageNet), this method beat almost all other existing methods, achieving the best results (State-of-the-Art) in many cases.

Summary

Think of it like this: Old methods were like a teacher guessing, "I'll give the quiet students extra credit because they seem shy." This new method is like a teacher who says, "I know the goal is for everyone to have the same average score. Let me calculate the exact amount of help each student needs right now to make that happen, and I'll adjust it every single day."

By turning the problem into a math equation with a clear target (equal loss), the authors created a system that naturally guides the AI to a perfectly balanced state.

Technical Summary: Rethinking Loss Reweighting for Imbalance Learning as an Inverse Problem

1. Problem Statement

Deep neural networks have achieved significant success on balanced datasets like ImageNet, but real-world scenarios often exhibit long-tailed distributions, where head classes dominate the sample space while tail classes have few samples. This imbalance causes models to over-bias toward head classes, leading to poor generalization on tail classes.

While loss reweighting is a widely used strategy to mitigate this by adjusting class-wise weights, existing methods are often heuristic and target-agnostic. They lack a well-specified ideal objective for what the class-wise loss distribution should look like at the end of training. Consequently, current strategies rely on ad-hoc designs (e.g., inverse frequency, effective number of samples) without a rigorous theoretical target to guide the reweighting process.

2. Methodology

The paper proposes a novel framework that rethinks loss reweighting through the lens of Neural Collapse (NC) and formulates it as an inverse problem.

2.1. Theoretical Foundation: Neural Collapse (NC)

The authors leverage the phenomenon of Neural Collapse, observed in the terminal phase of training on balanced datasets, where:

Within-class variability collapses: Features of the same class converge to their class mean.
Simplex ETF Geometry: Class means form an ideal Equiangular Tight Frame (ETF).
Self-Duality: Classifier weights align with the centered class means.
Nearest Center: Classification simplifies to finding the nearest class mean.

Key Theoretical Insight: The paper establishes that under NC geometry (specifically NC1-NC3), all classes must share the same average loss. The authors prove that if a loss imbalance coefficient ( $\rho$ ) remains non-zero, the network cannot converge to the ideal ETF structure. Therefore, the ideal objective for reweighting in long-tailed learning is to enforce equal per-class average loss to facilitate convergence to NC-consistent geometry.

2.2. Inverse Problem Formulation

Instead of prescribing weights heuristically (the "forward" view), the authors treat reweighting as an inverse problem:

Target: The ideal state is where the effective loss for every class equals the global mean loss ( $\bar{L}$ ).
Inference: Given the current class-wise average losses $L_c(W)$ , the method infers the optimal class weights $w_c$ that would drive the system toward this balanced state.

This is formulated as a Tikhonov-regularized inverse problem:
$\min_{w_c} (w_c L_c(W) - \bar{L})^2 + \alpha (w_c - w_c^{(0)})^2$
Where:

$L_c(W)$ is the current average loss for class $c$ .
$\bar{L}$ is the global average loss.
$w_c^{(0)}$ is a prior weight (e.g., from an existing method or uniform).
$\alpha$ is a regularization coefficient.

2.3. Closed-Form Solution

The optimization problem admits a closed-form solution for the optimal weight $w_c^*$ :
$w_c^* = \frac{\bar{L} L_c(W) + \alpha w_c^{(0)}}{L_c(W)^2 + \alpha}$
This allows for a plug-and-play implementation where weights are dynamically inferred at each training step.

2.4. Two-Level Reweighting Strategy

To address the issue that tail classes appear sporadically in mini-batches (leading to weak cumulative optimization strength), the method combines two components:

Batch-wise Inverse Reweighting: Calculates weights $w_c^*$ for classes present in the current batch using the closed-form solution.
Macro-Level Reweighting: Introduces a frequency-aware compensation factor $\beta_c \propto (B_c)^{-\gamma}$ , where $B_c$ is the number of batches in which class $c$ has appeared. This ensures tail classes receive sufficient cumulative optimization strength over the entire training process.

The final effective weight is $\hat{w}_c = \beta_c \cdot w_c^*$ .

3. Key Contributions

Identification of a Missing Target: The paper identifies that current reweighting strategies lack a defined ideal loss distribution. It proposes equal per-class average loss as the theoretical target, derived from the geometric constraints of Neural Collapse.
Inverse Problem Perspective: It casts long-tailed reweighting as an inverse problem, deriving a closed-form solution to dynamically infer class weights that satisfy the equal-loss objective.
Plug-and-Play Algorithm: The proposed method is implemented as a batch-wise, closed-form algorithm that can be integrated with existing long-tailed learning frameworks.
Theoretical Justification: The authors provide proofs showing that persistent loss imbalance prevents convergence to the ideal ETF geometry, motivating the equal-loss objective.

4. Experimental Results

The method was evaluated on standard long-tailed benchmarks: CIFAR-10-LT, CIFAR-100-LT, ImageNet-LT, and iNaturalist.

Performance vs. Reweighting Baselines: On CIFAR-100-LT (IF=100), the proposed method achieved 47.88% accuracy, outperforming the standard Cross-Entropy (41.62%) and other reweighting strategies like Inv-Freq, CB, Focal, and IB. It showed consistent improvements across imbalance factors (50, 100, 200).
Neural Collapse Metrics: The method demonstrated superior recovery of NC geometry. It achieved lower NC2 (distance to Simplex ETF) and NC3 (distance to self-duality) values compared to baselines, indicating the learned features and classifiers were closer to the ideal symmetric structure. It also achieved the lowest class-wise loss imbalance coefficient ( $\rho$ ).
State-of-the-Art (SOTA) Performance: When combined with strong baselines (e.g., GLMC, ETF-DR), the method consistently improved performance, achieving SOTA results on CIFAR-100-LT (63.9% with GLMC+Loss+LR) and ImageNet-LT (57.6% with GLMC+Loss+LR).
Ablation Studies: Experiments confirmed that both the batch-wise inverse component and the macro-level compensation are complementary, with the combination yielding the best results.

5. Significance and Claims

The paper claims that by explicitly defining equal per-class average loss as the target for reweighting, the method provides a principled alternative to heuristic designs. The significance lies in:

Theoretical Alignment: Connecting reweighting directly to the geometric properties of Neural Collapse, offering a clear explanation for why certain reweighting strategies work (or fail) based on their ability to drive the system toward an ETF structure.
Practical Efficiency: The closed-form solution introduces negligible computational overhead, making it highly efficient compared to methods requiring auxiliary networks or complex meta-learning procedures.
Generalizability: The method acts as a general reweighting module that can be applied to various existing long-tailed learning frameworks (e.g., contrastive learning, two-stage methods) to boost their performance.

The authors conclude that enforcing loss balance is a fundamental requirement for achieving the ideal NC geometry in imbalanced settings, and their inverse-view strategy effectively achieves this goal.