Domain Expansion: A Latent Space Construction Framework for Multi-Task Learning

AI summary

Overview

Research area: Machine learning, specifically representation learning and multi-task/multi-objective learning (cs.LG). The work sits at the intersection of latent-space design, multi-task optimization, and interpretable representation learning.

Technical level: Intermediate. The paper relies on concepts familiar to readers with some ML background — encoders and decoders, covariance and eigendecomposition, projection operators, contrastive and ranking losses, and gradient-conflict mitigation in multi-task learning. No advanced mathematics beyond linear algebra is required, and the core idea is geometric and easy to visualize.

Scope: The paper proposes and empirically validates a latent-space construction framework, built around a mechanism the authors call "orthogonal pooling," that assigns each learning objective its own mutually orthogonal subspace in order to prevent the failure mode the authors name "latent representation collapse," evaluated on ShapeNet, MPIIGaze, and a custom Rotated MNIST benchmark.

What This Paper Is About

When a single neural network is trained on several objectives at once, the objectives often produce conflicting gradients that pull shared features in opposing directions. The authors formalize the resulting failure as latent representation collapse: the network carves out only a small, compromised region of the latent space that partly satisfies every objective but excels at none, leaving entangled and uninterpretable features. The paper's goal is to prevent this interference structurally — by shaping the latent space itself rather than by re-weighting or projecting gradients during optimization.

Key Contributions

1. Formalization of latent representation collapse. The authors name and describe a critical failure mode in multi-objective representation learning, in which a single shared representation subjected to multiple objective losses is forced into a compromised state that is suboptimal for every task (illustrated in Figure 1(a) and Figure 2).

2. The Domain Expansion framework and orthogonal pooling. The paper introduces a lightweight architectural primitive that constrains features for different objectives to lie in mutually orthogonal subspaces. The framework computes the empirical mean and covariance of latent features, eigendecomposes the covariance to obtain an orthonormal basis, assigns the top M eigenvectors to the M target concepts, and decomposes each latent feature by projecting it onto those axes.

3. An explicit, interpretable, and compositional latent space. The paper shows that the resulting space supports a "concept algebra": a concept-specific adjustment operator (⊕ᵐ and ⊖ᵐ) that modifies one concept without affecting others, and a concept composition operator (⊕ and ⊖) that combines two full instantiated concepts through simple vector addition or subtraction in the latent space.

4. Evaluation across three benchmarks. The method is validated on ShapeNet (~30,000 rendered images), MPIIGaze (~30,000 images from 15 participants), and a custom Rotated MNIST dataset, on multi-objective problems combining classification with pose and gaze regression, against a weighted-sum baseline and the gradient-based methods Nash-MTL, FAMO, and IMTL.

Main Findings

• Naive multi-objective training does collapse the latent space. On ShapeNet Objective Set 2, the weighted-sum baseline scores a V-score of 0.00 for model ID while achieving 0.99 accuracy on that same classification task — a discrepancy the authors interpret as evidence that baselines learn a superficial shortcut for the predictive task while the underlying latent space remains collapsed. FAMO, Nash-MTL, and IMTL show the same pattern (V-score 0.00 under Objective Set 2 for model ID, with classification accuracy of 0.99 in all three cases).

• Domain Expansion substantially improves representation quality on ShapeNet. Under Objective Set 1, the proposed method reaches Spearman correlations of 0.95 (azimuth), 0.87 (elevation), and 0.85 (rotation), and V-scores of 0.99 (category) and 0.91 (ID). The strongest baseline values for those same metrics are 0.49/0.41/0.42 (FAMO, Nash-MTL) and 0.39/0.28 (IMTL). Under Objective Set 2, the proposed method reaches Spearman 0.95/0.87/0.85 and V-scores 0.98/0.96.

• Better predictive accuracy and error as well. On ShapeNet Objective Set 1, Domain Expansion records MAE of 0.08°, 0.08°, and 0.09° for the three regression concepts and accuracy of 0.99 (category) and 0.97 (ID). The baseline's MAEs for those concepts are 0.12°, 0.09°, and 0.09°, with accuracies of 0.28 and 0.37. Under Objective Set 2, the proposed method reaches MAE 0.07°, 0.08°, 0.09° and accuracy 0.98 and 0.94.

