Structured Modality-Aware Token Interaction for Multimodal Medical Imaging

1 Introduction

Transformer architectures have reshaped representation learning by enabling explicit modeling of long-range dependencies through self-attention [Vaswani et al., 2017]. In medical imaging, this property is particularly valuable for volumetric data, where global anatomical context plays a critical role. When applied to multimodal inputs, however, most existing architectures still rely on implicit fusion, typically implemented by concatenating modalities along the channel dimension at the network input. This strategy assumes that heterogeneous modalities can be processed within a single shared representation space without explicit structural constraints. In practice, standard self-attention tends to conflate modality-specific feature learning and cross-modal fusion into a single unstructured operation, forcing the network to simultaneously disentangle modality-specific representations and learn how to combine them. As a result, multimodal interaction remains implicit and difficult to control.

∗Corresponding author

Several recent works have explored explicit modalityaware designs based on multi-branch encoders, crossattention mechanisms, or modality-specific transformer streams [Xing et al., 2022; Zhang et al., 2022; Sun, 2025]. Although effective, these approaches typically increase architectural complexity, memory consumption, and computational cost, and frequently rely on quadratic cross-attention between modality-specific token sets. This makes it difficult to disentangle the benefits of architectural inductive bias from those of increased model capacity. These limitations are particularly evident in multimodal 3D medical image segmentation, where accurate delineation requires the coherent integration of complementary sequences over large spatial contexts. In this setting, transformer-based models such as UNETR [Hatamizadeh et al., 2022b] and SwinUNETR [Hatamizadeh et al., 2022a] have demonstrated strong performance, but still largely rely on implicit or heavyweight fusion mechanisms.

In this work, we argue that structured modality interaction should be treated as a first-class architectural design principle in multimodal learning. We propose Modality-Aware Token Interaction (MATI), an architecturally lightweight and backbone-agnostic module that operates directly on token embeddings by structuring the channel space into modalityaligned subspaces. MATI explicitly separates modalitypreserving refinement from controlled cross-modal information exchange, injecting a structured multimodal inductive bias into a single-stream transformer representation. We instantiate MATI within the UNETR framework and propose two architectures: ModaUNETRS and ModaUNETRE . ModaUNETRS applies MATI to selected intermediate token representations used for skip connections, whereas ModaUNETRE injects MATI directly inside the transformer encoder to progressively shape the representation hierarchy. Both variants preserve the overall UNETR topology and introduce only a moderate computational overhead. We evaluate the proposed models on the BraTS 2020 benchmark, a challenging and well-established testbed for multimodal representation learning in medical imaging.

In summary, our contributions are threefold:

• We introduce MATI, an architecturally lightweight and backbone-agnostic modality-aware interaction module that structures multimodal fusion inside the token embedding space.
• We propose ModaUNETRS and ModaUNETRE , two instantiations of MATI within a UNETR backbone using respectively skip-level and in-encoder interaction.
• We demonstrate on BraTS 2020 that explicitly structuring multimodal interaction improves accuracy without substantially increasing computational complexity or introducing additional attention streams.

2 RelatedWork

Convolutional encoder-decoder architectures such as U-Net and its 3D variants [Ronneberger et al., 2015; C¸ ic¸ek et al., 2016] have long dominated medical image segmentation. Frameworks such as nnU-Net [Isensee et al., 2021] further showed that carefully engineered pipelines can achieve strong and reproducible performance. However, convolution-based models remain inherently limited by local receptive fields, reducing their ability to capture long-range spatial dependencies and global anatomical context.

Transformer-based architectures address this limitation through self-attention mechanisms that explicitly model long-range interactions. Following the success of Vision Transformers (ViTs) [Dosovitskiy et al., 2020], several transformer-based models have been proposed for medical image segmentation, including UNETR [Hatamizadeh et al., 2022b], TransUNet [Chen et al., 2021], and SwinUNETR [Hatamizadeh et al., 2022a]. These works established transformers as a powerful alternative to purely convolutional designs for volumetric segmentation.

