ForAug: Mitigating Biases in Image Classification via Controlled Image Compositions

Abstract

Large-scale image classification datasets exhibit strong compositional biases: objects tend to be centered, appear at characteristic scales, and co-occur with class-specific context. By exploiting such biases, models attain high in-distribution accuracy but remain fragile under distribution shifts. To address this issue, we introduce ForAug, a controlled composition augmentation scheme that factorizes each training image into a foreground object and a background and recombines them to explicitly manipulate object position, object scale, and background identity. ForAug uses off-the-shelf segmentation and inpainting models to (i) extract the foreground and synthesize a neutral background, and (ii) paste the foreground onto diverse neutral backgrounds before applying standard strong augmentation policies. Compared to conventional augmentations and content-mixing methods, our factorization provides direct control knobs that break foreground-background correlations. Across 10 architectures, ForAug improves ImageNet top-1 accuracy by up to 6 percentage points (p.p.) and yields gains of up to 7.3 p.p. on fine-grained downstream datasets. Beyond accuracy, ForAug provides a controlled way to quantify and mitigate background reliance, foreground focus, center bias, and size bias. Our code and dataset are publicly available on GitHub.

Introduction

Image classification – teaching computers to label images – is a cornerstone of AI vision, powering everything from medical diagnosis to autonomous driving. Datasets like ImageNet have been crucial, especially with the rise of powerful models like Vision Transformers (ViTs).

The catch is that these datasets carry strong compositional biases: objects tend to be centered, appear at a characteristic size, and sit on class-specific backgrounds. Models happily exploit these shortcuts to boost in-distribution accuracy, but that makes them fragile once the composition shifts. ViTs are especially exposed here – unlike older Convolutional Neural Networks (CNNs), they don’t inherently understand that an object stays the same regardless of where it sits in the image (they lack “translation equivariance”) – and standard augmentations (flipping, cropping) don’t fix it.

To tackle these problems, we propose ForAug, a novel data augmentation to help ViTs learn invariances during training. The core idea? Make the spatial relationships explicit in the training data. ForAug achieves this by:

1. Separating foreground objects from their backgrounds in the dataset.
2. Recombining these objects with different backgrounds on-the-fly during training.
3. Controlling the object’s size and position during this recombination.

The results? Training ViTs with ForAug instead of standard ImageNet boosts accuracy by up to 6 percentage points on ImageNet classification, with gains of up to 7.3 p.p. on fine-grained downstream tasks.

Furthermore, ForAug provides powerful new ways to analyze model biases. Researchers can now precisely measure:

• Background Robustness: How much does the background influence the prediction?
• Foreground Focus: Does the model correctly focus on the main object?
• Center & Size Bias: Is the model overly reliant on objects being centered or a specific size?

Training with ForAug demonstrably reduces these biases, leading to more robust models.

ForAug (Method)

So, how does ForAug actually build these dynamic training images? The process, visualized above, involves two main stages: an offline Segmentation stage and an online Recombination stage.

Segmentation

The process kicks off with the Segmentation stage, a one-time, offline preparation step performed before model training even begins. Think of it as carefully prepping the visual ingredients. Here, we leverage the state-of-the-art Grounded SAM segmentation model, guiding it with the known class label of each image (e.g., instructing it to specifically find the ‘golden retriever’) to precisely isolate the main subject. Once the foreground object is digitally ‘cut out’, an object removal or ‘inpainting’ model intelligently fills the resulting hole in the original background, ensuring the backdrop looks natural and plausible. Crucially, not all generated assets make the cut; a filtering step employs other pre-trained AI models to assess quality. This ensures only clearly defined foregrounds and clean backgrounds – ones that don’t inadvertently give away the object’s identity or look overly artificial – are selected. This meticulous preparation yields the core assets for ForAug: a collection of ready-to-use foreground objects (with transparency) and a diverse pool of cleaned-up backgrounds.

Recombination

