Soft Task-Aware Routing of Experts for Equivariant Representation Learning

W1 & W2. 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. 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.

Regarding the reviewer’s comment on the conceptual clarity of invariant and equivariant learning, we will clarify these concepts in the context of our framework as follows:

• Invariant learning (or task) aims to extract features that remain unchanged under input transformations. It filters out transformation-specific signals while preserving semantic information.
• Equivariant learning (or task), in contrast, seeks to retain transformation-induced variations in the representation. The goal is to ensure that the representation changes in a structured and predictable manner with respect to the transformation applied.

We clarify that the terms “invariant task” and “equivariant task” refer to pretext tasks that induce these respective properties in the learned representations, not to tasks that are themselves invariant or equivariant. This clarification serves as a conceptual foundation that leads naturally into the second paragraph, where we introduce our modeling strategy and discuss baseline comparisons. We believe this revision significantly improves the clarity and coherence of the introduction.

W3. 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.

We thank the reviewer for the accurate interpretation. We will revise the sentence to better reflect that redundancy reduces the effective capacity each head can dedicate to truly task-specific signals. As the reviewer pointed out, even when using two independent heads, each can still capture both task-specific features and shared knowledge. The core issue, however, lies in the redundancy: both heads must re-encode shared information, which leads to inefficient use of the representational space. Our intention in L37 was to emphasize that this redundancy reduces the effective capacity each head can dedicate to truly task-specific signals. We will revise the sentence to better reflect this point.

Q1. 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?

We understand the reviewer’s question as asking whether the explicit equivariant learning mechanism is necessary, given that the MoE gating could in principle assign experts to learn augmentation-induced shifts as part of their specialization. Based on this understanding, we provide the following response.

1. Gating Alone Does Not Induce Equivariance

We agree that the mixture-of-experts (MoE) mechanism encourages each expert to focus on different portions of the data distribution. However, gating only selects which expert to use; it does not constrain the functional role of each expert. In particular, gating does not ensure that the selected expert encodes augmentation-specific transformations in a consistent and structured manner without equivariant objective.

2. Why An Explicit Equivariant Objective Is Necessary

The invariant loss encourages the model to remove augmentation-induced variations to enforce semantic consistency. This removes the signals, such as color and pose, needed to learn augmentation-aware representations. In contrast, the equivariant objective provides an explicit inductive bias that encourages experts to preserve these signals in the embedding space by modeling transformation-consistent feature shifts.

3. Gating Without Equivariant Loss

To isolate the effect of the equivariant objective, we performed an ablation where we removed the equivariant loss while keeping the MoE architecture with one gating network. The result shows that gating alone is insufficient for learning robust, generalizable features:

Despite identical architecture (except for gating network) and capacity, the performance drops significantly without the equivariant loss. This suggests that without an equivariant loss, the experts do not specialize in modeling augmentation-specific information, even when gating is present.

We hope this addresses the reviewer’s question. If we misunderstood the question, we kindly ask the reviewer to rephrase or elaborate so we may provide a more precise response.

Q2. 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?

The observed expert-wise separation in Figure 3(a) suggests that additional constraints enforcing disentanglement across heads would have limited effect, as our method already induces clear task-specific specialization via the MMoE gating mechanism.

To verify this, we conducted two complementary experiments:

• Orthogonality regularization: We explicitly enforced orthogonality across expert outputs and observed only a small drop in out-domain classification accuracy from 53.07 to 52.75, confirming that he additional orthogonality constraint introduces only minor adjustments to the parameters.
• Hard assignment: We replaced soft gating with hard expert selection (1 shared + 5 invariant + 2 equivariant), matching the effective expert allocation observed in Figure 3(a). The resulting model showed only a minor performance drop from 52.37 (MMoE with 8 experts) to 52.27, indicating that the learned gating already approximates hard selection.

These results support the conclusion that our MMoE structure effectively promotes disentangled, task-specific representation learning, even without explicit constraints.

We hope our response aligns with the reviewer’s intent, and we would welcome any clarification if we have misunderstood the question.

[W1] The empirical gains from using the method are often quite small.

