ChA-MAEViT: Unifying Channel-Aware Masked Autoencoders and Multi-Channel Vision Transformers for Improved Cross-Channel Learning

Incremental novelty: All core ideas (channel masking, memory tokens, hybrid fusion, channel tokens) have direct antecedents (CA-MAE, DiChaViT, ViT register tokens).

Thank you for the opportunity to clarify our contributions.

First, DCP Masking is not derived from CA-MAE but is a novel method, as noted by reviewers 4uNy and La1T. Unlike prior MAEs that assume cross-channel redundancy, DCP Masking addresses the low redundancy in multi-channel data by masking both channels and spatial patches to promote stronger cross-channel interactions. As shown in Tables 2 and 4, it outperforms CA-MAE's patch masking by 1-6% and 3-21% across three datasets with different backbones.

Second, as stated throughout the paper, memory tokens and Hybrid Fusion module are adapted from existing ideas but are for novel purposes with a non-trivial shift in functionality: enhancing cross-channel interactions in masked modeling. Memory tokens capture global cross-channel semantics, while the Hybrid Fusion module leverages reconstructed patch tokens to enrich the final representation.

Third, Channel-Aware Decoder introduces a scalable architecture that uses a shared lightweight decoder with channel tokens, unlike prior methods (e.g., CA-MAE) that require separate decoders per channel.

Importantly, there is no guarantee combining prior ideas will work at all. Indeed, Table 2 shows merging the strongest baseline (DiChaViT) with CA-MAE yields only modest gains and still lags behind ChA-MAEViT by 3.5-12.5%. Our work investigates the role of all the components, both existing, newly proposed, and their integration, offering insights on what matters to solve open problems in the field. We’ll revise the camera-ready version to clarify these points.

Method is overly heuristic: Alternating between patch- and channel-masking lacks a principled basis;

Alternating between masking strategies helps mitigate the learning difficulties caused by excessive information loss when channels and patches are masked simultaneously. As discussed in Appendix F.6 (line 652), a patch ratio of 0.25 when used with dynamic channel masking, results in about 65% of the spatial content being masked, comparable to the 75% masked in Alternate masking. Since each channel carries distinct semantics, masking both at once can be difficult to recover. By alternating the two, we observed more stable learning and up to a 5% performance gain on the JUMP-CP dataset (Table 4).

hyperparameters require extensive tuning and appear dataset-specific, reducing out-of-the-box usability.

ChA-MAEViT does not require extensive or dataset-specific hyperparameter tuning. All our experiments used the same hyperparameters without performing any separate tuning for each dataset. Specifically:

Reconstruction Lambda: fixed 0.99 for all experiments (line 308)

DCP Masking has three hyperparameters (mask ratio, p_patch, p_channel). As detailed in Sec. 3.1 (line 156), we used two fixed configurations: Alternate (0.75, 0, 0) and Combination (0.25, 0.5, 0.5), across all experiments. Table 4 shows both consistently outperform the best baseline in Table 1 by 2.5-21.0% and 3.5-16.5%, showing strong generalizability.

Memory Token: As shown in Fig. 3, a small number of tokens (e.g., 4) boosts performance, with diminishing returns beyond that. Results are stable, showing low sensitivity to this hyperparameter.

These results show that our default settings are effective across datasets, and can likely be used in new settings directly.

We'd like to emphasize that additional hyperparameters do not necessarily diminish the method’s out-of-the-box usability, e.g., DINOv2 has over 20 and is widely adopted for its stability and performance. Ours are stable, transferable, and require no per-dataset tuning.

To further support reproducibility and enhance understanding, we provide extensive analysis (Table 4, Fig. 3, and Appendix Figs. 9, 10, 12, 13), confirming ChA-MAEViT’s robustness and effective generalization across tasks.

Limited task scope: Evaluated only on classification; no evidence it extends to segmentation or regression in MCI.

Only JUMP-CP and So2Sat are classification tasks. In CHAMMI, methods are evaluated based on their ability to learn useful cellular representations across various axes of generalization (e.g., new plates, new cell lines), using nearest neighbor protocols based on feature sets from the CHAMMI’s authors. Thus, CHAMMI assesses representation learning, not a classification task. ChA-MAEViT performs well on both tasks (3-21.5% boost across three datasets), showing a broad applicability that merits publication in its own right. In addition, this is the same set of experimental settings of prior work [1,12,13,28], suggesting the community finds these two evaluation task types are effective.

