# Soft Task-Aware Routing of Experts for Equivariant Representation Learning

arXiv: https://arxiv.org/abs/2510.27222 · HTML: https://mlsanity.com/paper/2510.27222/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 55 references that are in the map.
- First posted on arXiv 31 Oct 2025.

## Builds on (most cited first, 30 of 32)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks](https://mlsanity.com/paper/1506.01497/) (2015)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Describing Textures in the Wild](https://mlsanity.com/paper/1311.3618/) (2013)
- [Fine-Grained Visual Classification of Aircraft](https://mlsanity.com/paper/1306.5151/) (2013)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [An Empirical Study of Training Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.02057/) (2021)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Contrastive Multiview Coding](https://mlsanity.com/paper/1906.05849/) (2019)
- [AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty](https://mlsanity.com/paper/1912.02781/) (2019)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Do Better ImageNet Models Transfer Better?](https://mlsanity.com/paper/1805.08974/) (2018)
- [TADAM: Task dependent adaptive metric for improved few-shot learning](https://mlsanity.com/paper/1805.10123/) (2018)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [Using Deep Learning for Image-Based Plant Disease Detection](https://mlsanity.com/paper/1604.03169/) (2016)
- [Scaling and Benchmarking Self-Supervised Visual Representation Learning](https://mlsanity.com/paper/1905.01235/) (2019)
- [Equivariant Contrastive Learning](https://mlsanity.com/paper/2111.00899/) (2021)
- [Improving Transferability of Representations via Augmentation-Aware Self-Supervision](https://mlsanity.com/paper/2111.09613/) (2021)
- [Self-supervised learning of Split Invariant Equivariant representations](https://mlsanity.com/paper/2302.10283/) (2023)
- [Structuring Representation Geometry with Rotationally Equivariant Contrastive Learning](https://mlsanity.com/paper/2306.13924/) (2023)
- [Neural Experts: Mixture of Experts for Implicit Neural Representations](https://mlsanity.com/paper/2410.21643/) (2024)

## Abstract (arXiv)

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.