Title: Soft Task-Aware Routing of Experts for Equivariant Representation Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: equivariant representation learning

self-supervised learning

contrastive learning
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 \textbf{S}oft \textbf{T}ask-\textbf{A}ware \textbf{R}outing (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 \url{https://github.com/YonseiML/star}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses jointly learning invariant and equivariant representations in self-supervised learning. The authors argue that treating them as fully separate leads to redundancy. Instead of using independent projection heads, they introduce a task-aware routing mechanism that dynamically routes features through shared and task-specific paths. This allows the model to better capture both shared and unique aspects of each objective, improving generalization and reducing representational overlap. Their method shows performance improvements over classification, few-shot learning, and object detection tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The experiments are comprehensive.\n- The experiments cover the claims made in the introduction.\n- Ablation studies show the effectiveness of the introduced components over previous work.\n\nWeaknesses:\n- The introduction section needs to be revised to clearly explain what this paper refers to as invariant and equivariant concepts. I initially assumed that based on figure 1 and L21, the invariance is with respect to a task but there is clear explanation and no buildup or flow in this section. One paragraph is not enough to explain your motivation and intuition before diving to specific examples and details in paragraph 2.\n\n- Following from the previous point, the second paragraph starts with a description of another work that is too technical for the introduction section. Same as before, the concept of invariance and equivariance are being referred to at the same time as invariant and equivariant tasks alongside features. So the introduction needs an overhaul to make it clear.\n\n- About statement in L37, “This not only results in inefficient use of model capacity, but also hinders each head’s ability to specialize in task-specific features.” Say each head learns both shared knowledge and task-specific knowledge. In that case, regardless of redundancy, you still have task-specific features. I’m assuming you’re making a point about heads not being able to only attend to task-specific knowledge and having to share their representational power with shared knowledge but L37 does not explain that clearly and it needs a revision.\n\nQUESTIONS:\n- What would be the point of a separate equivariant learning mechanism? Shouldn’t that be part of the experts’ job? As in, the model should assign an expert to learn the observed shifts in the dataset as its role?\n\n- Since a topic of focus is the separation of heads by adding a shared projector, would you expect forcing disentanglement between representations of separate heads to cause minimal degradation due to the results in figure 3.a?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the interesting problem of representational redundancy in self-supervised learning methods using both invariant and equivariant objective functions. The authors point out that predicting semantic category and predicting augmentation may require similar features. Training projection heads for each task independently may therefore cause similar features to be encoded multiple times in the representation, reducing overall model capacity. The authors propose a mixture-of-experts based gating mechanism in which several shared projection heads are trained simultaneously. Objective-based gating networks produce softmax distributions over the experts to appropriately weight the most fitting projections for each task. This, in addition to an equivariant learning mechanism using a residual connection, yields representations that produce superior empirical results to comparable equivariant learning methods. Additional analysis of the networks trained using their mechanism demonstrate some evidence that redundancy among projection heads is reduced can lead to poorer performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- The paper is well written, easy to follow, and has easy to interpret figures.\n- A large number of experiments are run, combining various tasks, datasets, evaluation types (in-domain, out-of-domain), and baselines.\n- The additional analyses presented in section 5.2 help the reader understand how different kinds of experts naturally emerge using the new mechanism.\n- The practicality of the method is considered, with Figure 5 demonstrating that the new method can produce stronger results in a shorter time.\n\nWeaknesses\n- [W1] The empirical gains from using the method are often quite small.\n- [W2] Not all comparisons may be fair. Are all encoders trained for 200 epochs (Section 5, setup)? If so, this would mean that the author’s method gets a significant computational advantage, since the method “incur(s) a higher cost per epoch”.\n- [W3] It is not clear why the method could actually hurt in-domain performance under some circumstances.\n- [W4] I’m not entirely convinced that redundancy/canonical correlation amongst the projection from the experts implies redundancy in the un-projected representation. See questions.\n\nQUESTIONS:\n- [W2] Are all encoders trained for 200 epochs (Section 5, setup)? Figure 5 potentially addresses this, but what about the other experiments in the paper?\n- [W4] Is the argument about representational redundancy concerning the encoder f(x) itself, or the outputs of the projection heads? Couldn’t it be the case that both the invariant and equivariant projection heads could use the same feature in the representation f(x) without that feature being encoded multiple times in the representation? Is the issue then concerning the capacity of the projection heads themselves?\n- [W3] Is there any justification for why the new method can hurt in-domain performance?\n\nLIMITATIONS:\nPending the weaknesses mentioned previously, the authors adequately discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis works introduces the use of Mixture of Experts module in place of the standard MLP projector in a SSL framework. The authors present this framework as a mechanism to support equivariant representation learning as the MoE allows for specialisation to invariant and equivariant tasks either independently or as a shared space of the two. They introduce a series of constructive losses to enforce both representation properties, with the equivariant component enforced with a predictor network. The authors report strong invariant performance on a variety of benchmark tasks, and demonstrate robustness in and out of distribution. The ablation studies perform an investigation into the behaviour of the MoE and its impact on equivariant loss.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:\n+ The investigation of equivariant SSL is timely and important to the computer vision community.\n+ The method presented is sensible, seemingly correct, and simple. It is presented as general purpose method that can be adapted to a variety of data domains, architectures and SSL transformations. \n+ The invariant representation quality is well explored with a wide variety of experimentation, tests for robustness and generalisation.\n+ The investigation into redundancy is a nice addition, and demonstrates well that the experts are able to specialise and improve performance without wasted resource.\n\n**Weaknesses**:\n- The main argument made at the outset of the paper is that using separate projection heads limits the expressivity of the network and limits shared information. However, your findings in Figure 3 (a) somewhat contradict your conjecture. In fact it can be seen that from the expert gating that experts are largely specialised to either invariant or equivariant, with only one being shared, but not balanced. Therefore, one could argue that the improved results are due to the increased capacity for the network to learn specialised equivariant and invariant features given MoE. It would be nice to prove that simply having a projector network for Transformation type in T does not reach the same performance, given that T is accessible during training.\n- The evaluation does not analyse or report any quantitative measures for equivariant representations. Therefore, there is no way to know if the objective function is actually learning equivariant representations. Some further analysis of the learnt representation space is needed before the claim of equivariance can be substantiated. For example, all experimentation evaluates the invariance of the representations, yet does not evaluate explicitly the transformation information in the representations.\n- The novelty of the approach is arguably limited, this method is essentially akin to EquiMod, or SEN but with a MoE as a projector. While I am keen for these findings to be distributed to the wider community, the novelty is narrow. \n- While it is nice to show faster convergence with wall time, it would be beneficial to see the memory consumption, iterations per second, FLOPs as a more quantitative evaluation of computer resources.\n- There are some omitted details in the manuscript that have been outlined in the questions, these is a minor concern that will improve the reproducibility of the work if included.\n\nQUESTIONS:\n- How does the k-nn retrieval perform on the output of each expert? Do you get specialised transformation encodings?\n- $\\psi$ from figure 2 doesn’t seem to be defined in the text, I assume this is some network to encoder the transformation parameters.\n- Is $\\phi_{T}$ a mlp predictor or a different type of network? This is not defined.\n- Why is there a drop in performance as a factor of training time with SimCLR (Figure 5), from experience with SimCLR this behaviour is unusual and instead the convergence behaviour of equimod is expected. While you state overfitting, the mean accuracy is low. Can you also confirm which dataset this is performed on?\n\nLIMITATIONS:\nLimitations are addressed, but could be expanded to include task specifics and broader implications.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new Invariant-Equivariant SSL paradigm that combines the invariant and equivariant losses through a mixture of experts’ approach. This allows the model to learn both components, but removes the redundant information across the two projector heads. Some of the learned projector heads become specialized for one task while others are shared across the tasks. The authors present experimental results on STL10 and Imagenet100 pretrained models, testing a wide range of out-of-domain classification tasks, and find that their method improves out of domain classification.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n* Incorporating MMoE into the projection model to select the “task” by using a gating network seems quite interesting as a way to mitigate the dependence on the lambda term for the two losses. \n* The paper presents results with other invariant learning objectives (Table 6) to show that the results generalize in other SSL frameworks. \n\n**Weaknesses**\n* The paper motivation states that the invariance and equivariant components are encoding “redundant” information, however this statement is not quantified. It would be nice to have seen this quantified in some way for the comparison models, and then show that it is improved for the presented models. The k-NN retrieval is qualitative and not quantitative. \n* The analysis of the “expert specialization” is performed only on the STL10 network, and the results seem somewhat qualitative and difficult to interpret. \n* Although the results for out-of-domain classification seem very reasonable and consistent, suggesting the model is doing something interesting, the margin of differences across models is very small for these OOD experiments (and has been reported to disappear for the equivariant models when trained on ImageNet1000, see below).\n\nQUESTIONS:\n* I found Figure 1 confusing. Naively, I would think the L/R flip of the objects would also be preserved as an augmentation, but this stands out as something that the invariance-only model is doing better at preserving than the equivariant model.  Could this be explained, or could the figure be made clearer? For instance, what would hypothetical images be that visually contain the augmentation-aware information but do NOT contain semantic structure? \n* The paper should cite and discuss CE-SSL from last years’ NeurIPS, which also combines equivariant and invariant losses (this is the paper that also talks about ImageNet1000 vs ImageNet100 training for equivariant model evaluation) https://proceedings.neurips.cc/paper_files/paper/2024/file/ae28c7bc9414ffd8ffd2b3d454e6ef3e-Paper-Conference.pdf \n* The authors talk about “representational redundancy” and how it is reduced by including a shared projector in the invariant and equivariant projectors, however, to my knowledge, these projectors are typically discarded anyway and the “representation” that is typically used is the backbone (which is shared across the different projectors). How does this learned representation (the backbone representation) relate to the assumption of independence of the projectors? Naively, I am not sure why having completely independent projectors should encourage one to learn a better representation, and as the text is currently written, the more general motivation for benefits of this change are not clear.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}