That said, to further demonstrate generalization, we evaluate segmentation performance on the 38-Cloud dataset [A], which includes four-channel (RGB + NIR) satellite images. Below are the average scores and std from three runs:

ChA-MAEViT outperforms all baselines across all evaluation metrics, demonstrating its effectiveness in enhancing multi-channel representation learning.

[A] Cloud-Net: An End-To-End Cloud Detection Algorithm for Landsat 8 Imagery, Mohajerani et al., IGARSS’19

Unclear robustness: Experiments report averages over random channel removals, but worst‐case drops when critical channels are missing are not shown.

Results in Table 1 follow established evaluation protocols from prior work, not random channel removals. CHAMMI’s varying channels stem naturally from differences in its three source datasets, i.e., we do not remove any channels. For So2Sat and JUMP-CP, we follow the established evaluation protocols of [2,3]. This can result in some critical channels never being masked, which is why we provided a more complete evaluation.

Specifically, in Table 3, we report the average over all possible channel combinations as stated in its caption and L279, not random masking. E.g., column "4" shows the mean and std across all 70 four-channel combinations (8C4 = 70). Due to space constraints, only the mean and std values are shown in the main paper. However, detailed result in Appendix Table 8 does show the worst-case drop, when channel 3 is removed. Even in this worst-case setting, our method achieves the best result, outperforming the SOTA by 23%.

To leave no further doubt, we extend Table 3 by including the minimum performance reported across all channel combinations below:

“avg” indicates the average across all possible combinations, while “min” represents the worst-case performance. Note that different from Table 3, here we report the average of three runs and the std for model variance. As seen above, our approach consistently outperforms the prior state-of-the-art.

Unreported compute costs: No training/inference time comparisons versus simpler supervised baselines or existing MAEs.

We report training runtimes in Appendix C (line 561), and FLOPs comparisons in Figure 5 of the Appendix. Specifically, on JUMP-CP, DINOv2 takes 6 days, SimCLR 2 days, and CAMAE/ChA-MAEViT 1 day. Figure 5 shows ChA-MAEViT has the lowest FLOPs, about 1/6 of DINOv2, while outperforming it.

How does the model handle truly unseen channel types at test time?

In Appendix F.11, we show that ChA-MAEViT can handle novel channels unseen during training. We initialize new channel tokens for these channels and fine-tune only these tokens using a small set of unlabeled test images with DCP, while freezing the rest of the model. This approach improves performance as more unlabeled images are used (Fig. 16).

If some channels are rare during training, DCP masking may not adequately teach reconstruction, leading to overfitting to “full” channel sets.

Rare channels do not lead to overfitting to full channel settings. DCP masking consistently uses a subset of channels during training (line 138), forcing the model to be robust to missing modalities and mitigating reliance on any specific channel configuration. Also, DCP masking outperforms patch masking by 4% (Table 4) on CHAMMI, in which some channels are always missing due to the nature of the dataset, meaning it never sees “full” channel sets. Rare channels, instead, reflect the natural data distribution and do not imply overfitting. We demonstrate robustness under these conditions in CHAMMI, where the “membrane” channel occurs only 31% as often as “nucleus.” Yet, our model outperforms the best baseline by 2% on images with the “membrane” channel (Table 7, Appendix), showing strong generalization under uneven channel distributions. Additionally, techniques like channel-aware sampling, oversampling, or importance-weighted loss can be easily integrated into our method to enhance learning if needed.

We appreciate the reviewer’s thoughtful comments and respond to the questions below.

While the DCP masking strategy is a clear and novel contribution, the other components are largely effective adaptations of existing ideas. Memory tokens are directly inspired by Register Tokens, the Channel-Aware Decoder builds upon the channel token concept from prior work like ChannelViT and DiChaViT, and the hybrid fusion module is a relatively standard approach.

Thank you for recognizing the novelty of our DCP masking strategy. We would like to clarify the contributions of the other components as follows.

First, as stated throughout the paper, the memory tokens and Hybrid Fusion module are adapted from existing ideas but are employed for novel purposes. Specifically, we apply them in the context of multi-channel representation learning to enhance cross-channel interactions in a masked modeling framework. This constitutes a non-trivial shift in functionality: the memory tokens serve as global carriers of multi-channel semantics throughout the encoder-decoder pipeline. Hybrid Fusion module further contributes by making use of fine-grained patch tokens learned during reconstruction to enrich the final representation.

