[Paper Review] Learning to Balance: Bayesian Meta-Learning for Imbalanced and Out-of-distribution Tasks

Haebeom Lee, Hayeon Lee+5 more|arXiv (Cornell University)|May 30, 2019

Domain Adaptation and Few-Shot LearningComputer Science22 references56 citations

TL;DR

The paper introduces Bayesian Task-Adaptive Meta-Learning (Bayesian TAML) that adaptively balances meta-knowledge and task-/class-specific learning under task and class imbalance and out-of-distribution (OOD) scenarios, using variational inference. It outperforms existing meta-learning methods on realistic imbalanced few-shot benchmarks.

ABSTRACT

While tasks could come with varying the number of instances and classes in realistic settings, the existing meta-learning approaches for few-shot classification assume that the number of instances per task and class is fixed. Due to such restriction, they learn to equally utilize the meta-knowledge across all the tasks, even when the number of instances per task and class largely varies. Moreover, they do not consider distributional difference in unseen tasks, on which the meta-knowledge may have less usefulness depending on the task relatedness. To overcome these limitations, we propose a novel meta-learning model that adaptively balances the effect of the meta-learning and task-specific learning within each task. Through the learning of the balancing variables, we can decide whether to obtain a solution by relying on the meta-knowledge or task-specific learning. We formulate this objective into a Bayesian inference framework and tackle it using variational inference. We validate our Bayesian Task-Adaptive Meta-Learning (Bayesian TAML) on multiple realistic task- and class-imbalanced datasets, on which it significantly outperforms existing meta-learning approaches. Further ablation study confirms the effectiveness of each balancing component and the Bayesian learning framework.

Motivation & Objective

• Motivate meta-learning under realistic task distributions with varying numbers of shots per task and per class.
• Develop a framework that adaptively balances meta-knowledge and task-specific learning for each task and class.
• Handle class imbalance, task imbalance, and out-of-distribution tasks within a unified Bayesian meta-learning model.
• Propose three balancing variables to modulate learning dynamics and initialization per task.
• Validate the approach on imbalanced few-shot benchmarks and analyze component contributions.

Proposed method

• Introduce task-dependent balancing variables to adjust inner-gradient updates: class-wise learning-rate weights, per-layer learning-rate multipliers, and task-dependent initial parameter modulation.
• Formulate a Bayesian meta-learning objective with an amortized variational posterior over balancing variables to capture uncertainty (variational lower bound with KL regularization).
• Use a hierarchical set-encoder StatisticsPooling to encode task representations from per-class samples, producing inputs for the balancing variables.
• Derive update rules that integrate balancing variables into the MAML-style inner-loop updates.
• Perform MC sampling during meta-testing to approximate the predictive distribution and capture uncertainty in balancing variables.

Experimental results

Research questions

• RQ1 Can a Bayesian, task- and class-adaptive balancing mechanism improve few-shot learning under realistic imbalances and distribution shifts?
• RQ2 How do task-dependent initialization, per-class gradient scaling, and per-layer learning-rate modulation contribute to performance under imbalanced and OOD tasks?
• RQ3 Does amortized variational inference over balancing variables yield robust task-specific predictors across diverse task distributions?
• RQ4 Is a hierarchical set-encoder effective at capturing task statistics for balancing variable generation?
• RQ5 Do the proposed balancing components provide measurable gains over strong meta-learning baselines on imbalanced and OOD benchmarks?

Key findings

• Bayesian TAML significantly improves over strong meta-learning baselines on imbalanced and OOD few-shot classification benchmarks.
• The three balancing variables enable effective handling of task imbalance (gamma), class imbalance (omega), and distributional shift (z) by adaptively modulating meta- and task-specific learning.
• Bayesian modeling and MC sampling of balancing variables yield substantial gains, especially for out-of-distribution tasks.
• A task-encoder based on hierarchical set statistics (mean, variance, cardinality) supports effective inference of balancing variables.
• Ablation studies demonstrate the contribution of each balancing component to overall performance.
• Model variants with Bayesian treatment of balancing variables outperform deterministic counterparts, showing the value of uncertainty in balancing decisions.