Leveraging Task Structures for Improved Identifiability in Neural Network Representations

TL;DR

This paper demonstrates that access to task distributions enhances the identifiability of neural network representations, enabling better recovery of latent factors and causal structures, with practical benefits shown on synthetic and molecular data.

Contribution

It introduces a theoretical framework for multi-task identifiability, showing how task distributions and causal assumptions improve latent factor recovery and enable simple optimization methods.

Findings

Linear identifiability is achievable in multi-task regression.

Task distributions reduce the equivalence class to permutations and scaling.

The proposed method outperforms unsupervised models on synthetic and real data.

Abstract

This work extends the theory of identifiability in supervised learning by considering the consequences of having access to a distribution of tasks. In such cases, we show that linear identifiability is achievable in the general multi-task regression setting. Furthermore, we show that the existence of a task distribution which defines a conditional prior over latent factors reduces the equivalence class for identifiability to permutations and scaling of the true latent factors, a stronger and more useful result than linear identifiability. Crucially, when we further assume a causal structure over these tasks, our approach enables simple maximum marginal likelihood optimization, and suggests potential downstream applications to causal representation learning. Empirically, we find that this straightforward optimization procedure enables our model to outperform more general unsupervised models in recovering canonical representations for both synthetic data and real-world molecular data.