🤖 machine learning

Neural Architecture Search of Sample Reweighting Networks for Complex Distribution Shift

This paper proposes enhancing the Meta-Weight-Net (MW-Net) framework for handling simultaneous label noise and class imbalance by employing neural architecture search to optimize its network structure and input selection, thereby overcoming the performance limitations of simple networks when relying solely on classification loss.

Original authors: Keisuke Sugawara, Kento Uchida, Shinichi Shirakawa

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Keisuke Sugawara, Kento Uchida, Shinichi Shirakawa

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

The Big Picture: Teaching a Robot to Ignore Bad Data

Imagine you are hiring a robot to sort a huge pile of mail. You want it to learn how to deliver letters to the right houses. However, you have two big problems:

The "Noise" Problem: Some of the letters have the wrong address written on them (Label Noise).
The "Imbalance" Problem: You have 1,000 letters for "House A" but only 10 letters for "House B" (Class Imbalance).

If you just let the robot learn normally, it will get confused. It might ignore the rare "House B" letters because there are so few of them, or it might get frustrated by the letters with wrong addresses and start guessing randomly.

The Old Solution: A Simple "Weight" Calculator

To fix this, researchers use a trick called Sample Reweighting. Think of this as giving the robot a "confidence score" for every letter it sees.

If a letter looks easy and correct, give it a high score (pay attention to it).
If a letter looks weird or wrong, give it a low score (ignore it).

The paper focuses on a specific tool called MW-Net (Meta-Weight-Net). This is a small, simple calculator that decides the score for each letter.

How it used to work: The old MW-Net was like a very simple calculator with just one button. It looked at how "wrong" the robot was about a letter (the "loss") and decided the score.
- The Flaw: In a messy situation (where you have both wrong addresses AND rare houses), a simple calculator gets confused. A rare house letter might be hard to read (high "wrongness"), and a wrong-address letter is also hard to read (high "wrongness"). The simple calculator can't tell the difference, so it treats both the same way, which hurts performance.

The New Idea: Letting the Robot Design Its Own Calculator

The authors asked: "What if we didn't just use a simple calculator? What if we let the robot design its own, more complex calculator to figure out the scores?"

They used a technique called Neural Architecture Search (NAS). Think of this as a "Master Architect" that tries out thousands of different blueprints for the calculator until it finds the perfect one.

The Master Architect looked for two things:

The Blueprint (Structure): How many "rooms" (layers) should the calculator have? How many "workers" (nodes) should be in each room?
- Analogy: Sometimes you need a tiny shed (1 layer) to do a simple job. Other times, you need a multi-story office building (5 layers) to handle complex logic.
The Input (What it looks at): The old calculator only looked at the "wrongness" score. The new one looks at two things:
- The "wrongness" score.
- The "Context": A snapshot of what the letter actually looks like (features) and what the address should be (the label).
- Analogy: Instead of just asking "Is this letter confusing?", the new calculator asks, "Is this letter confusing, AND does it look like a rare house letter, AND does it have the right zip code?"

How They Found the Best Design

They used a smart search method called TPE (Tree-structured Parzen Estimator).

Imagine the Master Architect is playing a game of "Hot and Cold."
It tries a blueprint (e.g., 3 rooms, 500 workers).
It tests the robot. If the robot does well, the Architect remembers that blueprint.
If the robot does poorly, the Architect forgets it.
Over time, the Architect learns the exact shape of the calculator needed for the specific type of messiness (e.g., "Flip Noise" vs. "Random Noise").

What They Discovered

They tested this on two famous image datasets (CIFAR-10 and CIFAR-100), which are like giant boxes of photos of cats, dogs, cars, etc., but they messed up the labels and made some categories rare.

The Search Works: The "Master Architect" consistently found better calculators than the old, simple one. The robot sorted the mail much better.
Different Messes Need Different Tools:
- Flip Noise (Systematic errors): When the wrong addresses were specific (e.g., all "Cats" were labeled as "Dogs"), the Architect built a deep, complex calculator (3+ layers). It needed depth to understand the specific pattern of the mistake.
- Uniform Noise (Random errors): When the wrong addresses were totally random, the Architect built a simpler calculator (1 or 2 layers) but made it wider (more workers). It didn't need deep logic, just a lot of eyes to spot the random errors.
Looking at the Right Place: The search also figured out where to look at the letter. Sometimes it needed to look at the very end of the robot's brain (the final layers), and sometimes it needed to look at the middle layers, depending on the type of noise.

The Bottom Line

The paper shows that when data is messy (noisy and unbalanced), you shouldn't just use a "one-size-fits-all" simple tool to fix it. Instead, you should use an automated search to design a custom tool that is perfectly shaped for the specific type of mess you are dealing with. This makes the learning robot much smarter and more accurate.

Technical Summary: Neural Architecture Search of Sample Reweighting Networks for Complex Distribution Shift

