UniMRSeg: Unified Modality-Relax Segmentation via Hierarchical Self-Supervised Compensation

Dear Reviewer mFFe,

We sincerely thank you for your valuable comments and questions, which help us a lot to improve our work. We address your questions as follows.

[W1, the suggested framework is predominantly based on the existing self-supervised frameworks; the main novelty is in the sequential orchestration of them and not that of the novelties themselves.]

Due to the character limit in this rebuttal, we kindly refer you to our response to Reviewer sDPL's [Q1 W1], where we address these concerns in detail.

[Q1, the introduced three-stage complexity ... a more direct, one-shot, end-to-end training regime that consolidates all the suggested losses, reconstruction, contrastive, and consistency, as multi-tasks ...]
[W2, the training pipeline is also quite complex ... ]
[W4, a direct comparison with the end-to-end architecture ...]

Thank you for your suggestion. Prior to submission, we actually conducted the single-model unified training experiment. However, we chose not to include it in the paper, as our intention was to help readers focus more clearly on the self-supervised benefits brought by each stage individually. This is because the distinction between multi-stage and single-stage training essentially reflects a difference between self-supervised representation learning and multi-task learning.

(1) Self-supervised Representation Learning (Three Stages)

Our work is built upon a self-supervised pretraining framework, with training initialized from scratch starting at Stage 1. For each task (e.g., BraTS 2020), self-supervised training is conducted solely on the task’s own training set, rather than relying on externally labeled datasets such as ImageNet-pretrained weights for initialization. Our three-stage design includes:

• Stage 1: Pretrain both the encoder and decoder through arbitrary modality reconstruction.
• Stage 2: Pretrain the encoder with contrastive learning. The segmentation task here is only used to guide the contrastive objective, not as a final goal.
• Stage 3: Perform the downstream segmentation task.

This design explicitly aims to reduce the representation gap between complete and incomplete modalities in the encoder-decoder space during the final segmentation stage. If all these tasks are trained jointly in a single-stage model, it would fall into the scope of multi-task learning, rather than self-supervised pretraining.

(2) Multi-task Learning (Single Stage)

In multi-task learning, the goal is to leverage multiple complementary clues for collaborative learning. These clues can come from the data level [1,2,3], or from structural supervision types [4,5,6]. Most approaches often incorporate deliberate designs for task-sharing and task-specific components to enable effective joint prediction across tasks. In this work, if the three stages are forcibly merged into a single training stage, the following issues will arise:

• The model is trained with inputs that involve Random Modality Dropout, Random Modality Shuffle, and Random Spatial Masking. The encoder is supervised using the NT-Xent contrastive loss, an adapter is incorporated, and the model is simultaneously tasked with both segmentation and reconstruction. Such a fully entangled training setup lacks a clear task hierarchy, and there is no explicit coordination among the input, encoder, decoder, and output.
• When all tasks are treated equally without ordering, joint optimization becomes highly difficult. In the unified single-stage training attempt, the total loss involves six parts. The results on BraTS 2020 are as follow:

We observed the following phenomena during the experiments:

• The single stage model failed to converge, with loss plateauing early.
• The optimization process was highly unstable, with different losses fluctuating in turn and failing to decrease consistently together.

These phenomena and performace clearly indicate that unified single stage training is unable to achieve effective coordination among the various designs and supervision signals. The lack of a clear training order, combined with competing optimization objectives, leads to mutual interference.

In summary, we would like to clarify two potential misconceptions:

1. The three-stage training is not overly complex. Both theoretically and practically, it is a simpler and more stable design, especially when compared to directly mixing all objectives and modules in a single unified training phase, which leads to optimization difficulties.
2. The three stages are not repeating the same task. Instead, they form a self-supervised learning (SSL) pipeline that does not rely on extra data, and progressively reduce the gap between incomplete and complete modalities at the input, feature, and output levels to improve segmentation performance.

References:
[1] Segment Anything, ICCV 2023.
[2] SegGPT: Towards Segmenting Everything in Context, ICCV 2023.
[3] Spider: A Unified Framework for Context-dependent Concept Segmentation, ICML 2024.
[4] Deep Multitask Learning with Progressive Parameter Sharing, ICCV 2023.
[5] MTMamba: Enhancing Multi-Task Dense Scene Understanding by Mamba-Based Decoders, ECCV 2024.
[6] Multi-Task Dense Prediction via Mixture of Low-Rank Experts, CVPR 2024.