We acknowledge the reviewer’s concern. Our method shows clear improvements across multiple benchmarks, and the performance gains are comparable to those of prior methods. For example, on STL10, we achieve a 7.1% relative gain over the previous best (EquiMod) while attaining the highest overall accuracy. On ImageNet100, we observe 1.3% gain over EquiMod.

Similar trends are observed in the few-shot classification and object detection tasks (see Table 3-4), where our method consistently outperforms the strongest baseline. These results collectively demonstrate the effectiveness of our approach across diverse downstream tasks.

[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?

Unless otherwise stated, all encoders are trained for 200 epochs. We chose this schedule because 200 epochs allows the baselines to reach their best performance and yields reasonable performance for our method as well.

Importantly, as shown in Figure 5, our methods already outperforms the strongest baseline even when trained for significantly fewer epochs, and its performance continues to improve with longer training.

A similar trend is observed in few-shot classification and object detection. In particular, under the few-shot classification setting, our method, trained with only 100 epochs, still outperforms the prior best method (AugSelf) on STL10-pretrained ResNet-18. Specifically, while AugSelf achieves an average accuracy of 53.54% (Table 2), our method reaches 55.27%, yielding a relative improvement of 3.23% despite fewer training iterations. This result highlights the sample efficiency and strong transferability of our approach.

[W3] Is there any justification for why the new method can hurt in-domain performance?

We attribute the slight drop in STL10 in-domain accuracy to the difference in dataset complexity and the effect of increased model capacity.

Compared to EquiMod, our method adopts a deeper predictor to better model complex augmentation-induced shifts, which enhances out-domain generalization. However, in low-diversity datasets like STL10, where semantic variation is relatively limited, the added capacity may introduce excessive degrees of freedom, leading to unnecessary feature dispersion and marginal degradation in in-domain accuracy.

To support this explanation, Table 5 presents an ablation using a shallower predictor (depth = 1). In this setting, the in-domain performance improves to 87.41%, confirming that deeper predictor can be a factor in performance drop under simpler data conditions.

Importantly, in the context of SSL, the quality of learned representations is typically assessed by their transferability rather than in-domain performance. From this perspective, our architectural choice prioritizes generalization, as evidenced by consistently strong out-domain classification results.

[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? We clarify that the redundancy we discuss arises at the projection head output level. However, this head-level redundancy negatively impacts the encoder. This may be harmless at inference, but during training, independent heads must each disentangle both shard and task-specific factors from entangled backbone features. This is inherently difficult and often results in suboptimal gradients being propagated back into the encoder, degrading representation quality.

We clarify that the redundancy we discuss arises at the projection head output level. However, this head-level redundancy negatively impacts the encoder.

This may be harmless at inference, but during training, independent heads must each disentangle both shard and task-specific factors from entangled backbone features. This is inherently difficult and often results in suboptimal gradients being propagated back into the encoder, degrading representation quality.

In contrast, explicitly separating shared and task-specific pathways (e.g., via a shared projector or MMoE gating) allows heads to specialize more easily. They converge faster and provide more reliable gradients. We confirmed this empirically: the Frobenius norm of differences in projector weights every 10 epochs decreased more quickly with our method than with EquiMod, indicating more stable optimization.

Finally, this issue cannot be explained by projection head capacity alone. Even when we increased the number of projectors from 2 to 8 while reducing the hidden dimension (512 →180) to keep the parameter count approximately the same, applying MMoE Projection still improved the performance from 49.54 to 50.34. This demonstrates that the performance gain comes from the inductive bias introduced by MMoE, not simply from increasing capacity.

W1. The main argument is that separate projection heads limit expressivity of the network and limits shared, but Figure 3(a) seems to contradict this. One could argue the gains come from increased capacity to learn specialized features. It would be nice to show that a projector per transformation type T would not match the same performance.

The imbalance observed in Figure 3(a), where only one expert is shared across tasks, does not contradict our conjecture. Our argument is not that experts must be balanced between task, but that having two independent projectors leads to redundant learning of shared information. In MMoE Projection, the observed specialization naturally emerges during optimization as the model allocates experts according to the invariant and equivariant objectives while avoiding repeated extraction of shared factors.

To address the possibility that the improved results simply stem from increased capacity, we conducted an ablation study using multi-projector baseline with six dedicated projection heads: one for invariance and the others for five transformation types. For fairness, the MMoE variant was also trained with six experts to match model capacity.

Despite identical capacity, MMoE Projection shows a relative improvement of 5.5% out-domain classification performance than the transformation-specific multi-projector baseline (52.35 vs. 49.61). This confirms that the gain arises not merely from increased capacity, but from the structured inductive bias and task-aware routing introduced by MMoE.

W2. The evaluation does not analyse or report any quantitative measures for equivariant representations. Some further analysis of the learnt representation space is needed before the claim of equivariance can be substantiated.

Our paper provides clear evidence that the learned features preserve augmentation-aware (i.e., equivariant) information, following the evaluation paradigm established in prior work such as AugSelf [3], where color consistency is used as a key qualitative indicator of equivariance. Based on this criterion, we provide two complementary pieces of evidence:

• Figure A.2: Retrieves neighbors with similar color attributes in k-NN retrieval.
• Table 1: Achieves the highest accuracy on the color-sensitive Flowers102 dataset.

These results indicate that our method effectively captures transformation-sensitive features.

To further strengthen this claim, we add explicit quantitative evaluations of transformation-aware information in the learned representations:

1. Transformation Parameter Regression We design a regression task using a frozen backbone and a single linear layer to predict augmentation parameters given the original and augmented images. We evaluate performance using Mean Squared Error (MSE).

• For Crop, we predict scale and aspect ratio.
• For Color, we predict brightness, contrast, saturation, and hue.
• When both are applied, we predict all parameters jointly.

MMoE achieves the lowest MSE across all transformation types, indicating that its representations more accurately encode transformation-aware information.

2. Transformation Type Classification We also perform a binary classification task to distinguish whether the applied transformation is Crop or Color, using the same frozen representation and a linear head.

MMoE attains the highest classification accuracy (95.76%), suggesting that it preserves clear discriminative signals between different transformation types in the learned representation.

These results collectively provide both quantitative and qualitative support for the claim that our method captures transformation-aware features in the learned representations.

W3. 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.

Our method proposes projection module designs that addresses the limitations of prior work such as EquiMod [1] and SEN [2] by explicitly handling redundancy and improving specialization:

• Diagnosing the redundancy problem between two independent heads in Equivariant SSL.
• Structuring the projection space to separate shared and task-specific components.
• Leveraging multi-gate routing aligned with invariance and equivariance objectives.

These designs lead to better expert specialization and improved downstream performance, offering both conceptual and practical advances over prior work.

W4. 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.

Thank you for the suggestion. We report a detailed comparison of memory usage, FLOPs, and iteration speed in the table below. While MMoE Projection has marginally higher memory consumption and FLOPs than EquiMod, the values remain largely comparable.

Importantly, as shown in Figure 5, MMoE Projection achieves strong performance even with fewer training iterations, demonstrating that the model is not only effective but also sample-efficient. This confirms that its improved performance stems not from increased computational cost, but from architectural advantages.

W5. 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. (Q2 & Q3 included)

We are committed to ensuring reproducibility, and to that end, we have included not only the definitions of $\psi$ and $\phi_T$, but also all relevant hyperparameter settings and configuration details in Appendix B. To make the manuscript more self-contained, we will move key definitions such as those of $\psi$ and $\phi_T$ into the main text in the revised version.

Specifically, as defined in Appendix B.2:

• $\psi$ is a single linear layer that takes the augmentation parameter vector as input and expands it to match the output dimensionality of the expert representations. This design follows EquiMod and emphasizes the role of augmentation parameters in guiding the transformation.

• $\phi_T$ is an MLP predictor.

Q1. How does the k-nn retrieval perform on the output of each expert? Do you get specialised transformation encodings?

Our analysis suggests that the experts are not specialized by transformation type, but rather align more closely with objective-level distinctions (i.e., shared vs. invariant vs. equivariant), as indicated by their k-NN retrieval results (Figure A.3).

To analyze this, we conducted k-NN retrievals within each expert’s embedding space on the STL10 test set. The results show that:

• Expert 1 (shared) retrieves semantically consistent images, reflecting general-purpose representation.
• Experts 2–6 retrieve samples invariant to changes induced by transformation, aligning with the invariant objective.
• Experts 7–8 retrieve samples with similar color tone, suggesting sensitivity to color-based transformations.

This suggests that the experts specialize based on the nature of the objective, rather than being explicitly aligned with specific transformation types.

Q4. 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?

The performance drop with longer training in Figure 5 is observed because we evaluate out-domain generalization, not in-domain performance. As shown in Figure A.1(a) in Appendix, the in-domain accuracy consistently improves with longer training, while out-domain accuracy degrades. The out-domain results in Figure 5 are evaluated on a collection of 11 out-domain datasets (details in Appendix B.1). This indicates that the drop is not a convergence problem but instead reflects a loss of generalization, likely stemming from the lack of explicit inductive bias in SimCLR for modeling augmentation-aware structures.

We will consider including the corresponding in-domain performance curve (currently in Figure A.1(a)) in the main text in the revised version for clarity.

Reference

[1] Devillers and Lefort. EquiMod: An Equivariance Module to Improve Self-Supervised Learning. In ICLR. 2023.

[2] Park et al. Learning Symmetric Embeddings for Equivariant World Models. In ICML. 2022.

[3] Lee et al. Improving Transferability of Representations via Augmentation-Aware Self-Supervision. In NeurIPS. 2021.

Thank you to the authors for their clarifications on these points, and apologies for the misunderstanding of Q1 and Q4.

My only concern still lies with the evaluation of equivariance, while results are presented on colour specific tasks and that there is some retrieval based evaluations, this is not necessarily cause for determining equivariance. Instead it can be argued that the work is augmentation aware, not necessarily equivariant without providing theoretical proof, or at the very least some indication that the latent space can be predictably transformed inversely to the input transformations.

I understand this may seem pedantic but the definition of equivariance here is not substantiated.

Given that my other concerns have been addressed, I will increase my score.

W1. 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.

We already quantify representational redundancy using the canonical correlation between the heads, as shown in Figure 4. This metric captures the linear dependency between representations produced by different heads and provides a direct measure of redundancy across projection spaces. The connection between head-level redundancy and representation is further discussed in our response to Q3.

As reported in Figure 4(b), EquiMod [1], which uses two independent projectors for invariant and equivariant learning, exhibits high Canonical Correlation (0.5830), indicating substantial overlap in learned embeddings. In contrast, our methods, including Single Shared Projection (0.5383) and MMoE Projection with 16 experts (0.3948), show significantly lower Canonical Correlation. This corresponds to a 7.6% and 32.3% relative reduction, respectively, reflecting reduced redundancy and enhanced specialization across heads.

W2. The analysis of the “expert specialization” is performed only on the STL10 network, and the results seem somewhat qualitative and difficult to interpret.

We further confirmed expert specialization patterns in our analysis on ImageNet100, which also show clear and interpretable expert allocation. We will include the ImageNet100 results in the revised version.

Figure 3 provides both quantitative and qualitative evidence of expert specialization.

Specifically, Figure 3(a) visualizes the gating weight distribution across experts, showing that during training, the model naturally assigns distinct roles: 1 shared expert, 5 invariant experts, and 2 equivariant experts among the 8 experts. The Min/Max ratio shows how evenly each expert is used across invariant and equivariant learning tasks. For instance, Expert 1, identified as the shared expert, has a Min/Max ratio of 0.34, indicating that it is assigned to both tasks relatively evenly. This balanced assignment encourages the expert to capture shared information beneficial to both objectives.

Figure 3(b) and Figure A.3 in Appendix visualize the k-NN retrieval results based on each expert’s embedding. Invariant experts retrieve samples with similar shapes, whereas equivariant experts tend to focus on color similarity. The shared expert, on the other hand, retrieves semantically consistent samples within the same class. These results indicate that different experts attend to distinct transformation sensitivity and semantic attributes, providing clear evidence of expert specialization.

W3. 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).

As commonly reported in the SSL and Equivariant SSL literature, performance gains on OOD tasks are often marginal.

For instance, SimSiam [2] reports similarly small margins (<1% absolute) over the prior best (e.g., MoCo v2) on VOC/COCO transfer learning tasks, yet these consistent gains are considered indicative of improved generalization.

Q2. 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)

