Soft Task-Aware Routing of Experts for Equivariant Representation Learning

AI summary

Overview

Research area: Self-supervised learning (SSL), specifically invariant and equivariant representation learning.

Technical level: Intermediate. The core idea is intuitive, but readers benefit from familiarity with contrastive learning, projection heads, and mixture-of-experts concepts.

Scope: This paper introduces a routing strategy called Soft Task-Aware Routing (STAR) that restructures the projection heads in invariant–equivariant SSL to share information and reduce redundancy, validated across 11 transfer learning benchmarks plus object detection and few-shot tasks.

What This Paper Is About

Modern self-supervised methods often learn two things at once: an invariant representation that ignores augmentations (capturing "what" is in an image) and an equivariant representation that tracks them (capturing "how" the image was transformed). The standard approach uses a shared encoder with two separate projection heads—one per objective—implicitly treating the two tasks as unrelated. The authors argue this is wrong: invariant and equivariant learning share substantial information (e.g., knowing a shape is a crater helps infer lighting direction, and vice versa), so two independent heads redundantly relearn the same features and waste model capacity. STAR fixes this by making the projection heads cooperate.

Key Contributions

1. Diagnosis of redundant feature learning: The paper empirically shows that invariant and equivariant objectives are interdependent, and that standard two-branch designs with separate projection heads encode overlapping shared information, leading to wasted capacity and weaker gradients.

2. STAR-SS (Single Shared projection): A simple design adding a third shared expert whose output is added to both the invariant and equivariant embeddings, forcing it to capture information useful to both tasks.

3. STAR-MMoE (Mixture-of-Experts projection): A more flexible design using multiple shared experts and two task-specific routers (invariant and equivariant) that softly weight experts per input, enabling dynamic specialization.

4. Empirical validation of reduced redundancy: The authors quantify redundancy via canonical correlation between experts and show that lower redundancy correlates with better generalization.

Main Findings

• Consistent transfer gains: On 11 out-of-domain datasets, STAR-MMoE achieves best or near-best results, topping 7/11 datasets with STL10-pretrained ResNet-18 and 10/11 with ImageNet100-pretrained ResNet-50.

• Object detection improvements: On PASCAL VOC07+12 with a frozen R50-C4 backbone, STAR-MMoE achieves the highest AP (48.85), AP50 (76.81), and AP75 (53.01) among compared methods.

• Few-shot superiority: Best performance across every 5-way 1-shot and 5-way 5-shot setting on FC100, CUB200, and Plant Disease.

• In-domain performance preserved: In-domain accuracy on STL10 (86.74) and ImageNet100 (84.83) matches or exceeds baselines, showing gains aren't at the cost of the pretraining distribution.

• Clear expert specialization: Routing weight analysis shows some experts serve mostly invariant objectives, some mostly equivariant, and one shared expert gets balanced weights from both.

• Reduced canonical correlation: Experts assigned to the same objective correlate more strongly with each other than with experts of the other objective; more experts generally means lower mean canonical correlation and higher downstream accuracy.

• Higher-quality gradients: Gradient cosine similarity between experts drops from 0.59 (EquiMod) to 0.30 (STAR-MMoE), and experts converge faster—providing stronger, more task-specific signals to the backbone.

• Stronger equivariance: STAR-MMoE reaches 0.98 R-equivariance and 0.27 P-equivariance, outperforming SimCLR, AugSelf, EquiMod, and CARE.

• Training efficiency: STAR reaches high accuracy in less wall-clock time than baselines, which start saturating or degrading after ~300 epochs.

Methodology in Plain English

The setup: take images, apply two random augmentations to each, and pass both through a shared encoder. Then comes the projection stage, where the design diverges from prior work.

• STAR-SS uses three experts instead of two. An invariant expert, an equivariant expert, and a shared expert. The invariant embedding is the sum of the invariant expert's output and the shared expert's output; the equivariant embedding is the sum of the equivariant expert's output and the shared expert's output. The shared expert is pulled toward both tasks automatically.

• STAR-MMoE replaces this with a bank of N experts and two routers (one per task). Each router outputs softmax weights over the experts, so each embedding is a weighted combination of all experts. Experts that help both tasks get high weight from both routers; task-specific experts get high weight from only one. Soft (non-sparse) routing is used deliberately because batch normalization inside experts makes top-k routing unstable in SSL.

Equivariant learning works by predicting the shift caused by an augmentation in embedding space: a small predictor takes the original encoding plus a projected augmentation parameter and predicts the shifted embedding, trained with an InfoNCE loss against the actual augmented embedding. Invariant learning uses standard InfoNCE between the two augmented views.

Crucially, the projection heads and experts are discarded after pretraining—only the encoder transfers downstream, eliminating the usual MoE inference-time routing burden.

Why This Matters

Impact on research: This paper challenges a tacit assumption in equivariant SSL—that invariant and equivariant objectives should be handled by independent heads. It offers a principled, empirically validated alternative and introduces canonical correlation as a diagnostic for measuring redundancy in multi-objective representation learning.

Real-world applications:

• Medical imaging, where scans vary in orientation, lighting, and contrast—models need both stable semantic recognition and awareness of acquisition conditions.
• Satellite and aerial imagery, where viewing angle, shadow direction, and orientation carry meaning (e.g., distinguishing terrain features).
• Robotics and autonomous driving, where objects appear under varied poses and lighting and downstream tasks depend on both identity and pose.
• Industrial inspection, where defects must be recognized regardless of orientation but the orientation itself can be diagnostic.

Industry relevance: Because only the encoder is transferred, STAR avoids the inference-time overhead that normally limits MoE deployment. This makes it practical for any pipeline where a pretrained backbone is fine-tuned or linearly probed—standard practice at companies relying on transfer learning. The wall-clock efficiency gains also matter for teams with limited compute budgets.

Future Directions

• Sparse routing without batch normalization instability: The authors note soft routing was chosen because top-k routing destabilizes training under batch norm. Finding normalization schemes that permit sparse, more efficient routing is an open problem.

• Scaling to larger backbones and datasets: Experiments use ResNet-18 and ResNet-50 on STL10 and ImageNet100. Whether STAR's benefits hold for ViTs, larger expert counts, or web-scale pretraining is untested.

• Alternative routing criteria: Routers currently condition only on the input feature. Conditioning on augmentation parameters, task difficulty, or training progress could yield finer specialization.

• Theoretical grounding for the redundancy–generalization link: The paper shows correlation between lower redundancy and better accuracy, but a formal account of when and why reducing canonical correlation improves downstream transfer remains open.

Target Audience

Machine learning researchers and graduate students working on self-supervised learning, equivariant representation learning, or mixture-of-experts architectures. Practitioners who pretrain backbones for transfer learning—especially in vision—will also benefit, since the method is architecture-agnostic at the projection-head level and only requires modifying pretraining, not downstream pipelines. Readers without familiarity with contrastive learning or InfoNCE loss should first review a SimCLR-style objective before diving in.

Authors’ 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.