• The latent space supports concept arithmetic. Concept composition, measured by average cosine similarity between reconstructed and ground-truth latent vectors, reaches 0.95 under ShapeNet Objective Set 1 and 0.93 under Objective Set 2, versus 0.22 and 0.42 for the baseline, 0.28 for both FAMO settings, 0.28/0.35 for Nash-MTL, and 0.14/0.28 for IMTL.

• Results generalize to MPIIGaze. On MPIIGaze Objective Set 1, the method achieves Spearman 0.72/0.92/0.79 across the three gaze axes, a V-score of 0.99 for participant ID, and composition similarity of 0.95. Under Objective Set 2 it achieves 0.73/0.92/0.81, a V-score of 0.98, and similarity of 0.89. MAE values for the gaze axes are 0.01/0.02/0.02 in Objective Set 1 and 0.01/0.02/0.01 in Objective Set 2.

• Results generalize to Rotated MNIST. With a ResNet-18 backbone, the method reaches a rotation Spearman correlation of 0.92, an ID V-score of 0.87, MAE of 0.14 rad, accuracy of 0.88, and composition similarity of 0.73 under Objective Set 1; and 0.93, 0.96, 0.12 rad, 0.99, and 0.71 under Objective Set 2. Notably, in Objective Set 2 the baseline also reports a composition similarity of 0.71, matching the proposed method on that single metric.

• Eigenvector alignment is necessary for stability. Without Hungarian alignment, consecutive-epoch cosine similarity of the learned eigenvectors is 0.49 (azimuth), 0.04 (elevation), 0.03 (rotation), 0.05 (category), and 0.98 (ID). With Hungarian alignment, all five concepts stabilize at 0.99.

• Qualitative structure matches the quantitative claim. PCA visualizations (Figure 6) show the baseline and Nash-MTL latent spaces as entangled and unstructured, while the proposed method's latent space shows clear directional alignment for each concept.

Methodology in Plain English

The authors keep the encoder architecture standard — a ResNet-50 backbone producing a 2048-dimensional latent feature (ResNet-18 for the Rotated MNIST experiments) — and change how that feature is used. Each objective gets its own single linear layer decoder.

At each training epoch, the framework does three things:

1. Find principal axes. It computes the mean and covariance of the latent features over the current batch or the full training set, then eigendecomposes the covariance to get an orthonormal set of eigenvectors.

2. Define the orthogonal domain. It picks the top M eigenvectors (largest eigenvalues) and assigns one to each target concept, giving each concept its own 1D subspace and projection operator.

3. Orthogonal pooling. Each latent feature is centered by subtracting the mean and projected onto each of those axes. The decoders then operate on the projections, and the total loss is the weighted sum of per-concept losses applied to the independent projected features.

Because each objective's updates are confined to a distinct subspace, the gradients are decoupled by construction — the problem is prevented rather than mediated after the fact. The paper contrasts this explicitly with gradient-level methods such as GradNorm, PCGrad, IMTL, CAGrad, and MGDA, which the authors characterize as reactive.

Training proceeds in two stages: the encoder is trained while the basis and mean are updated each epoch, then the encoder is frozen and the linear decoders are trained on the final fixed representations. To stop the basis from shuffling between epochs, the authors align eigenvectors across epochs with the Hungarian algorithm (resolving both permutation and sign ambiguity).

For supervision on the latent space, the authors use two representation-learning losses: Supervised Contrastive (SupCon), adapted by replacing inner-product similarity with L2 distance to suit the projection-based method, and Rank-N-Contrast (RNC). Both use temperature τ = 2.0, and the loss weights are set to 1.0 for RNC and 0.02 for SupCon. Objective Set 1 uses RNC for all concepts; Objective Set 2 uses RNC for the regression concepts and SupCon for the classification concepts.

The evaluation uses Spearman's rank correlation for regression concepts, V-measure (V-score) for the clustering quality of classification concepts, plus MAE and accuracy as predictive metrics, and average cosine similarity for the concept-composition test. In that composition test, the test set is split into halves P and Q; the conceptual difference between two samples is computed in concept space, applied to a source concept to create a synthetic target concept, reconstructed into a latent vector, and compared by cosine similarity against the ground-truth encoded vector.

The paper also describes how the inverse decoder is made tractable. Since decoders are linear and non-square (2048-dimensional input versus a concept dimension of 1 or a class-logit count), the inverse is ill-posed. The orthogonality constraint resolves this: an arbitrary solution is found (for example via pseudo-inverse) and then projected back onto the correct eigenvector's 1D subspace.

Why This Matters

