Objective-Specific Privileged Bases via Full-Prefix Matryoshka Learning

TL;DR

This paper introduces a novel Matryoshka Representation Learning method that learns task-specific privileged bases, aligning individual dimensions with task signals and recovering principal directions efficiently.

Contribution

It provides theoretical proof and empirical evidence that full-prefix MRL recovers ordered principal directions and yields task-aligned, informative representations.

Findings

Full-prefix MRL recovers ordered principal directions efficiently.

MRL produces representations where coordinate magnitude reflects informativeness.

Empirical results show consistent per-dimension structure aligned with task signals.

Abstract

Learned representations are often invariant to rotational transformations, leaving individual dimensions non-identifiable and interchangeable. We study how Matryoshka Representation Learning (MRL) induces a task-aligned privileged basis distinct from variance-based or regularizer-induced orderings. In the linear setting, we prove that full-prefix MRL recovers the ordered principal directions, and can be computed efficiently using shared statistics. Empirically, we demonstrate that MRL yields consistent per-dimension structure aligned with task signal, where coordinate magnitude reflects informativeness.