Soft Task-Aware Routing of Experts for Equivariant Representation Learning

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Jeon, Jaebyeong, Jang, Hyeonseo, Sohn, Jy-yong, Lee, Kibok
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915588436131840
author Jeon, Jaebyeong
Jang, Hyeonseo
Sohn, Jy-yong
Lee, Kibok
author_facet Jeon, Jaebyeong
Jang, Hyeonseo
Sohn, Jy-yong
Lee, Kibok
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_27222
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Soft Task-Aware Routing of Experts for Equivariant Representation Learning
Jeon, Jaebyeong
Jang, Hyeonseo
Sohn, Jy-yong
Lee, Kibok
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
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.
title Soft Task-Aware Routing of Experts for Equivariant Representation Learning
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2510.27222