Problem Statement

Distribution shifts, specifically the co-occurrence of label noise and class imbalance, pose significant challenges to standard supervised learning. While Sample Reweighting is a established approach to mitigate these issues, existing methods face limitations in complex scenarios:

Manual Tuning: Traditional methods (e.g., Focal Loss, self-paced learning) rely on manually designed weighting functions and hyperparameters, requiring expert knowledge.
Limitations of Meta-Weight-Net (MW-Net): MW-Net automates weight computation using a simple neural network (a single hidden layer of 100 nodes) that takes only the classification loss as input. While effective for single types of shifts, MW-Net struggles when label noise and class imbalance coexist. In such cases, both noisy samples (which should be down-weighted) and minority class samples (which should be up-weighted) often exhibit high loss values, making it difficult for a simple, loss-only network to infer appropriate weights.

Methodology

The authors propose enhancing MW-Net through Neural Architecture Search (NAS) to optimize both the network structure and the input features. The core modifications are:

Enriched Input Features: Instead of using only the scalar loss value, the input to the weighting network is replaced with a concatenation of:
- Intermediate feature representations ( $b_j$ ) from a selected block of the classification model (processed via global average pooling).
- One-hot encoded label information ( $y_i$ ).
- Formula: $\phi(x_i, y_i; w) = (\psi_{ave}(b_j(x_i)), \psi_{onehot}(y_i))$ .
Joint Optimization via TPE: The authors employ the Tree-structured Parzen Estimator (TPE), a Bayesian optimization method, to jointly search for:
- Network Architecture: The number of hidden layers (1–5) and the number of nodes per layer (64–1024, optimized on a log-scale).
- Input Source: The specific block index (11–15) of the ResNet-32 classifier from which features are extracted.
Training Framework: The method utilizes a meta-learning framework with two datasets:
- Training Dataset ( $D_{train}$ ): Contains noisy and imbalanced data.
- Meta-Dataset ( $D_{meta}$ ): A small, clean, and balanced dataset sampled from the target distribution.
- The classifier and the weighting network are updated alternately. The classifier minimizes the weighted loss on $D_{train}$ , while the weighting network minimizes the loss on $D_{meta}$ after the classifier update.

Key Contributions

First Application of NAS to MW-Net: To the authors' knowledge, this is the first study to apply Neural Architecture Search to sample weighting networks within a meta-learning setting.
Handling Complex Shifts: The proposed method addresses the specific challenge where noisy and minority samples share high loss values, a scenario where previous loss-based weighting networks fail.
Automated Feature Selection: The method automatically identifies whether early or late-stage features from the classifier are more suitable for weight estimation under different noise types (e.g., flip noise vs. uniform noise).

Experimental Results

The method was evaluated on CIFAR-10 and CIFAR-100 datasets modified with synthetic label noise (Flip and Uniform, 40% rate) and class imbalance (Imbalance Factors of 20 and 100).

Performance on CIFAR-10:
- The proposed NAS methods consistently outperformed the standard MW-Net baseline (1 hidden layer, 100 nodes) and a "Largest Architecture" baseline (5 layers, 1024 nodes).
- Flip Noise: The search tended to select deeper architectures (3+ layers) and features from the final blocks (14th/15th) of the classifier.
- Uniform Noise: The search favored architectures with fewer layers (1–2) and features from earlier blocks (11th–13th).
- Weight Distribution: Visual analysis showed that the NAS-derived networks successfully assigned low weights to noisy samples while maintaining high weights for clean minority samples, a distinction the baseline failed to make effectively.
Performance on CIFAR-100:
- While NAS methods generally outperformed the baseline, the performance gap was smaller compared to CIFAR-10.
- The authors attribute this to the extreme scarcity of clean minority samples in CIFAR-100 (e.g., only ~3 clean samples per minority class under high imbalance and noise), making the distinction between clean and noisy data inherently difficult regardless of architecture.

Significance and Claims

The paper claims that applying NAS to MW-Net enables the model to learn sophisticated sample weighting strategies that are robust to complex distribution shifts. By optimizing the depth, width, and input feature selection, the method adapts to the specific characteristics of the noise (e.g., structured flip noise vs. random uniform noise).

The authors maintain a modest tone regarding the results on CIFAR-100, acknowledging that the task difficulty limits the observable gains of architectural differences. They explicitly state that the primary goal was to examine the effect of applying NAS to MW-Net rather than achieving state-of-the-art absolute performance.

Future Directions:
The authors suggest future work should focus on:

Jointly optimizing the architectures of both the weighting network and the classifier.
Expanding the search space to include convolutional and pooling layers.
Reducing computational costs by adopting One-Shot NAS or other efficient NAS methods, as the current "train-from-scratch" approach for every candidate is computationally prohibitive for large-scale datasets.