GeoERM: Geometry-Aware Multi-Task Representation Learning on Riemannian Manifolds

TL;DR

GeoERM addresses robustness in multi-task learning by enforcing orthonormal, geometry-aware representations on the Stiefel manifold. It optimizes via a two-step procedure that combines Riemannian gradient calculations and polar retractions, ensuring updates stay on the manifold and preserve geometric structure. The approach yields improved estimation accuracy, reduced negative transfer, and resilience to adversarial label noise, demonstrated on synthetic simulations and a real-world HAR dataset. This geometry-centric framework offers a principled way to exploit cross-task structure while maintaining stability in heterogeneous settings, with practical implications for high-dimensional MTL applications. The work also identifies theoretical and scalability directions to further strengthen geometry-aware learning on manifolds.

Abstract

Multi-Task Learning (MTL) seeks to boost statistical power and learning efficiency by discovering structure shared across related tasks. State-of-the-art MTL representation methods, however, usually treat the latent representation matrix as a point in ordinary Euclidean space, ignoring its often non-Euclidean geometry, thus sacrificing robustness when tasks are heterogeneous or even adversarial. We propose GeoERM, a geometry-aware MTL framework that embeds the shared representation on its natural Riemannian manifold and optimizes it via explicit manifold operations. Each training cycle performs (i) a Riemannian gradient step that respects the intrinsic curvature of the search space, followed by (ii) an efficient polar retraction to remain on the manifold, guaranteeing geometric fidelity at every iteration. The procedure applies to a broad class of matrix-factorized MTL models and retains the same per-iteration cost as Euclidean baselines. Across a set of synthetic experiments with task heterogeneity and on a wearable-sensor activity-recognition benchmark, GeoERM consistently improves estimation accuracy, reduces negative transfer, and remains stable under adversarial label noise, outperforming leading MTL and single-task alternatives.

Figures (7)

• Figure 1: A geometric illustration of the two-step optimization process on the Stiefel manifold $\mathcal{M}$. Starting from the $k$-th iteration point of the $t$-th task, $\boldsymbol{A}_{k}^{(t)} \in \mathcal{M}$, the Euclidean gradient (black arrow) is first orthogonally projected onto the tangent space $T_{\boldsymbol{A}_{k}^{(t)}}\mathcal{M}$ (light gray plane) to obtain the Riemannian gradient (red arrow). The subsequent retraction (dotted line) maps this gradient back onto the manifold, producing the updated point $\boldsymbol{A}_{k+1}^{(t)} \in \mathcal{M}$. This process makes sure the updates $\boldsymbol{A}_{k+1}^{(t)}$ remain on $\mathcal{M}$ at every iteration, thus preserving the geometric structure of the representation.
• Figure 2: Maximum error across varying heterogeneity parameter $h$, evaluated under two outlier proportion settings: $\epsilon = 0$ (left) and $\epsilon = 0.1$ (right). Simulations are conducted with $n = 100$, $T = 50$, $p = 30$, and $r = 5$. Evaluation metrics and computational settings are described in Sections \ref{['evalmetrics']}.
• Figure 3: Maximum error across varying $h$, under $\epsilon = 0$ (left) and $\epsilon = 0.1$ (right). Simulations: $n = 100$, $T = 50$, $p = 50$, $r = 5$. Evaluation metrics and computational settings are in Sections \ref{['evalmetrics']}.
• Figure 4: Maximum error across varying $h$, under $\epsilon = 0$ (left) and $\epsilon = 0.1$ (right). Simulations: $n = 100$, $T = 50$, $p = 80$, $r = 5$.
• Figure 5: Maximum error across varying $h$, under $\epsilon = 0$ (left) and $\epsilon = 0.1$ (right). Simulations: $n = 150$, $T = 50$, $p = 80$, $r = 5$.

Theorems & Definitions (9)

• Proposition C.1: Stiefel manifold structure
• proof
• Remark C.2
• Remark C.3
• Remark C.4
• Proposition C.5: Well-defined retraction
• proof
• Proposition C.6: Uniqueness of the polar retraction
• proof