[Q2, ... What exactly is the reversed? What does it mean by explicitly compensate the weak perceptual semantics? Comparing this design with a more standard attention mechanism ... an ablation analysis would help in understanding the need of making the modification.]

(1) Definition of “Reverse” in Reverse Attention

"Reverse" refers to inverting the standard self-attention response. Traditional self-attention emphasizes high-activation (salient) regions. However, under missing modality conditions, some important semantic cues may receive low activation and be underrepresented. We compute reverse attention as 1 - self-attention to explicitly highlight these under-attended regions, enabling targeted semantic compensation.

(2) Purpose of Compensating Weak Perceptual Semantics

Ideally, if the features from incomplete and complete modalities are well aligned, their self-attention maps should exhibit similar activation distributions. In practice, due to missing inputs, incomplete modality features often lack activation in key semantic regions. These low-response areas typically contain complementary or suppressed information that standard attention overlooks. To address this, we introduce an adapter in Stage 3, trained to recover semantic cues from these weakly activated regions. Guided by reverse attention, the adapter focuses on residual perceptual gaps and compensates them explicitly. This enhances semantic completeness and improves prediction consistency under missing-modality scenarios.

(3) Ablation Study

We have conducted a detailed ablation analysis of the Reverse Attention Adapter (RAA) as shown in Table 4. The "mutual attention" setting corresponds to a standard attention mechanism implemented with a 3D Swin Transformer. The results show that removing reverse attention ("w/o Reverse Attention") leads to a significant performance drop for UniMRSeg, while the performance gap between "w/o Reverse Attention" and "w/o Reverse Attention + Mutual Attention" is minimal. This strongly supports the effectiveness of the proposed reverse attention design.

In addition, we refer you to our response to Reviewer sDPL’s [Q1 W1] "(3) Lightweight Reverse Attention Adapter (RAA)", which further clarifies the motivation and advantages of the RAA.

[Q3, It would be useful to have these ablations in order to isolate contribution of individual stages better ...]
[W3, This means that the ablation insights provided by the authors ... cannot explain the need of such ordered pipeline.]

We fully agree on the importance of isolating the contribution of individual stages. Accordingly, we have already conducted detailed ablation experiments covering each stage independently, as well as all their combinations, which are presented in Table 5 of the submitted version. We kindly invite you to check it.

[Q4, ... Had the baseline models been trained end-to-end as per their papers or had they been trained under a multi-stage or pre-training regime to be fairly comparable? ... performance provided by UniMRSeg could, in some part, be attributed to large pre-training and not only to architectural or methodological advances.]

First, we would like to clarify that UniMRSeg does not use any external data in any stage of training. The different UniMRSeg models used for various segmentation tasks (Brain Tumor Segmentation / RGB-D SOD / RGB-T SOD / RGB-D Semantic Segmentation) are trained on the corresponding task’s training set in all three stages.
Therefore, the performance gains of UniMRSeg are not attributable to large-scale pretraining, but rather to the proposed architectural and learning design.

Second, to ensure fair comparison across methods in Tables 1, 2, 6, and 7, all baseline models are evaluated under the same training dataset splits. Specifically:

• If a method reports results using the same dataset splits as ours, we directly cite the numbers from the original paper.
• If the results are missing (e.g., for missing modality settings), or if different splits are used, we retrain the models using their released code, following the original training protocols as closely as possible.

This ensures that all results are evaluated under consistent and fair condition.

Dear Reviewer mFFe,

Thanks for your constructive comments and valuable suggestions to improve this paper.

As we are entering the final days of the author-reviewer discussion period, we are not sure whether our rebuttals have addressed your concerns. We would be greatly appreciated if you could spend some of your time for a reply.

Thank you again for your time and consideration.

Best regards,

The Authors

Dear Reviewer ife8,

Thank you very much for your valuable comments and suggestions, which are greatly helpful for improving both our current and future work. We address your questions as follows.

[W1, authors mainly focus on the analysis of modality combinations. It might be insightful to show the semantic importance. Whether some modalities (e.g., Flair vs T1ce) are inherently more informative or complementary to others.]

