[Paper Review] Unshuffling Data for Improved Generalization

Damien Teney, Ehsan Abbasnejad+1 more|arXiv (Cornell University)|Feb 27, 2020

Multimodal Machine Learning Applications82 references41 citations

TL;DR

Partition training data into multiple environments to reduce reliance on spurious correlations, train environment-specific classifiers with a shared feature extractor, and use a variance regularizer to promote stable, environment-invariant features for better OOD generalization in VQA tasks.

ABSTRACT

Generalization beyond the training distribution is a core challenge in machine learning. The common practice of mixing and shuffling examples when training neural networks may not be optimal in this regard. We show that partitioning the data into well-chosen, non-i.i.d. subsets treated as multiple training environments can guide the learning of models with better out-of-distribution generalization. We describe a training procedure to capture the patterns that are stable across environments while discarding spurious ones. The method makes a step beyond correlation-based learning: the choice of the partitioning allows injecting information about the task that cannot be otherwise recovered from the joint distribution of the training data. We demonstrate multiple use cases with the task of visual question answering, which is notorious for dataset biases. We obtain significant improvements on VQA-CP, using environments built from prior knowledge, existing meta data, or unsupervised clustering. We also get improvements on GQA using annotations of "equivalent questions", and on multi-dataset training (VQA v2 / Visual Genome) by treating them as distinct environments.

Motivation & Objective

• Motivate and address the poor OOD generalization caused by dataset biases and spurious correlations in vision-and-language tasks.
• Propose a method to partition data into training environments where spurious patterns vary while reliable patterns stay stable.
• Train multiple environment-specific classifiers with a shared feature extractor and a variance regularizer to encourage invariance across environments.
• Demonstrate the method on VQA-related tasks, including resilience to language biases (VQA-CP), invariance to equivalent questions (GQA), and multi-dataset training.
• Provide empirical analyses and sensitivity studies to hyperparameters and partitioning strategies.

Proposed method

• Partition the training data into E disjoint environments such that spurious correlations vary across environments while reliable correlations remain stable.
• Train a shared feature extractor f_theta across environments and separate classifiers W_e for each environment, with a variance regularizer to drive W_e toward a common value.
• Optimize the objective: minimize the sum of environment-specific losses plus a penalty lambda * Var_e(W_e), where Var_e(W_e) is a variance measure of the environment-specific classifiers.
• At test time, use the averaged classifier weights or prediction: Phi*(x) = _theta(x).
• Adopt either absolute or relative variance formulations to stabilize training; optionally use alternating optimization (updating theta and W_e separately) after a warm-up phase.

Experimental results

Research questions

• RQ1 How can data be partitioned into environments so that spurious correlations vary across environments while true task signals remain stable?
• RQ2 Can a shared feature extractor combined with environment-specific classifiers, regularized by variance, learn invariant predictors that generalize better to out-of-distribution data?
• RQ3 What is the impact of different environment construction strategies (ground-truth question types, unsupervised clustering) on OOD performance in VQA?
• RQ4 How does the proposed method perform across VQA-CP, GQA with equivalent questions, and multi-dataset VQA settings?
• RQ5 How sensitive are the results to the variance regularizer weight and optimization scheme?

Key findings

Method	Val. set (Other)	Test set (Other)
Baseline	54.74	43.33
Environments: random; rel. var., no alt. opt.	53.34	43.51
Environments: clustered questions; rel. var., no alt. opt.	54.10	46.35
Environments: question groups ; rel. var., no alt. opt.	53.87	47.60
+ Alternating optimization (0 warm-up epoch)	54.00	47.71
+ Alternating optimization (2 warm-up epochs)	53.90	47.82
+ Alternating optimization (4 warm-up epochs)	53.98	48.06
+ Alternating optimization (6 warm-up epochs)	53.86	47.38
Without variance regularizer	40.76	39.14
With absolute variance regularizer	51.44	46.17

• Significant improvements over baselines on VQA-CP, especially for 'Other' questions, with the proposed environment-based method.
• Using ground-truth question-type environments yields strong gains; unsupervised clustering of questions also yields notable improvements, though slightly lower than ground-truth types.
• The variance regularizer is crucial; relative variance regularization performs slightly better than absolute variance, and an alternating optimization scheme provides modest extra gains.
• Training across environments that randomize environment partitions (i.e., random environments) does not yield improvements, underscoring the need for informative environment construction.
• The method remains competitive on standard VQA splits and can complement ensembling; improvements are most pronounced for out-of-distribution generalization tasks (VQA-CP).
• On GQA, using equivalent-question annotations improves robustness; on multi-dataset VQA (VQA v2 / Visual Genome), treating datasets as separate environments provides small gains.