Axial-Centric Cross-Plane Attention for 3D Medical Image Classification

This paper proposes an axial-centric cross-plane attention architecture that leverages a frozen MedDINOv3 foundation model and directional transformer encoders to align 3D medical image classification with clinical workflows, demonstrating superior performance over existing methods by prioritizing the axial plane while integrating complementary coronal and sagittal information.

The Gist

Imagine you are trying to understand a complex 3D object, like a human heart, but you can only look at it through a series of 2D slices, similar to looking at the pages of a book one by one.

In the real world, doctors don't just stare at the whole book at once. They have a specific way of reading it:

1. The Main Page (Axial Plane): They start with the "top-down" view (the axial plane). This is their primary reference, like the main chapter of a story.
2. The Side Views (Coronal & Sagittal Planes): If they see something interesting or confusing on the main page, they quickly flip to the "front view" (coronal) or "side view" (sagittal) just to get a little extra context. They don't treat these side views as equally important; they are just helpers to clarify the main story.

The Problem with Old AI

For a long time, computer programs (AI) trying to do this job were a bit clumsy. They treated every single slice and every single angle (top, front, side) as if they were equally important. It was like a student trying to read three different books simultaneously, giving equal attention to every word in all of them. This wasted a lot of brainpower (computing power) and often missed the subtle clues because it didn't know which "page" was the most important.

The New Solution: The "Smart Detective"

This paper introduces a new AI architecture called Axial-Centric Cross-Plane Attention. Think of it as a Smart Detective who knows exactly how a human doctor thinks.

Here is how it works, using simple analogies:

1. The Expert Librarian (MedDINOv3)
First, the AI uses a pre-trained "Expert Librarian" (a model called MedDINOv3) that has already read millions of medical scans. This librarian is frozen in place—it doesn't learn new things during this specific task, but it's incredibly good at recognizing what a slice of a heart or lung looks like. It acts as the eyes for all three angles.

2. The Three Teams
The AI looks at the 3D scan from three angles: Top (Axial), Front (Coronal), and Side (Sagittal).

• The Intra-Team Meeting: Before they talk to each other, each team (Top, Front, Side) holds a private meeting to understand their own specific slice. They figure out, "Okay, in this top-down view, we see a tumor here."

3. The "Boss" and the "Consultants" (The Core Innovation)
This is the magic part. In the old AI, everyone shouted their opinions at once. In this new AI:

• The Top View (Axial) is the Boss. It holds the main decision-making power.
• The Front and Side Views are Consultants. They don't get to vote on the final answer; instead, they are there to answer questions the Boss has.

The AI uses a mechanism called "Cross-Plane Attention." Imagine the Boss (Top View) asking the Consultants: "Hey, I see this weird shape in my view. Does the Front View see anything that explains this?"
The Consultants look at their data and send back only the relevant information to help the Boss. The Boss then updates its understanding based on that specific help.

4. The Final Verdict
After the Boss has gathered the right clues from the consultants, it makes the final diagnosis. The AI doesn't just average the opinions of all three views; it lets the "Boss" view lead the conversation, using the others only to fill in the gaps.

Why Does This Matter?

The researchers tested this "Smart Detective" on six different medical datasets (like looking at broken bones, organs, or blood vessels).

• The Result: It beat almost every other AI model out there.
• The Reason: By mimicking how human doctors actually work (focusing on the main view and using side views only for help), the AI became more accurate and needed less data to learn. It stopped wasting energy on irrelevant details and focused on the most important clues.

In short: Instead of forcing the computer to treat all angles equally, this new method teaches it to be a "Chief Detective" who knows which angle is the most important and how to ask the right questions to the others. This makes the AI smarter, faster, and more like a human doctor.

Technical Summary

1. Problem Statement

Current 3D deep learning methods for medical image classification often fail to align with clinical workflows.

• Clinical Reality: Radiologists interpret 3D volumes (e.g., CT scans) using multiple anatomical planes (axial, coronal, sagittal). However, the axial plane is the primary acquisition and diagnostic reference, while coronal and sagittal planes serve as auxiliary views to provide complementary spatial context. This creates an inherently asymmetric dependency among planes.
• Limitations of Existing Models:
  • Holistic Processing: Many 3D models process the entire volume as a single block, ignoring the specific utility of individual planes.
  • Symmetric Multi-Plane Models: Existing multi-plane approaches often treat all planes with equal importance or fuse them symmetrically. This fails to capture the clinical reality where the axial plane is the "query" and other planes are "context," leading to suboptimal feature integration and high computational costs.
  • Feature Transferability: General Vision Foundation Models (VFMs) pretrained on natural images often exhibit a feature representation gap when applied to medical data.

