Sampling and Loss Weights in Multi-Domain Training

TL;DR

This paper investigates how sampling and loss weights in multi-domain training influence model performance, demonstrating their roles in variance reduction and generalization through theoretical and empirical analysis.

Contribution

It provides a rigorous study of sampling and loss weights, revealing their complementary roles and how they can be combined to improve training outcomes.

Findings

Sampling and loss weights reduce gradient variance.

They improve generalization by lowering the generalization gap.

Joint dynamics of weights enhance understanding of data mixing effects.

Abstract

In the training of large deep neural networks, there is a need for vast amounts of training data. To meet this need, data is collected from multiple domains, such as Wikipedia and GitHub. These domains are heterogeneous in both data quality and the diversity of information they provide. This raises the question of how much we should rely on each domain. Several methods have attempted to address this issue by assigning sampling weights to each data domain using heuristics or approximations. As a first step toward a deeper understanding of the role of data mixing, this work revisits the problem by studying two kinds of weights: sampling weights, which control how much each domain contributes in a batch, and loss weights, which scale the loss from each domain during training. Through a rigorous study of linear regression, we show that these two weights play complementary roles. First, they can reduce the variance of gradient estimates in iterative methods such as stochastic gradient descent (SGD). Second, they can improve generalization performance by reducing the generalization gap. We provide both theoretical and empirical support for these claims. We further study the joint dynamics of sampling weights and loss weights, examining how they can be combined to capture both contributions.