🧬 biology

3D Masked Autoencoders are Robust Learners of Volumetric and Multimodal Cellular Representations for Microscopy

This paper demonstrates that 3D masked autoencoders, particularly when enhanced with cross-modal alignment to protein language models and specialized 3D-aware architectural components, significantly outperform 2D-based approaches in learning robust volumetric representations for single-cell microscopy tasks such as protein localization and interaction prediction.

Original authors: Amirhossein Kardoost, Lion Gleiter, Tingying Peng, Carsten Marr

Published 2026-06-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Amirhossein Kardoost, Lion Gleiter, Tingying Peng, Carsten Marr

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are trying to understand the layout of a bustling city. Most previous computer programs studying cells (the tiny building blocks of life) were like taking a single, flat photograph of the city from a drone. They would squish all the tall buildings and deep basements into one flat image. While this gives you a general idea, you lose all the depth, the hidden alleyways, and how things stack on top of each other.

This paper introduces a new way for computers to "see" cells. Instead of flattening the image, the researchers built a system that looks at the cell as a full, 3D volume—like holding a transparent block of the city in your hands and walking around it.

Here is a breakdown of their approach and findings using simple analogies:

1. The "Blurry Photo" vs. The "3D Model"

The researchers compared two types of AI students:

Student A (2D): This student only studies flat, 2D photos of cells. Even if the cell is a 3D object, Student A squishes it flat to study it.
Student B (3D): This student studies the full 3D "block" of the cell, preserving the depth and layers.

The Result: Student B (the 3D model) consistently learned better. When asked to identify where specific proteins (the workers) are located inside the cell, or to guess if two proteins are friends (interacting), Student B was much more accurate. The paper claims that keeping the full 3D shape provides a much richer "memory" of the cell than flattening it out.

2. The "Fill-in-the-Blanks" Game (Masked Autoencoders)

To teach these students, the researchers used a game called a "Masked Autoencoder." Imagine showing a student a picture of a cell, but then covering up 75% of it with black tape. The student has to guess what is hidden underneath based on the small visible parts.

By forcing the AI to reconstruct the missing 3D parts, it learns the rules of how cells are built.
The paper found that the 3D student got much better at this game than the 2D student, proving it understood the structure of the cell much better.

3. Adding a "Dictionary" (Protein Language Models)

Cells have a "blueprint" written in their DNA, which is a sequence of letters (like a language). The researchers gave their 3D student a special dictionary (a pre-trained protein language model called ESM2) that understands this biological language.

The Analogy: Imagine trying to identify a specific tool in a toolbox. If you only look at the tool's shape (the image), it might be hard. But if you also read the label on the tool (the protein sequence), it becomes much easier.
The Result: When the 3D student used this "dictionary" alongside the 3D images, it learned even faster and more accurately. The paper notes that this "multimodal" approach (combining images and text) helped the 3D model significantly more than it helped the 2D model.

4. The "Frequency" Filter

The researchers also added a special rule to the training game. They told the AI: "Don't just guess the general shape; make sure the tiny, fine details (like the texture of a cell wall) look sharp."

They used a mathematical trick (called FFT) to check if the "fine print" of the image was being reconstructed correctly. This helped the AI focus on the tiny, critical structures inside the cell rather than just the big blobs.

The Bottom Line

The paper concludes that for understanding cells, 3D is better than 2D.

On the "Protein Localization" task (finding where proteins live): The best 3D model achieved a score of 0.952, beating previous top methods.
On the "Protein Interaction" task (guessing if proteins work together): The 3D model scored 0.865, also beating previous methods.

In short, the researchers showed that if you want a computer to truly understand the complex, 3D world inside a cell, you have to let it see the whole 3D picture, not just a flat shadow of it. They also proved that giving the computer a "dictionary" of protein names helps it understand that 3D picture even better.

Technical Summary: 3D Masked Autoencoders for Cellular Representation Learning

Problem Statement
Fluorescence microscopy provides rich, high-dimensional data capturing the subcellular organization of cells, often in the form of volumetric z-stacks (e.g., OpenCell, WTC-11). However, existing representation learning approaches for single-cell analysis, such as Subcell and DINO4Cell, predominantly operate on 2D maximum-intensity projections or individual slices. This practice discards inherent depth-resolved structural information, potentially limiting the model's ability to learn discriminative cellular representations. Furthermore, purely visual models often lack integration with the biological priors encoded in protein sequences. The paper investigates whether native 3D modeling and multimodal alignment with protein language models can overcome these limitations to improve downstream tasks like protein localization and protein-protein interaction (PPI) prediction.