In multimodal medical image segmentation, however, most transformer-based approaches still rely on implicit fusion by concatenating modalities at the input level [Myronenko, 2019; Fareed et al., 2025]. While effective, this strategy treats modalities as homogeneous information sources and offers no explicit control over modality-specific representation and interaction. To overcome this limitation, several works introduced explicit modality-aware designs based on multi-branch encoders, cross-attention, or auxiliary objectives, including mmFormer [Zhang et al., 2022], NestedFormer [Xing et al., 2022], and cross-modality distillation approaches [Lin et al., 2023]. Although these methods often improve performance, they typically require multiple modality-specific streams or full cross-attention modules, substantially increasing architectural complexity and computational cost.

In contrast, our approach embeds modality-aware structure directly within a single-stream transformer by factorizing the token embedding space into modality-aligned subspaces and regulating cross-modal interaction within it. This preserves the simplicity and scalability of architectures such as UNETR while introducing an explicit multimodal inductive bias, offering a lightweight alternative to both implicit early fusion and complex multi-stream transformer designs.

3 Methodology

3.1 Backbone

ModaUNETRS and ModaUNETRE are built upon UNETR [Hatamizadeh et al., 2022b], a transformer-based encoder-decoder architecture for volumetric medical image segmentation that combines global context modeling through a ViT encoder with high-resolution reconstruction via a convolutional decoder.

Following the standard UNETR configuration, input volumes of size 128 × 128 × 128 are partitioned into nonoverlapping 16 × 16 × 16 patches, producing 512 tokens embedded in a 768-dimensional space and processed by a 12-layer transformer encoder with 12 attention heads and MLP dimension 3072. Intermediate token representations are projected into spatial feature maps and forwarded to a hierarchical convolutional decoder through skip connections for multi-scale feature fusion. This design contains approximately 1.02 × 108 parameters and is kept unchanged across all experiments, providing a stable backbone for isolating the effect of the proposed modality-aware interaction.

3.2 ModaUNETRS

In multimodal medical imaging, UNETR processes different modalities by concatenating them along the channel dimension before tokenization. After patch embedding, modality information becomes fully entangled within the token embeddings, and all subsequent transformer layers operate in a shared representation space. Although effective, this design provides no explicit mechanism to preserve modality-specific structure or regulate cross-modal information exchange.

ModaUNETRS introduces an explicit yet low-overhead modality-structured inductive bias by augmenting UNETR with MATI blocks applied to selected transformer skip representations. MATI does not modify the ViT encoder itself. Instead, it refines intermediate token representations before they are projected and injected into the convolutional decoder through the standard UNETR skip connections.

Placement. As illustrated in Figure 1, MATI blocks are applied to the transformer outputs Z6 and Z9, corresponding to the outputs of the 6th and 9th ViT encoder blocks. These token sequences are the same representations used by UNETR to build two of the multi-scale skip connections into the convolutional decoder, hence IMATI = {6, 9}. This placement targets intermediate semantic representations, where global contextual information has already been established but before token-to-feature-map projection and spatial decoding. At these depths, tokens capture high-level semantic information while still preserving localized spatial correspondence, making them suitable for modality-aware refinement without interfering with either early low-level processing or late-stage spatial reconstruction.

Subspace-structured interaction. Let Z ∈ RB×N×C denote a transformer token sequence selected for refinement. MATI introduces an explicit modality-structured inductive bias by factorizing the channel dimension into M fixed subspaces, each associated with one input modality. In practice, we use M = 4 subspaces corresponding to the four Magnetic Resonance Imaging (MRI) modalities:

MATI then performs two complementary operations: intrasubspace refinement, which preserves modality-aligned processing through independent self-attention within each subspace, and inter-subspace mixing, which enables controlled cross-modal information exchange via a lightweight global gating mechanism. The complete MATI operation is summarized in Algorithm 1.

Intra- vs. inter-subspace. For each modality subspace, MATI applies a residual self-attention block:

This allows token interaction within each modality-aligned subspace while preserving structured multimodal representations.

To enable controlled cross-modality interaction without introducing full cross-attention over all tokens, MATI summarizes each refined subspace by average pooling over the token dimension, concatenates the resulting summary vectors, and predicts a set of mixing weights through a small gating network. These weights are used to form a global mixed feature that is broadcast back to all subspaces and added residually to each token representation. The final refined token sequence is then reconstructed by concatenation:

This design explicitly factorizes modality-preserving refinement from controlled modality interaction, instantiating a structured multimodal inductive bias while operating entirely within the original UNETR token space.

Complexity and integration. MATI introduces no additional token streams, parallel modality-specific encoders, or quadratic cross-attention across modalities. It preserves the token sequence length N and is applied only to a small number of skip representations. The additional parameters are limited to the per-subspace attention projections and the lightweight mixing network. Consequently, performance gains can be attributed to the introduced structural bias rather than a substantial increase in backbone capacity, whereas computational complexity remains dominated by the original UNETR architecture. After MATI refinement, the updated token representations are projected into spatial feature maps through the standard UNETR token-to-feature projection layers and injected into the convolutional decoder via unchanged skip connections. The complete forward pass of ModaUNETRS is summarized in Algorithm 2.

3.3 ModaUNETRE

While ModaUNETRS applies MATI only to selected transformer outputs later used as UNETR skip tokens,

Algorithm 1 Modality-Aware Token Interaction (MATI).

Require: Token sequence Z ∈ RB×N×C , number of modalities M

Ensure: Refined token sequence Z ′ ∈ RB×N×C

• 1: // Split channels into modality-aligned subspaces
• 2: {Z (1) , . . . ,Z (M) } ← SplitChannels(Z)
• 3: // Intra-subspace refinement
• 4: for m= 1 toM do
• 5: U (m) ←Z (m) +MHAm(LN(Z (m) ))
• 6: end for
• 7: // Compute subspace summaries
• 8: for m= 1 toM do
• 9: s(m) ←Meann(U (m) )
• 10: end for
• 11: s ←Concat(s(1) , . . . , s(M) )
• 12: // Predict mixing weights
• 13: α← σ(MLP(s))
• 14: // Inter-subspace mixing
• 15: G← ∑Mm=1 αm · U (m) {broadcasted to all tokens}
• 16: // Residual fusion and reconstruction
• 17: for m= 1 toM do
• 18: ˆ Z (m) ← U (m) + G
• 19: end for
• 20: Z ′ ←Concat( ˆ Z (1) , . . . , ˆ Z (M) )
• 21: return Z ′

Algorithm 2 ModaUNETRS forward pass with MATI skiptoken refinement.

Require: Multimodal input X ∈ RB×M×D×H×W Ensure: Segmentation prediction ˆ Y

• 1: // ViT encoder (standard UNETR)
• 2: {Z1, . . . ,Z12},Zout ←ViT(X)
• 3: // Extract UNETR skip tokens
• 4: S ← {Z3,Z6,Z9,Z12}
• 5: // MATI refinement on selected skip tokens
• 6: // MATI refinement indices: IMATI = {6, 9}
• 7: for ℓ ∈ IMATI do
• 8: Zℓ ←MATI(Zℓ)
• 9: end for
• 10: // Token-to-feature projection (standard UNETR)
• 11: E1 ←Encoder1(X)
• 12: E2 ←Encoder2(Proj(Z3))
• 13: E3 ←Encoder3(Proj(Z6)) {Z6 is MATI-refined}
• 14: E4 ←Encoder4(Proj(Z9)) {Z9 is MATI-refined}
• 15: E5 ←Proj(Zout)
• 16: // Convolutional decoder (standard UNETR)
• 17: D4 ←Decoder5(E5,E4)
• 18: D3 ←Decoder4(D4,E3)
• 19: D2 ←Decoder3(D3,E2)
• 20: D1 ←Decoder2(D2,E1)
• 21: ˆ Y ←Out(D1)
• 22: return ˆ Y

ModaUNETRE injects the same interaction mechanism directly inside the ViT encoder, allowing modality-aware interaction to shape the evolving token stream.

In-encoder MATI injection. Let {B0, . . . ,B11} denote the twelve transformer blocks of the ViT encoder and let Z0 = PatchEmbed(X) be the embedded token sequence. ModaUNETRE modifies the ViT forward pass by inserting MATI after the 3rd and 6th blocks, as illustrated in Figure 2. In our implementation, this corresponds to the 0-indexed set JMATI = {2, 5}. The token stream update becomes:

When applied, MATI updates the running token stream, which is then propagated through all subsequent transformer blocks. As a result, downstream layers operate on modalitystructured tokens rather than on the original entangled representation.

Skip extraction frommodified stream. UNETR uses four transformer outputs as skip tokens, then projected to feature maps. Consistent with the UNETR design, we tap the token stream at block indices {2, 5, 8, 11}, yielding {Z3, Z6, Z9, Z12} from the MATI-refined stream. These representations are projected into spatial feature maps and injected into the unchanged convolutional decoder through the standard UNETR skip connections.

Complexity and integration. ModaUNETRE preserves the overall UNETR topology (encoder-decoder structure, skip scheme, and token-to-feature projections) and only augments the ViT forward pass with two MATI modules. Consequently, the parameter and computational overhead remain limited, while modality-aware interaction influences both the decoder inputs and the internal transformer representation hierarchy. The complete forward pass of ModaUNETRE is summarized in Algorithm 3.

Algorithm 3 ModaUNETRE forward pass with MATI injected inside the ViT encoder.

Require: Multimodal input X ∈ RB×M×D×H×W Ensure: Segmentation prediction ˆ Y

• 1: // Patch embedding (standard UNETRViT front-end)
• 2: Z ←PatchEmbed(X)
• 3: H ← [ ] {list of tapped skip tokens}
• 4: // ViT encoder blocks withMATI injection
• 5: for i = 0 to 11 do
• 6: Z ←Bi(Z) {standard transformer block}
• 7: if i ∈ JMATI then
• 8: Z ← MATI(Z) {JMATI = {2, 5} (after 3rd and 6th blocks)}
• 9: end if
• 10: if i ∈ {2, 5, 8, 11} then
• 11: append Z to H {H = [Z3,Z6,Z9,Z12] from the modified token stream}
• 12: end if
• 13: end for
• 14: Zout ←LN(Z)
• 15: // Token-to-feature projection (standard UNETR)
• 16: E1 ←Encoder1(X)
• 17: E2 ←Encoder2(Proj(H[0])) {Z3}
• 18: E3 ←Encoder3(Proj(H[1])) {Z6}
• 19: E4 ←Encoder4(Proj(H[2])) {Z9}
• 20: E5 ←Proj(Zout)
• 21: // Convolutional decoder (standard UNETR)
• 22: D4 ←Decoder5(E5,E4)
• 23: D3 ←Decoder4(D4,E3)
• 24: D2 ←Decoder3(D3,E2)
• 25: D1 ←Decoder2(D2,E1)
• 26: ˆ Y ←Out(D1)
• 27: return ˆ Y

4 Experiments

4.1 Data

All experiments are conducted on the BraTS 2020 benchmark [Menze et al., 2015; Bakas et al., 2017; Spyridon et al., 2017; Bakas et al., 2018], which provides pre-operative multimodal MRI scans of gliomas acquired across multiple institutions and imaging protocols. Each case includes four MRI sequences: T1-weighted (T1), T1-weighted contrastenhanced (T1ce), T2-weighted (T2), and FLAIR. Groundtruth annotations define three tumor subregions: Enhancing Tumor (ET), Tumor Core (TC), and Whole Tumor (WT).

We use the official training set with five-fold crossvalidation and patient-level splits, and report state-of-the-art comparisons on the BraTS 2020 validation set.

4.2 Preprocessing

Volumes are provided co-registered, resampled to 1 mm3 isotropic resolution, and skull-stripped. For computational efficiency and architectural compatibility, all inputs are further resized to 128×128×128. Image intensities are resampled using trilinear interpolation, while segmentation labels use nearest-neighbor interpolation to preserve discrete class values.

4.3 Training Protocol

All models are implemented in PyTorch and trained on an NVIDIA A100 with mixed precision under a unified setup.

We use Adam with initial learning rate 10-4 , decayed by a factor of 0.5 every 10 epochs, and train for 100 epochs per fold with batch size 1. Standard 3D augmentations (flips, affine transforms, intensity perturbations, and bias-field augmentation) are applied only to training volumes. The loss combines voxel-wise cross-entropy (0.3) and soft Dice loss (0.7).

