Stochastic Nested Compositional Bi-level Optimization for Robust Feature Learning

TL;DR

This paper introduces stochastic algorithms for nested compositional bi-level optimization problems, achieving efficient solutions without matrix inversions, and demonstrates their application in robust feature learning under covariate shift.

Contribution

The paper develops a novel stochastic approximation algorithm for complex nested bi-level problems, addressing multiple sources of bias and providing theoretical convergence guarantees.

Findings

Achieves $ ilde{O}_T(1/epsilon^{2})$ oracle complexity for $epsilon$-stationary solutions.

Handles three sources of bias in stochastic gradients due to composition, bi-level structure, and Neumann series approximation.

Demonstrates effectiveness in robust feature learning for deep neural networks under covariate shift.

Abstract

We develop and analyze stochastic approximation algorithms for solving nested compositional bi-level optimization problems. These problems involve a nested composition of $T$ potentially non-convex smooth functions in the upper-level, and a smooth and strongly convex function in the lower-level. Our proposed algorithm does not rely on matrix inversions or mini-batches and can achieve an $\epsilon$-stationary solution with an oracle complexity of approximately $\tilde{O}_T(1/\epsilon^{2})$, assuming the availability of stochastic first-order oracles for the individual functions in the composition and the lower-level, which are unbiased and have bounded moments. Here, $\tilde{O}_T$ hides polylog factors and constants that depend on $T$. The key challenge we address in establishing this result relates to handling three distinct sources of bias in the stochastic gradients. The first source arises from the compositional nature of the upper-level, the second stems from the bi-level structure, and the third emerges due to the utilization of Neumann series approximations to avoid matrix inversion. To demonstrate the effectiveness of our approach, we apply it to the problem of robust feature learning for deep neural networks under covariate shift, showcasing the benefits and advantages of our methodology in that context.