Methodology
The authors propose a systematic comparison and enhancement of Masked Autoencoder (MAE) frameworks for volumetric microscopy data.

Architectural Variants:
- MAE-2D: The input 3D volume is collapsed into a 2D maximum-intensity projection ( $I_{2D}$ ). The model processes non-overlapping 2D patches using a 2D ViT backbone with sinusoidal positional embeddings.
- MAE-3D: The full 3D volume ( $I_{3D}$ ) is retained and divided into 3D patches ( $p_z \times p_x \times p_y$ ). The model utilizes a 3D ViT backbone with 3D sinusoidal positional embeddings, preserving the depth dimension throughout the encoder-decoder pipeline.
- Channel Cross-Attention (CCA): To explicitly model inter-channel interactions, the authors introduce a dual-stream encoder-decoder where each channel is a separate token stream. A shared random masking pattern is applied across channels. Cross-attention is implemented where a token at spatial index $i$ in channel $c_i$ queries the token at the same index in channel $c_j$ . A sigmoid gating mechanism replaces the softmax operation to handle the single-key constraint.
- Frequency-Domain Regularization (FFT Loss): To preserve fine subcellular structures, an L1 loss is applied in the frequency domain between the magnitude spectra of the original and reconstructed images. This loss is computed over 2D axes for MAE-2D and all three volumetric axes for MAE-3D. It is linearly ramped up during training to stabilize early learning.
Multimodal Alignment:
The framework integrates a pretrained protein language model (ESM2) to align visual features with protein sequence information.
- Decoder Conditioning: Frozen ESM2 embeddings are projected to the decoder dimension and inserted as a single token into each channel-specific decoder sequence (with zero positional embedding). This allows masked patches to attend to protein context via self-attention.
- Contrastive Objective: A symmetric InfoNCE loss ( $L_{CLIP}$ ) aligns the normalized image embeddings from the encoder with the protein embeddings. This alignment is computed only on visible tokens to encourage robustness.
- Total Loss: The final objective combines Mean Squared Error (MSE) for reconstruction, FFT loss, and CLIP loss: $L = L_{MSE} + w_{FFT}L_{FFT} + w_{CLIP}L_{CLIP}$ .

Key Contributions

Systematic 3D vs. 2D Comparison: The authors demonstrate that native 3D MAE models consistently outperform 2D max-projection and slice-based variants across downstream tasks, even when the 2D models are restricted to the same number of slices.
Enhancement via CCA and FFT: The paper identifies channel cross-attention and frequency-domain regularization as critical components for leveraging 3D spatial context and preserving structural fidelity.
Multimodal Integration: The study establishes that aligning image features with protein sequence embeddings (via ESM2) significantly improves representation quality, with the gains being substantially larger for volumetric (3D) models than for 2D models.

Results
Experiments were conducted on the OpenCell dataset (1,310 proteins, ~6,300 volumes) using five-fold cross-validation.

Protein Localization:
- MAE-3D vs. MAE-2D: MAE-3D consistently outperformed MAE-2D across all mask ratios. Even when restricted to a central slice range (45–55), MAE-3D maintained superior performance (AUCmicro: 0.91 vs. 0.90).
- Multimodal Gains: The best 3D model (MAE-3D⋆) achieved state-of-the-art results with an AUCmicro of 0.952 and F1micro of 0.742, improving over previous approaches by +0.003 and +0.010, respectively. The integration of ESM2 provided larger performance boosts for the 3D model compared to the 2D variant.
- Attention Analysis: Visualizations showed that integrating ESM2 concentrated attention on protein-relevant structures (e.g., actin cytoskeleton) across relevant z-slices, whereas models without ESM2 or 2D models exhibited scattered attention or lack of depth specificity.
Protein-Protein Interaction (PPI):
- MAE-3D⋆ achieved a ROC-AUC of 0.865, outperforming prior methods (Subcell: 0.85, DINO4Cell: 0.84) by up to +0.025.
- Notably, the 3D model achieved these results on the smaller OpenCell dataset (~6k volumes) compared to models pretrained on the much larger HPA dataset, highlighting the efficiency of volumetric modeling in data-limited settings.

Significance and Claims
The paper claims that preserving full 3D volumetric context yields more informative and discriminative representations for single-cell microscopy than 2D projections. It posits that the synergy between volumetric morphology and protein-level semantics is strongest when using native 3D architectures. The authors conclude that native 3D modeling combined with multimodal alignment is essential for developing robust foundation models in cellular imaging. They emphasize that these advantages hold even without massive-scale pretraining, suggesting that architectural choices (3D vs. 2D) and multimodal integration are critical factors for performance.