Smarter Training Data, Not Algorithms, Is the Key to Easier Artificial Intelligence

Researchers are increasingly recognising that deliberately altering the training data distribution can dramatically improve machine learning outcomes. Marko Medvedev from University of Chicago, Idan Attias from Toyota Technological Institute at Chicago and University of Illinois Chicago, Elisabetta Cornacchia from Inria, DI/ENS, PSL, Theodor Misiakiewicz from Yale University, Gal Vardi from Weizmann Institute of Science, and Nathan Srebro from Toyota Technological Institute at Chicago, et al., demonstrate a framework termed ‘Positive Distribution Shift’ (PDS), challenging the conventional view of distribution shifts as inherently detrimental. This work posits that a carefully selected training distribution can actually simplify learning, transforming computationally intractable problems into those solvable with standard techniques. By formalising PDS and illustrating its benefits with specific examples, the authors highlight a shift in focus towards designing effective training distributions, rather than solely optimising algorithms, and reveal the often-overlooked computational advantages this approach provides.

This work challenges the conventional view of distribution shift as a hindrance, instead revealing its potential to transform computationally intractable problems into tractable ones.

The study formalizes different variants of PDS and establishes how certain complex classes become readily learnable under these conditions, drawing connections to membership query learning. This innovation centres on computational efficiency rather than statistical gains, allowing standard gradient-based training methods to succeed where they previously failed.
The research addresses a common scenario where obtaining data from the target distribution is difficult or expensive, necessitating training on an alternative, more accessible distribution. Traditionally, the focus has been on mitigating the negative effects of this shift, but this work proposes actively leveraging it for benefit.

By carefully selecting the training distribution, algorithms can overcome computational barriers, even with fixed model architectures and training algorithms like Stochastic Gradient Descent. This approach aligns with current machine learning trends, where finding optimal training distributions is often more impactful than developing novel algorithms or models.

Specifically, the study investigates how PDS can facilitate learning in cases where information-theoretic learnability exists, but worst-case computational complexity prevents practical implementation. Researchers demonstrate that with an appropriate distribution shift, even sophisticated deep models can be trained efficiently using standard methods.

Recent work introduced the notion of positive distribution shift for mixture distributions, and this study extends that concept to covariate shift in supervised classification, where the goal is to learn a predictor under a target distribution, but training occurs using a different distribution with the same conditional distribution. The framework defines PDS learning, outlining conditions for a learning rule to achieve a desired error rate and runtime.

Investigations focus on hypothesis classes that are computationally hard to learn under standard conditions, exploring whether PDS can enable tractable learning through tractable algorithms, stylized SGD, or empirical results with standard SGD on standard networks. This work provides a theoretical foundation for understanding and harnessing the benefits of positive distribution shift, offering a new perspective on the challenges and opportunities in modern machine learning.

Layerwise training and covariance loss for a two-layer ReLU network

A two-layer network with ReLU activations forms the core of the experimental setup used to investigate Positive Distribution Shift (PDS). The network, denoted as fNN(x; θ) = Σj∈[N] ajReLu(⟨wj, x⟩+ bj), comprises N neurons, where θ = (a, W, b) represents the network parameters. Layerwise training is implemented, initially optimising the bottom layer while the top layer remains fixed, followed by training the top layer with the bottom layer held constant.

This process utilizes a covariance loss function, Lcov(x, y, f) = (1 −cy y)(1 −y f(x))+, calculated from a sample of m labelled examples {(xi, yi)}i∈[m], where y represents the average label and c is a positive constant satisfying c·y The second layer weights are initialised to κ for i ∈[N/2] and −κ for i ∈[N/2], where κ is a positive constant. Biases are set to κ for all i ∈[N], and subsequently updated during the first training step to values drawn i.i.d. from a uniform distribution Unif[−L, L], with L ≥κ.

To construct the positively shifted distribution D′, a vector μ is sampled from a uniform distribution Unif[−1, 1]⊗d, defining a distribution Dμ where coordinates xi follow a Rademacher distribution Rad((μi + 1)/2). D′ is then created as a mixture of Dμ and the uniform distribution D0, with equal probabilities of 1/2.

The research demonstrates that k-sparse juntas, a class of functions with limited dependence on input variables, are learnable under PDS with a sample complexity of n = O(dk + 2k) queries on D′. Gradient descent is performed for T = Ω(ε−2(1−2η)−2) steps on a network of width N = Ω(ε−1(1−2η)−1), using m = Ω(d log(1/ε)ε−2(1−2η)−2) samples from D′ to achieve an error of EMDES′∼(D′,f)m[L(D,f)(fNN(x; θ(T)))].

Statistical and computational complexity of parity learning with uniform and biased distributions

Parity functions, under a uniform input distribution, require O(k log(d)) samples to identify the true support with high probability, where ‘k’ denotes the support size and ‘d’ represents the input dimension. Even with label noise present, the statistical sample complexity for learning these parities remains at poly(k log d) and does not benefit from distribution shift from an information-theoretical perspective.

However, computational hardness persists, with statistical query models demonstrating a requirement of dΩ(k) cost for learning parities and a widely held conjecture suggesting no poly(d, k, ε) time algorithm exists for learning k-sparse parities over a uniform distribution with noise η. Training on a shifted distribution, D′, instead of the uniform distribution D, offers significant computational advantages by revealing structural information not visible under uniform inputs.

Specifically, employing a product distribution D′ with non-zero bias for all bits creates a substantial correlation between coordinates within the support S and the target parity function χS. This increased correlation allows for a poly-time algorithm to identify the support by computing the correlation of each coordinate with the target and outputting the parity over those with the highest correlations, generalising back to the original distribution D.

This enhanced correlation also facilitates tractable learning with gradient descent, enabling the identification of the support before transforming the problem into a linear predictor, subsequently learnable by standard stochastic gradient descent on a neural network. This approach aligns with the concept of positive distribution shift, where the target function appears as a staircase in the Fourier basis of D′, making it easily learnable even with SGD and generalising to the initial test distribution D. The research focuses on identifying target functions learnable with respect to a starting distribution D using SGD on a neural network, leveraging a shifted training distribution D′.

Formalising and demonstrating benefits of training on mismatched data distributions

Researchers are investigating a phenomenon termed Positive Distribution Shift (PDS), where deliberately training machine learning models on data drawn from a different distribution than the target application can surprisingly improve learning efficiency. This contrasts with conventional thinking which views distribution shifts as detrimental to performance.

The core idea is that a carefully chosen training distribution can make computationally challenging problems more tractable, even when using standard training algorithms like stochastic gradient descent on conventional neural networks. This work formalises several variants of PDS and demonstrates its benefits across different learning scenarios, including those with perfect data, noisy labels, and general agnostic learning.

Evidence supporting PDS learnability is presented through theoretical proofs using simplified algorithms, proofs with tailored algorithms and networks, and empirical results using standard methods. Specifically, the researchers show that parity functions, known to be computationally difficult under a uniform input distribution, become efficiently learnable when trained on a product distribution with biased inputs.

This is because the altered distribution reveals correlations between input features and the target function that are hidden in the uniform case, simplifying the learning process. The authors acknowledge that rigorously proving learnability with standard techniques remains a significant challenge. Their current findings offer evidence through various approaches, rather than a single, universally applicable proof.

Future research directions include extending the analysis to more complex function classes and exploring the limits of PDS in different computational models. The implications of this work suggest a shift in focus for machine learning research, potentially prioritising the design of beneficial training distributions over algorithmic improvements, particularly for problems where computational cost is a major bottleneck.