4.4 Evaluation Protocol

Performance is evaluated using Dice Similarity Coefficient (DSC), Intersection over Union (IoU), and Volume Similarity (VS) under five-fold cross-validation on the BraTS 2020 training set.

For state-of-the-art comparison, we report DSC on the BraTS 2020 validation set following standard benchmark practice. Computational efficiency is evaluated on the same validation set in terms of peak GPU memory consumption and inference time, providing a practical measure of the accuracy-cost trade-off across architectures.

4.5 Baseline Comparison

We compare ModaUNETRS and ModaUNETRE against representative convolutional and transformer-based architectures for 3D medical image segmentation. UNETR [Hatamizadeh et al., 2022b] serves as the primary transformer baseline and backbone reference. SwinUNETR [Hatamizadeh et al., 2022a] is included as a hierarchical transformer with shiftedwindow attention, representing a strong high-capacity design, whereas UNet++ [Zhou et al., 2018] provides a competitive convolutional baseline with nested skip connections.

All models are trained and evaluated under the same preprocessing, training, and evaluation protocols to ensure fair comparison.

4.6 Results

Quantitative vs. qualitative. Table 1 reports DSC, IoU, and VS for all subregions under five-fold cross-validation on the BraTS 2020 training set, together with DSC and efficiency metrics (GPU memory and inference time) on the official, held-out validation set. Figure 3 shows representative ground truth and predicted segmentations on the BraTS 2020 training set under five-fold cross-validation.

SwinUNETR performs very well on all subregions, but at higher computational cost. UNet++ yields the lowest resource usage but also the weakest segmentation performance, whereas UNETR provides a strong intermediate baseline. Both ModaUNETRS and ModaUNETRE consistently outperform UNETR across all subregions, with larger gains on TC and WT. In particular, ModaUNETRE further improves over ModaUNETRS while maintaining substantially lower inference time and memory footprint than SwinUNETR. The qualitative results mirror the quantitative findings. SwinUNETR produces spatially coherent segmentations but oversmooths boundaries. UNet++ often misses small or disconnected enhancing areas, whereas UNETR shows occasional under-segmentation. In contrast, ModaUNETRS yields more consistent tumor delineations with reduced fragmentation, and ModaUNETRE further improves structural coherence and tumor extent coverage in challenging cases.

State-of-the-art comparison. Table 2 directly compares ModaUNETRS and ModaUNETRE against representative state-of-the-art methods on the BraTS 2020 validation set in terms of DSC for ET, TC, and WT. One compared method [Wang et al., 2020] corresponds to a top-performing BraTS 2020 challenge entry, while the others [An et al., 2024; Kunjumon et al., 2025; Zhou et al., 2025] represent recent architectural trends.

Both ModaUNETRS and ModaUNETRE achieve competitive performance across all subregions. In particular, ModaUNETRE attains the best or tied-best DSC on ET and TC and matches the top WT performance, while ModaUNETRS remains consistently close to the strongest competing methods.

4.7 Ablation Study

To isolate the effect of architectural structure from parameter count, we compare UNETR, ModaUNETRS , and ModaUNETRE under controlled settings. UNETR serves as the baseline without explicit modality-aware interaction. ModaUNETRS applies MATI to selected skip-token representations, whereas ModaUNETRE injects MATI inside the ViT encoder for progressive token refinement. All models share the same backbone and differ only in the presence and placement of MATI.

As shown in Table 3, ModaUNETRS consistently improves over UNETR across all subregions, while ModaUNETRE further improves over ModaUNETRS with only moderate increases in memory usage and inference time. Since no modality-specific encoders, additional token streams, or cross-attention mechanisms are introduced, these gains cannot be attributed to increased capacity alone. Instead, they suggest that progressively deeper modality-aware interaction improves performance, with in-encoder injection proving even more effective than skip-level refinement.

5 Discussion

This work investigates whether multimodal interaction in medical imaging can be improved through an explicit architectural inductive bias, without relying on multi-branch encoders or quadratic cross-attention between modality-specific token streams.