We thank the reviewer for highlighting CE-SSL [3]. CE-SSL introduces a contrastive equivariant SSL learning framework that explicitly adds equivariance-related objective to complement standard invariant SSL losses. Unlike our method, which proposes a new projection module design employing MMoE to reduce redundancy between projectors, CE-SSL [3] addresses the problem at the loss-design level.

We will revise the manuscript to cite and discuss this work. Regarding the findings on ImageNet1000 vs. ImageNet100, we were unable to include a verification in the main text due to resource limitations, but we will consider exploring whether similar trends hold in our setup if resources allow.

Q1. 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?

We first note that Figure 1 illustrates the existence of shared information between representations learned exclusively with the invariant or equivariant objective. Both models retrieve samples from the same class, indicating the shared information (e.g., semantic structure), while the equivariance-only model additionally preserves augmentation-related information, suggesting that the two objectives rely on a common set of information while capturing complementary aspects.

Nonetheless, we agree that the LR flip shown in retrieval results may be confusing. We conjecture that this implies that the LR flip is considered less important when measuring similarities on learned representations. To confirm that the LR flip is also captured in learned representations, we extended the experiment with LR flipped STL10 test dataset, essentially doubling the number of test data.

We observed that the equivariance-only model retrieves samples in a pattern that highlights flip-aware behavior:

• 1-NN: A flipped version of the query image.
• 2-NN: A sample with a similar pose to the query.
• 3-NN: A flipped version of the second image.
• A similar trend is observed in the pair of 4- and 5-NN, and so on.

In addition, the equivariance-only model retrieves a greater number of neighbors that share both consistent color and similar pose compared to the invariance-only model.

This pattern suggests that the model captures flip-related variation in representation. In contrast, the invariance-only model tends to retrieve samples with random pose, highlighting its preference for transformation-agnostic similarity.

Q3. 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.

We would like to clarify how the independence or separation of the projection heads affects the quality of the learned backbone representation. Our analysis is based on two complementary perspectives:

1. Redundant gradients from independent heads interfere with task-specific learning.

When the invariant and equivariant heads are completely independent projectors, both must extract shared factors separately. This leads to redundant supervision on shared components, making the backbone repeatedly optimize for the same features. Consequently:

• Model capacity is used inefficiently.
• Task-specific gradients become weaker, as redundant gradients dominate.
• The expressivity of the backbone is limited.

We empirically confirmed this by computing the cosine similarity between gradients from different heads with respect to the backbone: it decreased from 0.59 (EquiMod [1]) to 0.34 (Single Shared Projection) and 0.30 (MMoE Projection), indicating reduced redundancy and improved specialization.

2. Explicit disentanglement leads to easier optimization and faster convergence.

When each head is forced to handle both shared and task-specific information simultaneously, optimization becomes more difficult. By introducing a shared pathway (e.g., a shared projector or MMoE gating), each head can specialize more effectively, leading to better optimization signals. This improves:

• Gradient quality early in training, when learning rate is high.
• Convergence speed and stability of head parameters.

To support this, we measured the Frobenius norm of weight changes in projectors every 10 epochs. Our method showed smaller changes compared to EquiMod [1], indicating faster convergence of each head.

These two analyses collectively demonstrate that our design not only reduces redundant supervision but also facilitates more effective and stable optimization, leading to improved representation.

Reference

[1] Devillers and Lefort. EquiMod: An Equivariance Module to Improve Self-Supervised Learning. In ICLR. 2023.

[2] Chen et al. Exploring Simple Siamese Representation Learning. In CVPR. 2021.

[3] Yerxa et al. Contrastive-Equivariant Self-Supervised Learning Improves Alignment with Primate Visual Area IT. In NeurIPS. 2024.