Thank you for this insightful comment. We agree that understanding the semantic importance of each modality is an important aspect of multi-modal learning. Our current work primarily focuses on evaluating robustness and generalization across arbitrary missing modality combinations, which explains our emphasis on combination-level performance (mean and standard deviation). However, we acknowledge that some modalities may be inherently more informative or complementary, especially in specific clinical tasks. In fact, as shown in Table 1 of the paper, we observe that combinations including FLAIR generally lead to higher Dice scores, indirectly suggesting its stronger contribution. This observation aligns with clinical knowledge where FLAIR is often more sensitive to edema and tumor boundaries. We will discuss the semantic importance in the final version.

[W2, Although UniMRSeg is evaluated across distinct domains (e.g., medical and natural scenes), the paper does not explore whether a model trained on one task (e.g., BraTS) could transfer to another task (e.g., RGB-D segmentation). A discussion of cross-task generalizability would further strengthen the work.]

We agree that cross-task generalization is an important topic, especially for unified segmentation models. In our current work, we design UniMRSeg to support task-specific but modality-relax generalization, i.e., within each segmentation domain (e.g., BraTS, SUNRGBD, STERE), the model handles all possible missing modality combinations using a single shared model. This is already a significant step toward practical unification, compared to previous works that require one model per combination. However, transferring a model trained on one task (e.g., brain tumor segmentation) to a fundamentally different task (e.g., RGB-D salient object segmentation) is non-trivial, due to large domain gaps in image appearance, semantics, and supervision forms. Such cross-task generalization is out of the current scope, but we agree it is a promising future direction. That said, our architecture is compatible with such a generalization setting. In future work, we plan to explore whether the shared encoder in UniMRSeg can be reused across domains via task-specific adapters or prompting mechanisms.

[W3, the current model does not explicitly infer which modalities are missing at test time. It would be helpful to understand whether this lack of modality-awareness could affect reliability in certain deployment settings.]

We would like to clarify that our UniMRSeg does not rely on explicit modality-type awareness during inference. To this end, we explicitly avoid designing an input pipeline that assumes modality-to-channel correspondence or requires modality labels. Instead, in Stage 1, we introduce a Random Modality Shuffle mechanism, where the order of input modalities is randomly permuted at each training iteration. This forces the model to learn modality-invariant representations rather than relying on fixed or predefined modality identity. In the main paper (Table 9), we also demonstrate that our model maintains stable segmentation performance across five different random shuffle orders during inference without any prior knowledge of modality types. This shows that the learned representation is robust to modality ambiguity, and the model behaves as if it had already inferred or become invariant to missing or shuffled modality types. Hence, while we do not adopt an explicit modality classification or selection module (which may introduce additional complexity or assumptions), our design intentionally achieves the goal of modality-agnostic reasoning in a simpler and more scalable way.

Dear Reviewer PWbw,

Many thanks to your professional, detailed, and valuable reviews. We're going to response to your concerns one by one.