Interaction inside token embedding space. MATI structures multimodal interaction directly within the token embedding space by factorizing channels into modality-aligned subspaces. This explicitly separates two roles often entangled in standard self-attention: intra-subspace attention preserves coherent modality-specific representations instead of immediately merging them into a shared space, while inter-subspace mixing enables controlled global interaction via a compact gating mechanism. This separation reduces the burden on generic self-attention to jointly infer modality structure and semantic fusion.

Skip-level vs. in-encoder injection. In ModaUNETRS , MATI refines only selected transformer outputs used as skip tokens, leaving the ViT encoder unchanged. In contrast, ModaUNETRE injects MATI directly into the ViT encoder, allowing the refined token stream to propagate through subsequent transformer blocks and influence the entire representation hierarchy. The consistent improvement of ModaUNETRE over ModaUNETRS across all folds and subregions (Table 1) suggests that early and intermediate structured interaction is more effective than applying modalityaware refinement only at the decoder interface.

Accuracy-efficiency trade-off. MATI-equipped models provide a favorable middle ground, narrowing the accuracy gap with high-capacity transformers such as SwinUNETR while maintaining lower resource usage (Table 1). This trend is further supported by the state-of-the-art comparison (Table 2), where ModaUNETRE matches or achieves top performance despite using a simpler single-stream architecture. These results indicate that structured modality-aware interaction can recover much of the performance of heavier multimodal designs at lower computational burden.

Structure vs. capacity. A central motivation of this work is to disentangle architectural inductive bias from brute-force scaling. The ablation study (Table 3) comparing UNETR, ModaUNETRS , and ModaUNETRE shows monotonic performance gains despite only moderate increases in memory usage and inference time. Since these improvements are achieved without increasing the number of tokens, introducing additional attention streams, or using modality-specific encoders, the gains are likely driven by the imposed architectural structure rather than increased model capacity.

Stability and robustness. Although we do not conduct a dedicated multi-seed variance analysis, the improvements of ModaUNETRS and ModaUNETRE over UNETR are consistent across all folds and subregions (Tables 1 and 3), and we did not observe significant fold-to-fold variance or unstable training despite batch size 1.

Limitations and future work. MATI currently assumes a fixed modality-aligned subspace partition, which fits the BraTS setting where modalities are well-defined and consistently available. In scenarios with missing, corrupted, weakly correlated, or highly redundant modalities, this assumption may be partially violated and may also reduce the interpretability of fixed subspace assignments, motivating adaptive or learned allocation strategies. Moreover, this work is limited to a single backbone (UNETR) and benchmark (BraTS 2020). Since MATI is backbone-agnostic and operates directly on token embeddings, future extensions will validate it across additional architectures and datasets, also investigating adaptive handling of missing or unreliable modalities and dynamic subspace allocation.

6 Conclusions

This work addresses a key limitation of current transformerbased multimodal models: the lack of explicit architectural control over modality interaction in the representation space. We thus introduced MATI, an architecturally lightweight and backbone-agnostic interaction module that injects a structured multimodal inductive bias by factorizing token embeddings into modality-aligned subspaces and regulating crossmodal information exchange. We instantiated MATI within UNETR in two variants: ModaUNETRS applies structured interaction to selected skip-token representations, whereas ModaUNETRE injects it directly into the transformer encoder for progressive modality-aware refinement.

Extensive experiments on BraTS 2020 show that both ModaUNETRS and ModaUNETRE consistently improve over UNETR under five-fold cross-validation and official, held-out validation, while preserving a single-stream architecture and maintaining a favorable accuracy-efficiency tradeoff. More broadly, these results suggest that structured interaction should be treated as a first-class architectural principle in multimodal transformers rather than an emergent property of generic self-attention. Explicitly encoding modality structure within the token embedding space enables more coherent, robust, and efficient multimodal representations.

Acknowledgements

This research was partially funded by the Horizon Europe Research and Innovation Programme under Grant No. 101070537, project CROSSCON-Cross-Platform Open Security Stack for Connected Devices. Additional support was provided by the PNRR project FAIR-Future AI Research (PE00000013), within the NRRP MUR Programme funded by NextGenerationEU, and by the AI@TN2.0 project funded by the Autonomous Province of Trento, Italy.