Multi-task optimization improved by geometry aware weight updates
Geometry-Aware Multi-Objective Optimization: Norm-Constrained Oracles with Entropic Weight Updates
Machine Learning
Summary
Many computer programs learn to perform several tasks at once, but figuring out how to improve all tasks together is tricky. This paper looks at a new way to adjust priorities and directions for multi-task learning that respects different mathematical shapes (geometries). The authors show their method can steadily improve with certain guarantees, even when randomness is involved. They test how well it works compared to other approaches and suggest better ways to pick directions and update priorities.
What this means in practice
- •For machine learning engineers: Improve multi-task learning models by more effectively balancing tasks using norm-aware weight updates and momentum to enhance training stability.
- •For optimization software developers: Implement new optimization algorithms incorporating geometry-aware oracles and entropic updates to support complex multi-objective problems.
Authors
El Mahdi Chayti, Alireza Mirrokni, Martin Jaggi
Abstract
Multiple-gradient descent is typically formulated in Euclidean geometry, whereas matrix-aware optimizers use directions induced by other norms. We analyze a single-loop multi-objective method combining a norm-constrained linear minimization oracle with entropic updates of task weights. For any finite-dimensional norm pair, we establish an $O(T^{-1/2})$ ergodic stationarity rate in the deterministic setting, deriving the required trajectory gradient bound from smoothness and lower-boundedness in the unregularized case. With stochastic gradients and a bounded-gradient assumption, one persistent blended-gradient momentum buffer yields an $O(T^{-1/4})$ rate whose leading term has no explicit dependence on the number of tasks. We also establish static regret for the inner weighting objective on arbitrary bounded gradient sequences and, with entropy regularization and a weight floor, iteration-wise inner-gap control after burn-in while retaining ergodic stationarity, conditional on a trajectory gradient bound. The optimal blended gradient can be nonunique in spectral and sign geometries, motivating a selection-independent inner gap. Our analysis extends the spectral construction studied by concurrent MOON to general norms. Experiments compare the method with existing multi-task optimizers and examine its oracle, weighting, and momentum choices.