[W1, the training pipeline's complexity raises questions about its necessity. While the ablation studies (Tables 3 and 5) demonstrate that each stage contributes to overall performance, it remains unclear why these components cannot be integrated into a single, unified loss function for joint optimization.]

Thank you for your suggestion. Prior to submission, we actually conducted the single-model unified training experiment. However, we chose not to include it in the paper, as our intention was to help readers focus more clearly on the self-supervised benefits brought by each stage individually. This is because the distinction between multi-stage and single-stage training essentially reflects a difference between self-supervised representation learning and multi-task learning.

(1) Self-supervised Representation Learning (Three Stages)

Our work is built upon a self-supervised pretraining framework, with training initialized from scratch starting at Stage 1. For each task (e.g., BraTS 2020), self-supervised training is conducted solely on the task’s own training set, rather than relying on externally labeled datasets such as ImageNet-pretrained weights for initialization. Our three-stage design includes:

• Stage 1: Pretrain both the encoder and decoder through arbitrary modality reconstruction.
• Stage 2: Pretrain the encoder with contrastive learning. The segmentation task here is only used to guide the contrastive objective, not as a final goal.
• Stage 3: Perform the downstream segmentation task.

This design explicitly aims to reduce the representation gap between complete and incomplete modalities in the encoder-decoder space during the final segmentation stage. If all these tasks are trained jointly in a single-stage model, it would fall into the scope of multi-task learning, rather than self-supervised pretraining.

(2) Multi-task Learning (Single Stage)

In multi-task learning, the goal is to leverage multiple complementary clues for collaborative learning. These clues can come from the data level [1,2,3], or from structural supervision types [4,5,6]. Most approaches often incorporate deliberate designs for task-sharing and task-specific components to enable effective joint prediction across tasks. In this work, if the three stages are forcibly merged into a single training stage, the following issues will arise:

• The model is trained with inputs that involve Random Modality Dropout, Random Modality Shuffle, and Random Spatial Masking. The encoder is supervised using the NT-Xent contrastive loss, an adapter is incorporated, and the model is simultaneously tasked with both segmentation and reconstruction. Such a fully entangled training setup lacks a clear task hierarchy, and there is no explicit coordination among the input, encoder, decoder, and output.
• When all tasks are treated equally without ordering, joint optimization becomes highly difficult. In the unified single-stage training attempt, the total loss involves six parts. The results on BraTS 2020 are as follow:

We observed the following phenomena during the experiments:

• The single stage model failed to converge, with loss plateauing early.
• The optimization process was highly unstable, with different losses fluctuating in turn and failing to decrease consistently together.

These phenomena and performace clearly indicate that unified single stage training is unable to achieve effective coordination among the various designs and supervision signals. The lack of a clear training order, combined with competing optimization objectives, leads to mutual interference.

In summary, we would like to clarify two potential misconceptions:

1. The three-stage training is not overly complex. Both theoretically and practically, it is a simpler and more stable design, especially when compared to directly mixing all objectives and modules in a single unified training phase, which leads to optimization difficulties.
2. The three stages are not repeating the same task. Instead, they form a self-supervised learning (SSL) pipeline that does not rely on extra data, and progressively reduce the gap between incomplete and complete modalities at the input, feature, and output levels to improve segmentation performance.

References:
[1] Segment Anything, ICCV 2023.
[2] SegGPT: Towards Segmenting Everything in Context, ICCV 2023.
[3] Spider: A Unified Framework for Context-dependent Concept Segmentation, ICML 2024.
[4] Deep Multitask Learning with Progressive Parameter Sharing, ICCV 2023.
[5] MTMamba: Enhancing Multi-Task Dense Scene Understanding by Mamba-Based Decoders, ECCV 2024.
[6] Multi-Task Dense Prediction via Mixture of Low-Rank Experts, CVPR 2024.

[W2, It is still not clear why fine-tuning of the encoder degrades performance (line 310-311). Could you please elaborate on that ?]

We respond from two perspectives: the inheritance between Stage 2 and Stage 3, and the design rationale of the lightweight reverse attention adapter.

(1) Inheritance between Stage 2 and Stage 3

Unlike general SSL methods (e.g., SimCLR, MoCo, MAE) that focus on learning generic representations without targeting specific downstream tasks, our contrastive learning in Stage 2 is task-aware. As stated in Lines 197–199, we co-train the segmentation head to guide the encoder toward learning modality-invariant features that are beneficial for segmentation, rather than general-purpose representations. This task-guided design ensures that the learned contrastive space aligns with the downstream segmentation objective. As shown in Table 3, this approach significantly boosts segmentation performance. Therefore, fine-tuning the encoder in Stage 3 would undermine the task-guided contrastive representations we carefully established.

(2) Lightweight Reverse Attention Adapter (RAA)

The RAA is designed as a residual correction to bridge the encoder representation gap between incomplete and complete modality inputs during Stage 3. Formally:

where:

• $f_{\text{inc}}$: encoder features from incomplete modality input,
• $f_{\text{com}}$: encoder features from the complete modality input,
• $\mathcal{A}(\cdot)$: the learnable adapter.

During Stage 3, both $f_{\text{inc}}$ and $f_{\text{com}}$ are frozen, and only $\mathcal{A}$ is trained. This well-constrained setup ensures that the adapter focuses purely on compensating the missing information, enabling stable and efficient optimization. Once the encoder is unfrozen, all three components become variables, making their roles unclear and the optimization unstable. This is why the "finetune" setting in Table 4 leads to significant performance degradation.

In conclusion, freezing the encoder in Stage 3 is a deliberate and necessary design choice to preserve the task-guided contrastive representations learned earlier. We will make this motivation and its empirical evidence more explicit in the revised manuscript. Thank you again for the insightful question.

[W3, NestedFormer does not handle missing modality; it is unclear how the benchmark is done (Table 1).]

Thank you for pointing this out. NestedFormer is not designed for missing modality settings. We include it solely as a comparative baseline to highlight that segmentation methods designed for complete modalities perform significantly worse under missing modality conditions, due to the lack of explicit modeling for missing inputs. In contrast, all other methods in Table 1 are specifically designed to handle missing modality. We will clarify this purpose in the revised paper.

Dear Reviewer PWbw,

Thanks for your constructive comments and valuable suggestions to improve this paper.

As we are entering the final days of the author-reviewer discussion period, we are not sure whether our rebuttals have addressed your concerns. We would be greatly appreciated if you could spend some of your time for a reply.

Thank you again for your time and consideration.

Best regards,

The Authors

Dear Reviewer sDPL,

Thank you for your time and careful review of our work. Below we address your questions and weaknesses mentioned:

[Q1 W1, the article may lack certain novelty, such as in the design of the Perturbation and Mask components, as these techniques have likely been extensively explored in domains like computer vision, natural language processing, or time series analysis. The article does not appear to clearly highlight a more distinctive motivation or design, mainly extending existing technologies. Additionally, the incorporation of lightweight reverse attention adapters in the model may also lack sufficient innovation, so the core innovative aspects of the paper need to be better emphasized.]

We understand the concern regarding the novelty of individual components. We would like to emphasize that the novelty of UniMRSeg does not lie in standalone algorithmic innovations, but rather in the hierarchical, task-driven integration of multiple self-supervised mechanisms, jointly optimized to address the unique challenges of modality-relax segmentation. Specifically:

(1) Hierarchical Self-Supervised Compensation (HSSC) Framework

(as already acknowledged and appreciated in your previous review)

UniMRSeg unifies pixel-level modality reconstruction, feature-level contrastive learning, and prediction-level label distillation into a single annotation-free compensation framework. Unlike previous works that treat self-supervised tasks in isolation, HSSC employs a three-stage curriculum, jointly optimized with a single shared model per task, enabling robust generalization across arbitrary missing-modality combinations.

(2) Hybrid Data Perturbation

While data perturbation techniques have been explored in various tasks, we design a hybrid perturbation mechanism specifically tailored to the modality-relax setting:

• Random modality dropout simulates missing inputs;
• Modality shuffle removes the reliance on fixed modality-channel mapping, encouraging modality-agnostic encoding;
• Spatial masking mimics local degradation.

As discussed in Lines 585–606 and validated in Table 9, conventional models depend on strict input order and modality metadata, which limits their robustness in real-world scenarios. Our strategy removes these assumptions, improving deployment flexibility and generality.

(3) Lightweight Reverse Attention Adapter (RAA)

The RAA is designed as a residual correction to bridge the encoder representation gap between incomplete and complete modality inputs during Stage 3. Formally:

where:

• $f_{\text{inc}}$: encoder features from the incomplete modality input,
• $f_{\text{com}}$: encoder features from the complete modality input,
• $\mathcal{A}(\cdot)$: the learnable adapter.

During Stage 3, both $f_{\text{inc}}$ and $f_{\text{com}}$ are frozen, and only $\mathcal{A}$ is trained. This well-constrained setup ensures that the adapter focuses purely on compensating the missing information, leading to stable and efficient optimization.

To further enhance the adapter’s ability to capture under-represented semantic cues, we apply reverse attention to amplify the residual signal. In this way, RRA can focus on the non-salient or poorly attended regions that conventional self-attention may overlook.

Overall, our original intention was not to pursue isolated innovations on specific components or modules by introducing increasingly "fancy" designs in the hope of achieving strong results for unified modality-relax segmentation. While such improvements may be valuable directions for future work, they are not the focus of this paper. Instead, this work aims to demonstrate that by strictly adhering to the motivation of Hierarchical Self-Supervised Compensation (HSSC) and designing around it in a principled manner, UniMRSeg is able to achieve strong and consistent performance across tasks. We hope that the motivation and explanation above can be well understood and appreciated.

[Q2 W3, the article requires improvement in the visual presentation aspects, such as in the rendering of Figure 1 and Figure 2. The figures currently contain limited effective information or may not fully reflect the model's own structure or characteristics, necessitating further refinement. ]

Thank you for your suggestion. In the revised version, we will improve Figure 1 and Figure 2 as follows:

(1) Figure 1

We will split the current figure into two subfigures for clarity:

• Figure 1a: A clear, labeled block diagram of the full architecture, including the encoder, decoder, perturbation module, contrastive learning module, and reverse attention adapter.
• Figure 1b: The benchmark box plot will be retained, with expanded labels and annotations to improve readability.

(2) Figure 2

We will restructure Figure 2 to better highlight the three-stage training pipeline and the interconnections between stages, including:

• Inputs and outputs at each stage;
• The flow of gradients and supervision signals (e.g., reconstruction, contrastive, and segmentation losses);
• The dependency and progression between stages, as discussed in Response to [Q1 W1], emphasizing how Stage 1 prepares robust features, Stage2 aligns them with segmentation objectives, and Stage 3 adapts to incomplete modalities using frozen representations.

[Q3 W2, the article should compare its proposed model against a more competitive set of baseline models, and expand the experimental scale to further elucidate the performance advantages of the model presented in the paper.]

Thank you for your suggestion. We fully agree that strong baselines and broad evaluations are essential to demonstrate the effectiveness and generality of our proposed method.

(1) On Baseline Competitiveness

We would like to emphasize that Tables 1, 2, 6, and 7 already include state-of-the-art and highly competitive methods published in top-tier conferences/journals, such as:

• PASSION (MM 2024)
• GateNet (IJCV 2024)
• EonTRINet (PAMI 2024)
• MaskMentor (MM 2024)
• M3FeCon (MICCAI 2024)

These methods represent the most advanced techniques in multi-modal/ missing modality segmentation. Most importantly, they are among the few methods with publicly released code and pre-trained models, enabling fair and reproducible comparisons under consistent settings.

While other recent approaches do exist, many do not provide released implementations, making standardized evaluation infeasible. We are cautious about comparing only reported numbers without aligned settings.

In order to fully address your suggestion and further strengthen the persuasiveness of our results, we report a performance comparison with the recently published model LS3M [1] (CVPR 2025, code not yet released). Although its implementation is unavailable, the reported results on BraTS 2018 are relatively aligned with our evaluation setup, allowing for an approximate comparison.

[1] Incomplete Multi-modal Brain Tumor Segmentation via Learnable Sorting State Space Model, CVPR 2025.

(2) On Evaluation Scope

UniMRSeg has been evaluated across diverse domains, including:

• Medical segmentation: BraTS 2018, BraTS 2020
• RGB-D semantic segmentation: SUN-RGBD
• RGB-D salient object detection: STERE
• RGB-T salient object detection: VT1000

Each dataset is tested under all possible missing modality combinations, with a single unified model per task, to assess both task-level and modality-level generalization. The scope of evalution is broader than most previous works.

(3) Planned Expansions in Final Version

To further strengthen the evaluation, we will:

• Include comparisons with generalist segmentation models (e.g., SAM2, SegGPT) under in-context setups using prompt masks.
• Extend the evaluation to additional RGB-D SOD and SS datasets, such as SIP, NJUD, NLPR, and NYUD.

In summary, our current evaluation already reflects a wide and competitive scope. We are committed to further expanding the benchmarks and baselines to reinforce the value of our work. Thank you again for the constructive suggestion.

[Limitations: The author mentions the limitations of the study in the article, mainly that the three-stage training strategy increases the complexity of the training process and introduces additional training overhead, which may pose challenges for practitioners with limited time or resources. I consider this explanation quite adequate, and such trade-offs are within an acceptable range for the sake of better performance.]

Thank you very much for acknowledging our explanation regarding the three-stage training strategy. We also kindly refer you to our response to Reviewer PWbw’s [W1], where we further elaborate on the necessity of the three-stage pipeline and demonstrate its advantages over a unified single-stage alternative. We appreciate your understanding of the trade-off between performance and complexity.

Thanks for the author's positive response. It has addressed all my concerns. The effort put into such a detailed explanation is highly appreciated, and I have improved my overall rating.