With the assets prepared, the real action unfolds during the Recombination stage, which happens dynamically online while the Vision Transformer is training. This is where ForAug truly comes alive, creating new training examples on the fly. For every foreground object prepared in the first stage, the system randomly selects a background to pair it with. This background might be the object’s original one, perhaps one from another image belonging to the same object class, or even a completely unrelated background drawn from the entire dataset to maximize contextual variety. The chosen foreground object is then randomly resized (within sensible limits based on its original appearance) and placed at a random position onto this background canvas. To create a more seamless integration, a subtle smoothing effect is applied to the object’s edges where it meets the new background. Only after this dynamic composition is complete does the resulting image undergo the standard data augmentation techniques commonly used in AI training, like random color shifts or minor flips. This constant mixing-and-matching means that each time the AI cycles through the training data, it encounters familiar objects in entirely new visual contexts. This directly forces the ViT to learn robust features that identify the object itself, effectively teaching it the spatial invariance that doesn’t come built-in, by demonstrating repeatedly that appearance, not specific placement or background, is what defines the object.

Experiments

Image Classification Results

We compare training on ImageNet with and without ForAug for 10 different models:

We find that training with ForAug increases the accuracy of every model by up to 6 p.p. It also combats the overfitting problem of larger models.

When finetuning these models on 5 fine-grained down-stream datasets, we find that the ForAug-pretrained models consistently outperform the ImageNet-pretrained ones. Especially when looking at the transformer-based models.

Model Robustness

We also evaluate multiple robustness metrics.

Distribution-Shift Robustness

Breaking the foreground–background correlation during training pays off exactly where compositional shortcuts hurt: out of distribution. On standard distribution-shift benchmarks, ForAug-trained models gain up to 19 p.p. over their ImageNet-trained counterparts, with the largest gains on the ViT backbones.

Background Robustness

We check the background robustness of models, by inspecting the accuracy-change when evaluating with ForAug using backgrounds from the same class compared to backgrounds from all classes:

Training with ForAug reduces the Background Gap for all transformer models.

Foreground Focus

Since we have the foreground segmentation masks, we can also investigate the foreground focus of the trained models. For this, we utilize different input-importance metrics like GradCAM and IntegratedGradients (IG). We define a models foreground focus, by how much more it focuses on the foreground object compared to a uniform distribution:

We find that training with ForAug mostly significantly improves the foreground focus of all models.

Center Bias

Since we can freely change the object’s position and size, we can evaluate the model bias when the position changes. For this, we subdivide the image into $3 \times 3$ sections (nonants) and place each object only in one nonant. We then compare the accuracy of a model when an object is in a specific nonant to when it’s in the center nonant.

Our center-bias score is defined at the mean of (1) the worst accuracy in a corner and (2) the worst accuracy on an edge, relative to the accuracy in the center.

We visualize the center bias for 3 instantiations of each model. Training with ForAug significantly reduces the center bias; especially of larger transformers. We also find that when training on ImageNet, model consistently perform better when an object is on the right side of an image compared to the left side (even though we use 50% random flipping during training of all models).

Size Bias

We vary the object size by an additional factor of $f_\text{size}$ to see how the model accuracy changes relative to $f_\text{size} = 1$.

Using ForAug significantly reduces the accuracy drop-off when going towards smaller objects. These gains come on top off the overall better accuracy (at $f_\text{size} = 1$).

Conclusion

So, what’s the big takeaway from ForAug? This research introduces a genuinely novel data augmentation scheme designed specifically to enhance how Vision Transformers learn to classify images. By cleverly separating foreground objects from their backgrounds and dynamically recombining them during training, ForAug tackles a key characteristic of Transformer models head-on.

As the results clearly demonstrate, this dynamic approach pays off significantly. Training models with ForAug leads to substantial performance boosts on the standard ImageNet benchmark and translates to impressive gains on related fine-grained classification tasks downstream.

But the impact of ForAug extends beyond just improving accuracy scores. It also provides a powerful and much-needed framework for analyzing model behavior and uncovering hidden biases. Crucially, the experiments show that training with ForAug doesn’t just highlight these biases; it actively reduces them. This results in models that are not only more accurate but also more robust, reliable, and generalizable to varied real-world conditions.

Citation

If you use this work, please cite our paper:

@misc{nauen2025foraug,
  title = {ForAug: Mitigating Biases in Image Classification via Controlled
           Image Compositions},
  author = {Tobias Christian Nauen and Brian Moser and Federico Raue and
            Stanislav Frolov and Andreas Dengel},
  year = {2025},
  eprint = {2503.09399},
  archiveprefix = {arXiv},
  primaryclass = {cs.CV},
}