2. Methodology

The authors propose an Axial-Centric Cross-Plane Attention Architecture that explicitly models the asymmetric relationship between anatomical planes. The pipeline consists of four main stages:

A. Feature Extraction (Frozen VFM)

• Backbone: The architecture uses MedDINOv3, a medical vision foundation model pretrained on 3.87 million axial CT images via self-supervised learning.
• Implementation: MedDINOv3 is kept frozen and used as a feature extractor for all three planes (axial, coronal, sagittal).
• Input Processing: 3D volumes are sliced into 2D planes. Each slice is processed by MedDINOv3 to generate dense feature embeddings. These slice-wise embeddings are stacked along the depth dimension to form a sequence ($B \times D \times E$).

B. Positional Context Modeling (RICA Blocks)

• Challenge: While MedDINOv3 provides semantic features, it lacks explicit volumetric positional context for the stacked slices.
• Solution: RICA (Residual Inception with Channel Attention) blocks are applied to the stacked feature sequences. These CNN-based modules interpret the sequence as a pseudo-2D map to emphasize positional context and improve generalization.

C. Intra-Plane Contextual Self-Attention

• Mechanism: Independent Transformer Encoders are applied to the feature sequences of each plane (axial, coronal, sagittal).
• Goal: To capture long-range dependencies between non-adjacent slices within the same anatomical plane.
• Output: This produces three distinct global contextual embeddings (one per plane), preserving plane-specific semantics before fusion.

D. Axial-Centric Cross-Plane Attention (Core Innovation)

• Mechanism: Two separate Cross-Plane Transformer Encoders are used to fuse information.
  • Query (Q): Derived exclusively from the Axial plane embeddings.
  • Key (K) & Value (V): Derived from the Coronal and Sagittal plane embeddings (auxiliary planes).
• Asymmetry: This design forces the model to condition the primary axial features on the complementary information from auxiliary planes, rather than fusing them symmetrically.
• Residual Connection: Notably, the residual connection after cross-attention is omitted to prevent bias toward the original axial features and to enforce effective conditioning.
• Classification: The updated CLS tokens from the fused axial-coronal and axial-sagittal representations are passed to separate MLP heads. The final logits are averaged to produce the classification prediction.

3. Key Contributions

1. Axial-Centric Architecture: A novel design that mirrors clinical reasoning by treating the axial plane as the primary query and auxiliary planes as context, capturing asymmetric inter-plane dependencies.
2. Data-Efficient Foundation Model Integration: The use of a frozen, medical-specific VFM (MedDINOv3) allows for robust feature extraction with minimal trainable parameters, enhancing data efficiency.
3. Directional Fusion Mechanism: The introduction of a specific cross-attention mechanism where axial features query auxiliary features, validated by ablation studies to be superior to symmetric or sequential fusion.

4. Experimental Results

The model was evaluated on the MedMNIST3D benchmark, covering six diverse datasets (Organ, Fracture, Vessel, Synapse, Nodule, Adrenal) with modalities including CT, MRA, and electron microscopy.

• Performance: The proposed architecture achieved the highest accuracy on 5 out of 6 datasets and the highest AUC on 3 out of 6 datasets.
  • Notable Gains: Outperformed the second-best method by 3.9% on FractureMNIST3D and 4.4% on AdrenalMNIST3D.
• Comparison: It consistently outperformed existing 3D models (ResNet-18, ViViT, Medformer, etc.) and multi-plane baselines.
• Ablation Studies:
  • QKV Allocation: Reversing the allocation (using auxiliary planes as queries) caused significant performance drops, confirming the necessity of the axial-centric approach.
  • Fusion Strategy: Replacing the dual parallel cross-attention with a sequential fusion (Axial $\to$ Coronal $\to$ Sagittal) resulted in lower performance, proving that independent conditioning is more effective.
  • Capacity: Reducing model size had a minor impact compared to the architectural design, suggesting the design itself is the primary driver of performance.

5. Significance

• Clinical Alignment: The paper demonstrates that aligning deep learning architecture with human clinical interpretation workflows (axial-primary) significantly improves model robustness and accuracy.
• Efficiency: By leveraging a frozen medical foundation model and focusing on efficient cross-plane attention, the method achieves state-of-the-art results without the computational overhead of training massive 3D models from scratch.
• Generalizability: The success across diverse modalities (CT, MRA, Microscopy) suggests that axial-centric representation learning is a universal strategy for 3D medical image analysis.

In conclusion, this work bridges the gap between clinical practice and AI design, proving that asymmetric, clinically-inspired attention mechanisms are superior to symmetric or holistic approaches for 3D medical image classification.