Second, Channel-Aware Decoder introduces a new architectural design by integrating channel tokens with a lightweight shared decoder. This contrasts with prior work such as CA-MAE, which does not use channel tokens and instead relies on separate decoders for each channel. By enabling a unified decoding process, our approach improves scalability to a larger number of channels and yields better performance.

Importantly, there is no guarantee that just combining prior ideas will work at all. For example, in Table 2 of the main paper, we show that integrating the strongest baseline (DiChaViT) with CA-MAE yields only modest improvements and still significantly falls short of ChA-MAEViT by 3.5%-12.5% across three datasets. Our work investigates the role of all the components, both existing, newly proposed, and their integration. The paper presents insights on what matters to solve open problems in the field.

We will revise the camera-ready version to clarify these points.

Hyperparameter Complexity of DCP Masking: The proposed DCP masking introduces several new hyperparameters that appear sensitive. The appendix reveals that the two main configurations ("Combination" and "Alternate") require different settings for the patch ratio (0.25 vs 0.75, respectively).

First, we would like to point out that the same masking ratios were used for every dataset in our evaluation, i.e., these hyperparameters are stable and generalize well to new settings, making them easy to use. That said, these two settings use different hyperparameters, but result in the same amount of information being masked out in input images. In other words, the choice of patch ratio stems from ensuring the same amount of information is masked out, rather than the sensitivity of the hyperparameters.

Alternate masking most closely resembles the prior work like MAE and CA-MAE and, thus, we adopt a patch masking ratio of 0.75, consistent with standard practice for MAE and CA-MAE.

Combination differs in that we mask entire channels, and, thus, we reduce the patch ratio to 0.25 due to the additional channel-level masking. As discussed in Appendix F.6 (line 652), a patch ratio of 0.25, when used with dynamic channel masking, results in approximately 65% of the spatial content being masked on average, which is comparable to the 75% masked in Alternate masking. Thus, applying a high patch ratio (e.g., 0.75) on top of channel masking would cause excessive information loss, making reconstruction extremely difficult and unstable. It is also reasonable to argue that masking entire channels is inherently more challenging than masking individual patches, as it removes all modality-specific information rather than local details, suggesting that a slightly smaller total masking ratio of 65% is warranted in this setting.

Furthermore, Figure 10 shows a non-trivial trade-off between performance on full vs. partial channel settings based on the balance between patch and channel masking. This suggests that applying ChA-MAEVIT to a new dataset might require careful and potentially expensive tuning of the masking strategy.

As noted above ChA-MAEViT does not require expensive hyperparameter tuning as all experiments across all three datasets were performed using the same masking hyperparameters. As shown in Table 4 of the main paper, both configurations significantly and consistently outperform the SOTA in Table 1 for both full and partial channel settings. Specifically, the Alternate configuration achieves boosts of 2.5-21.0%, while the Combination improves by 3.5-16.5%. This highlights the robustness and generalizability of the hyperparameters in our model across both satellite and microscopy imaging domains.

Fig. 10 illustrates the effectiveness of combining both channels and patch masking. It shows that using both methods leads to better performance compared to using either one alone. Additionally, it provides insights into the trade-offs associated with each type of masking, which does not mean it requires costly tuning for new datasets. Thus, we can recommend these masking ratios as reliable defaults for applying ChA-MAEViT to new datasets.

Thank you for your insightful feedback! We address the reviewer’s questions below.

This is really an engineering paper that presents improved architectures, with reasonable arguments for them but without a corresponding theory

Dynamic Channel-Patch (DCP) Masking involves removing both patches and entire channels from multi-channel inputs, where each channel typically contains distinct and complementary information, e.g., different sensor modalities or spectral bands. When these channels are combined, they often create emergent representations, which are features that arise solely from interactions between the channels. Masking disrupts not only the information within each channel but also the crucial synergy between channels.

This concept can be formalized using mutual information. Let $X_{\text{full}} represent the complete multi-channel input, and \X_{\text{observed}}$ denote the unmasked subset. The relationship can be expressed as follows:

$$I(Z; X_{\text{full}}) = I(Z; X_{\text{observed}}) + I(Z; X_{\text{masked}} \mid X_{\text{observed}})$$

When part of the input is masked, the conditional mutual information term $I(Z; X_{\text{masked}} \mid X_{\text{observed}})$ becomes zero, indicating that the model loses access to the synergistic cross-channel dependencies that could otherwise be recovered from the full input.

Reconstruction loss, such as Mean Square Error, compels the model to predict the missing patches and channels:

$$\mathcal{L}\_{\text{rec}} = \mathbb{E}\_{X} \left[ \sum_{c \in \mathcal{M}} \sum_{p} \left \| X\_{c,p} - \hat{X}_{c,p} \right \|_2^2 \right]$$

Here, $\mathcal{M} represents the set of masked channels, and \p$ indexes spatial patches. This objective forces the encoder to capture both intra-channel structure and inter-channel dependencies, resulting in a more robust representation that retains emergent semantics and is resilient to missing modalities during inference. We shall update these discussions in our camera ready.

Reconstruction quality in Figure 15 looks quite low. I understand that the goal of the paper is not the reconstruction but the learned representation; still I am not sure if this is to be expected or if related work produces higher quality results.

Thank you for the insightful comment! Indeed, the low-fidelity reconstructions are expected and align with observations from prior work on Masked Autoencoders (MAEs), e.g., [A,B,C,D].

As the reviewer noted, the primary objective of ChA-MAEViT is to learn high-quality representations, not to generate detailed reconstructions. Our design encourages the encoder to capture high-level semantic and cross-modal dependencies, while the decoder remains intentionally lightweight to verify whether essential latent information has been encoded. Additionally, following prior MAEs (e.g., [A,B,C,D]), we employ the mean squared error (MSE) loss for reconstruction, which often results in blurry outputs by averaging plausible solutions [D, E]. This leads to reconstructions dominated by low-frequency components, an expected outcome of the design choice.

In contrast, models designed to generate high-fidelity images (e.g., [F, G, H]) typically rely on deeper, more expressive decoders. In such models, scaling the decoder is critical [H], and specialized loss functions such as perceptual loss [I] are employed to enhance visual quality.

[A] Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology, Krau et al., CVPR’24

[B] Masked Autoencoders Are Scalable Vision Learners, He et al., CVPR’22

[C] Rethinking Transformers Pre-training for Multi-Spectral Satellite Imagery, Noman et al., CVPR’24

[D] Scale-MAE: A Scale-Aware Masked Autoencoder for Multiscale Geospatial Representation Learning, Reed et aI., CCV’23

[E] Training a Task-Specific Image Reconstruction Loss, Mustafa et al., WACV’22

[F] Taming Transformers for High-Resolution Image Synthesis, Esser et al., CVPR’21

[G] TokenFlow: Unified Image Tokenizer for Multimodal Understanding and Generation, Qu et al., CVPR’25

[H] GigaTok: Scaling Visual Tokenizers to 3 Billion Parameters for Autoregressive Image Generation, Xiong et al., ICCV’25

[I] The Unreasonable Effectiveness of Deep Features as a Perceptual Metric, Zhang et al., CVPR’18

Why is the "Dynamic Channel-Patch Masking" called "Dynamic"? This implies that the masking strategy somehow dynamically adapts to the data. I don't think this is the case. Am I missing something?

The term "Dynamic" refers to the varying number of channels masked during training, rather than adaptation based on the specific input data. In other words, the masking strategy dynamically adjusts the number of masked channels and patches for each image, but this adjustment is independent of the actual content of the input samples. We will clarify this in the camera-ready version to avoid confusion.

Thank you for your encouraging feedback! We address the reviewer’s questions below.

Are there error bars that could be included in Table 2, Figure 3 and 4, to give an idea of statistical significance?

We report the mean and standard deviation (std) over three runs for the Table and Figures as follows for those results we were able to complete in the time constraints of the rebuttal period.

Table 2

Fig. 3b

Fig. 4

The relatively low standard deviations (std) across experiments suggest the results are stable, supporting the robustness of our findings. We will include the corresponding error bars in the camera-ready version.

Could ChA-MAEViT be applied to consider OOD channels, as this is a persistent issue when training models with fluorescence microscopy data.

We completely agree that OOD channel generalization is a critical and ongoing challenge in fluorescence microscopy. In Appendix F.11 we demonstrate that ChA‑MAEViT can handle novel channels never seen during training in a preliminary experiment by adjusting its weights via learning to reconstruct the new channels at test time. Specifically, we adapt to novel channels by mapping them to existing ones using a few unlabeled examples and our self-supervised objectives, resulting in a 2.5% boost using just 640 examples.

Line 175 - can you please clarify what channel tokens were used in the encoder? The paper would read better if a brief description was included here.

The channel tokens are learnable embeddings, each representing an input channel. Since all channels share a common projection layer, these tokens play a crucial role in capturing channel-specific features. They are concatenated with the spatial patch tokens and jointly processed by the transformer encoder and decoder. We will clarify this in the camera-ready version.