Soft Task-Aware Routing of Experts for Equivariant Representation Learning

TL;DR

This paper introduces STAR, a routing strategy for projection heads in equivariant representation learning, enabling experts to specialize in shared or task-specific features, leading to more efficient and effective representations.

Contribution

STAR models projection heads as experts to reduce redundancy and improve the learning of shared and task-specific features in equivariant representations.

Findings

Lower canonical correlations between invariant and equivariant embeddings.

Consistent improvements across diverse transfer learning tasks.

Abstract

Equivariant representation learning aims to capture variations induced by input transformations in the representation space, whereas invariant representation learning encodes semantic information by disregarding such transformations. Recent studies have shown that jointly learning both types of representations is often beneficial for downstream tasks, typically by employing separate projection heads. However, this design overlooks information shared between invariant and equivariant learning, which leads to redundant feature learning and inefficient use of model capacity. To address this, we introduce Soft Task-Aware Routing (STAR), a routing strategy for projection heads that models them as experts. STAR induces the experts to specialize in capturing either shared or task-specific information, thereby reducing redundant feature learning. We validate this effect by observing lower canonical correlations between invariant and equivariant embeddings. Experimental results show consistent improvements across diverse transfer learning tasks. The code is available at https://github.com/YonseiML/star.