Impact on research. The paper reframes a widely studied problem. Most multi-task learning work attacks conflicting gradients during optimization; this work argues for changing the geometry of the representation instead. If the approach holds up, it suggests that orthogonality constraints on latent subspaces could be a general design principle for multi-objective representation learning, and that interpretability and compositionality can be architectural byproducts rather than post-hoc add-ons. The paper's concept-algebra framing also connects representation learning to work on disentanglement and concept manipulation.

Real-world applications (the first two are named in the paper; the others follow from the benchmarks used):

• Algorithmic fairness, which the authors list as an application area for more controllable and interpretable models.
• Controllable multi-modal content generation, also named by the authors as a target application.
• 3D object understanding, where the ShapeNet experiments combine object category and model identification with azimuth, elevation, and in-plane rotation estimation — relevant to robotics and scene understanding.
• Gaze estimation interfaces, the MPIIGaze task, where a model must predict 3D gaze direction while remaining aware of individual participant identity.
• Recognition under pose variation, the Rotated MNIST task, where classification must succeed while the model also represents rotation angle.

Industry relevance. Systems that must satisfy several objectives at once — a vision model that classifies, localizes, and estimates orientation simultaneously, or a multimodal model aligning several signals — are exactly the setting where the paper claims interference degrades shared features. The framework adds only lightweight machinery (a covariance computation, an eigendecomposition, and projection operators) around an otherwise standard encoder-decoder. The compositional property also points toward interfaces where user-level edits to a single attribute can be applied as vector arithmetic without retraining, which is attractive for controllable generation and personalized systems.

Future Directions

• Decoding abstract composed concepts. The authors identify this as their main limitation. They note the framework can represent a novel concept such as "chair" ⊕ "boat" via the operation f_chair + f_boat, but they cannot currently decode such abstract compositions into human-understandable output. They propose pairing the encoder with a generative model, such as an LLM or diffusion model.
• Continual learning with new tasks. Appendix A.4 describes an experiment for adding N new objectives (C_{M+1}, ..., C_{M+N}) to a model already trained on M tasks, by freezing existing axes, eigendecomposing a residual feature space to find new orthogonal axes, and adding an L2 regularization loss to prevent catastrophic forgetting. The provided text is truncated mid-description, so the results of this experiment are not reported in the available content.
• Scaling to larger and more diverse concept sets. The three benchmarks use four to five concepts. Whether the eigendecomposition-based basis remains stable and informative with many more objectives, or with concepts whose natural directions are not orthogonal, is not addressed.
• Comparing against a wider set of gradient-based methods. The comparison covers Nash-MTL, FAMO, and IMTL, and the related work discusses GradNorm, PCGrad, CAGrad, and MGDA; the experiments reported here do not include all of those.

Target Audience

The paper is most useful to:

• Multi-task learning researchers looking for alternatives to gradient-manipulation methods, and to anyone who wants a concrete, testable definition of the representation-collapse failure mode.
• Interpretability and disentanglement researchers, because the method produces an explicitly structured latent space where individual axes correspond to single concepts, and where algebraic operations on concepts map to vector arithmetic.
• Practitioners building multi-objective systems in vision, robotics, or multimodal settings, who need a lightweight way to keep per-task features from interfering, and who value the ability to edit one attribute without disturbing others.
• Graduate students and advanced undergraduates in machine learning, since the core mechanism — eigendecomposition plus orthogonal projection — is conceptually accessible and the paper provides both theory and multi-benchmark empirical evidence.

The paper assumes familiarity with encoder-decoder architectures, contrastive and ranking losses, and the basics of multi-task optimization; readers without that background will need to consult the cited work on SupCon, RNC, and gradient-conflict methods.

Authors’ abstract

Training a single network with multiple objectives often leads to conflicting gradients that degrade shared representations, forcing them into a compromised state that is suboptimal for any single task--a problem we term latent representation collapse. We introduce Domain Expansion, a framework that prevents these conflicts by restructuring the latent space itself. Our framework uses a novel orthogonal pooling mechanism to construct a latent space where each objective is assigned to a mutually orthogonal subspace. We validate our approach across diverse benchmarks--including ShapeNet, MPIIGaze, and Rotated MNIST--on challenging multi-objective problems combining classification with pose and gaze estimation. Our experiments demonstrate that this structure not only prevents collapse but also yields an explicit, interpretable, and compositional latent space where concepts can be directly manipulated.