Why mamba is effective? Exploit Linear Transformer-Mamba Network for Multi-Modality Image Fusion
Chenguang Zhu Shan Gao Huafeng Chen Guangqian Guo Chaowei Wang Yaoxing Wang Chen Shu Lei Quanjiang Fan


[
    Amy X. Zhang
    September 5, 2024
=====================





§ ABSTRACT 

Multi-modality image fusion aims to integrate the merits of images from different sources and render high-quality fusion images. However, existing feature extraction and fusion methods are either constrained by inherent local reduction bias and static parameters during inference (CNN) or limited by quadratic computational complexity (Transformers), and cannot effectively extract and fuse features. To solve this problem, we propose a dual-branch image fusion network called Tmamba. It consists of linear Transformer and Mamba, which has global modeling capabilities while maintaining linear complexity. Due to the difference between the Transformer and Mamba structures, the features extracted by the two branches carry channel and position information respectively. T-M interaction structure is designed between the two branches, using global learnable parameters and convolutional layers to transfer position and channel information respectively. We further propose cross-modal interaction at the attention level to obtain cross-modal attention. Experiments show that our Tmamba achieves promising results in multiple fusion tasks, including infrared-visible image fusion and medical image fusion.  Code with checkpoints will be available after the peer-review process.
 














§ INTRODUCTION
 
Image fusion can integrate important image information from different data sources to render fused images with rich information.
Traditional image fusion approaches typically employ hand-crafted feature extraction and fusion rules. With the advancement of deep learning, image fusion methods based on deep learning have attained remarkable results <cit.>. In recent years, the networks used for feature extraction and image fusion are mostly built by CNN and Transformer. CNN-based methods have difficulty in capturing global context due to their limited receptive field, which makes them challenging to obtain high-quality fused images <cit.>. The convolution kernel of CNN has static parameters during the inference phase, which reduces the adaptability of CNN to different inputs <cit.>.
The Transformer-based methods perform well in global modeling, but the quadratic complexity of self-attention leads to high computational overhead. People use CNN and Transformer to form the dual-branch network to achieve the complementarity of network structure, but the problems of low input adaptability of CNN in the inference stage and quadratic complexity of Transformer still exist, which limits its application in the field of image fusion.

Recently, Mamba <cit.> provides us more options. Compared with CNN, Mamba has global modeling capabilities, and its selective scanning mechanism enables it to have high input adaptability. Compared with Transformer, Mamba has linear complexity. However, compared with the dual-branch network composed of CNN and Transformer, the ability of Mamba to extract features is obviously insufficient. The richness of features has a great influence on the effect of image fusion task. Therefore, how to customize a dual-branch network for Mamba suitable for multi modality image fusion is a compelling research issue.

In this paper, we propose the Tmamba block and use it to build a fusion network. It is a dual-branch network tailored for Mamba based on image fusion requirements. We considered three main points when choosing the branch to pair with Mamba: First, Mamba is good at processing long sequence inputs, so the branch paired with it should maintain acceptable computational complexity under long sequence inputs. Second, the branch paired with Mamba needs to be complementary and optimized with Mamba. Third, the branch paired with Mamba cannot undermine Mamba's advantages in image fusion. We choose Restormer block <cit.> and Vmamba block <cit.> to build the dual-branch Tmamba block.  Restormer block is a Transformer with linear complexity, which allows us to use pixel-level sequences while maintaining reasonable computational complexity. The design principles and network structures of Restormer block and Vmamba block are different, which enables them to extract features of different patterns as shown in Figure 1 (a) and (b). Diverse features can provide richer information for fusion. In addition, different network structures also make the features carry different information. Restormer block implicitly implements global attention through channel interactions, so that the features it extracts carry channel attention information. Mamba’s forgetting mechanism gives position information to the features it extracts <cit.>.  Unlike traditional dual-branch networks, two branches in our network are not completely independent. We added an interaction structure between the two branches to help complementary information pass to each other. Both Transformer and Mamba can adjust parameter matrices according to different inputs for targeted inference. This enables our method to fully extract and preserve information from different modalities. As shown in Figure 1 (c), entropy (EN) reflects the amount of information in the fused image, and mutual information (MI) reflects the similarity between the fused image and the input images. Our method can not only extract rich information, but also well preserve the features of the input images.




Furthermore, considering that the original self-attention of a single modality may be highly restricted by the modality information, we designed a cross-modal interaction at the attention level to get cross-modal attention.

In summary, the contributions of our work are as follows:



  * We innovatively built a Transformer-Mamba hybrid framework for multi-modality fusion tasks, and designed a hierarchical interaction strategy between Mamba and Transformer to optimize the features extracted by each other.

  * We proposed a cross-modal interaction at the attention level to break the limit of single modality information on attention, and get a more favorable  attention for image fusion .

  * Our model is highly adaptable to image inputs of multiple modalities and is able to perform targeted inferences based on the input images. This enables us to obtain state-of-the-art results on multiple datasets across two tasks by training on merely one dataset.







§  RELATED WORK



 §.§ DeepLearning-based Methods for Image Fusion

With the advancement of artificial intelligence technology, DL-based methods have become the mainstream in the field of image fusion.  DL-based image fusion methods can be roughly divided into three categories: modals based on generative method <cit.>, Autoencoder based models <cit.>  and models that combine downstream tasks <cit.>. (1)Among the image fusion methods based on generative methods, GAN-based methods are more commonly used. GAN-based models integrate different visual information by establishing an adversarial game between the original image and the fusion result, generating image results with richer and more diverse content. Recently, <cit.> proposed generative image fusion networks based on diffusion models that leverage powerful generative priors to address challenges such as training instability and lack of interpretability of GAN-based generative methods. (2)Autoencoder-based models extract image features using the encoder, then fuse the extracted image features, and finally output the final image through the decoder. (3)Many works combinine multi-modal image fusion with downstream tasks, <cit.> makes the model obtain more semantic information by introducing segmentation loss,  <cit.> explores methods to combine image fusion with detection, and <cit.> achieves good results in both tasks through the joint optimization of fusion network and segmentation network.



 §.§ State Space Model

SSM is a fundamental scientific model used in control theory. In recent years, efforts have been made to apply it to deep learning related tasks. LSSL <cit.> is the first to introduce SSM into the field of deep learning, demonstrating its significant advantages in handling long-sequence speech classification tasks. S4 <cit.> introduces low-rank correction adjustment based on LSSL and reduces the computational complexity of SSM. Subsequent works S5 <cit.> and H3 <cit.> improve SSM to make it better adapted to tasks in the field of deep learning. Recently, the emergence of Mamba <cit.> has once again triggered a craze for SSM in the field of deep learning. It can selectively extract input features and also integrate  hardware-aware algorithm. In the computer vision community, many models based on Mamba have also emerged. Vision mamba <cit.> and Vmamba <cit.> provide spatial perception capabilities for Mamba through bidirectional sequence modeling and cross-scanning, respectively. Many mamba-based models <cit.> have also emerged and achieved good results In the multi-modality image fusion. They all choose to use Mamba instead of Transformer and CNN, but did not try to combine them.





§ METHOD




 §.§ Overview


Our network can be specifically divided into five parts. Shallow feature extractor, Tmamba block for further feature encoding and feature interaction, Cross-modal interaction module for achieving cross-modal attention, Tmamba fusion module for fusing images of different patterns separately and Decoder for rendering the final fused image. All Transformer structures in our network are built by Restormer blocks <cit.> and all Mamba structures in our network are built by Vmamba blocks <cit.>. Our network is a generic image fusion network and we will explain our work using  Infrared-Visible Image Fusion as an example.



 §.§ Shallow feature extractor
 

Here we define the input paired infrared and visible images as I∈ℝ^ H× W and V∈ℝ^H× W.  The shallow feature extractor, Mamba and the Transformer blocks are represented by 𝒮(·), ℳ(·) and 𝒯(·) respectively.

We first use Transformer blocks to form a shallow feature extractor to extract shallow features {Φ_I^S, Φ_V^S} from infrared and visible inputs {I , V} :


    Φ_I^S = 𝒮(I) ,Φ_V^S = 𝒮(V).



In order to take advantage of Mamba in processing long sequence features, the shallow feature extractor processes the input image into 128×128×64 features. 



 §.§ Tmamba Block
 

We build Tmamba blocks with two layers of Mamba blocks(Mamba branch) and two layers of Transformer blocks(Transformer branch). We add an interaction structure between the two branches so that the two branches can transfer information to each other while maintaining feature differences.



  §.§.§ Transformer Branch.
  The shallow features extracted by the shallow feature extractor are input into the Transformer branch:


    Φ_I^trans,Φ_V^trans = 𝒯(Φ_I^S,Φ_V^S),

where Φ_{I,V}^trans represent features extracted from the first Transformer block. 

Restormer block can implicitly model global features by self-attention across channels.   Φ_{I,V}^trans are encoded as Q∈ℝ^ HW× C, K∈ℝ^C× HW, and  V∈ℝ^HW× C. We interact Q and K in the channel dimension to get the channel attention matrix with dimensions C × C :


    𝒜 = Softmax(K × Q  / α),

where 𝒜∈ℝ^C× C, α is a scaling parameter.
The elements in this matrix reflect the attention relationship between the corresponding channels. After interacting with V, this information will be passed to Φ_{I,V}^trans



  §.§.§ Mmaba Branch.
 The input features of the Mamba branch are the same as those of the Transformer branch:


    Φ_I^vm,Φ_V^vm = ℳ(Φ_I^S,Φ_V^S),

where Φ_{I,V}^vm represent features extracted from the first Mamba block respectively.

 Mamba’s forgetting mechanism ensures that the previous hidden state is always attenuated until the current token
is reached, which makes Mamba sensitive to the order of the input sequence and has position awareness <cit.>. This allows each feature pixel in Φ_{I,V}^vm to sense the positions of other feature pixels.



  §.§.§ T-M Interaction.
 












The information extracted by the two branches can complement each other. The channel attention information of the features extracted by the Transformer branch is exactly what the Mamba branch lacks, and the position information of the features extracted by the Mamba branch can also help the Transformer better extract and parse the input image without position encoding.

We designed different interaction structures based on the characteristics of information to help the two types of  information transmit between the Transformer branch and the Mamba branch.
To be specific, we choose a method similar to positional encoding to pass positional information  to the Transformer block. We align and add the features with position information  to the features extracted by the Transformer branch. In the process of addition, we set a global learnable parameter ω, which controls the proportion of information transmitted from the Mamba branch to the Transformer branch without interfering with the position information  encoded by Mamba:


    Φ_I^T = 𝒯(ω·Φ_I^vm+(1-ω)·Φ_I^trans),
    
    Φ_V^T = 𝒯(ω·Φ_V^vm+(1-ω)·Φ_V^trans) ,

where Φ_{V,I}^T are the final outputs of the Transformer branch

For features with channel attention, we use 1×1 convolution kernel to mix them  with features extracted by Mamba branch, and then use 3×3 convolution kernel to aggregate local features in space and send the aggregated features into Mamba block:


    Φ_I^M = ℳ(𝒞onv(Φ_I^vm,Φ_I^T)),
    
    Φ_V^M = ℳ(𝒞onv(Φ_V^vm,Φ_V^T)),

where Conv(·) stands for convolution operation and Φ_{V,I}^M are the final outputs of the Mamba branch.

The interactive process allows the features of the two branches to be optimized while maintaining their differences, which enables the Tmamba block to extract rich and high-quality image features.



 §.§ Cross-modality Interaction

After passing through the two branches of the Tmamba block, each modality (infrared and visible) has two different patterns of features due to the different network structures of the two branches.  Before entering the fusion module, we integrate and process the same pattern features of different modalities. For the features extracted from the Mamba branchs, we align and add the features at element wise to obtain the pre-fusion features:


    Φ_P^M = Φ_I^M + Φ_V^M,

where Φ_P^M is the pre-fusion feature of the Mamba branchs.

For the features extracted from the Transformer branchs, we adopt the  cross-modality interaction at attention-level , which simply refers to 𝒜ℱ.

The 𝒜ℱ block consists of two attention blocks and a weight calculation block 𝒲(·). Specifically, we first obtain the  attention matrices of the infrared and visible modality through the operation of channel interaction respectively. In the process of calculating attention, the features extracted by the network are encoded as Q∈ℝ^ HW× C,K∈ℝ^C× HW, and  V∈ℝ^HW× C. We multiply K and Q matrix to get  attention matrix with dimensions C × C.


    𝒜_V    = Softmax(K_V × Q_V  / α),
    
    𝒜_I    = Softmax(K_I × Q_I  / β) ,

where 𝒜_V,𝒜_I∈ℝ^C× C, α and β are scaling parameters.

Since there is interaction between channels in the calculation process of the Transformer branch’s attention, the attention matrix reflects the attention of each channel to other channels. Obviously, for different modalities of the same scene, the attention relationship between the various channels is different. We choose the DenseASPP to build  𝒲(·) to generate weights and get the cross-modal attention. 

Specifically, we first use DenseASPP blocks to further encode the features of the two modalities, and then connect the encoded features in the channel dimension and send them into the full connection layer to obtain two weights:


    ω_1 , ω_2 = 𝒲(Φ_V^T,Φ_I^T).
    
    𝒜 = ω_1  ·𝒜_V + ω_2 ·𝒜_I ,



Finally, we apply the cross-modal attention to the V of the two modalities separately, and then add them at element wise to obtain the pre-fusion feature:


    Φ_P^T = 𝒜· V_I + 𝒜· V_V

where Φ_P^T is the pre-fusion feature of the Transformer branchs, V_I and V_V ∈ℝ^HW× C are the encoded  matrices.





 §.§ Tmamba Fusion


The two patterns of features extracted by the Transformer and Mamba branches are fused separately in the fusion layer. Considering that the inductive bias of feature fusion should be similar to that of feature extraction, we still adopt Tmamba block for fusion. For the fusion block of Mamba's features, we only select the output of Mamba branch as the final output. and for fusion block of Transformer's features, we only select the output of Transformer branch as the final output:


    Φ_M = ℱ_ℳ(Φ_P^M)  ,Φ_T = ℱ_𝒯(Φ_P^T),

where ℱ_ℳ(·) and ℱ_𝒯(·)
 are the fusion blocks of Transformer's features and Transformer's features respectively.



 


 §.§ Decoder
 
We keep the decoder structure consistent with the design of shallow feature encoder, using the Restormer block as the basic unit of the decoder.

Our training process is divided into two stages. In the first stage, the model does not fuse images, but performs image restoration. In the second stage, we add the Cross-modality Interaction module and Tmamba fusion module to the training to obtain high quality fusion images:

    Stage1: V = 𝒟(Φ_V^T,Φ_V^M) ,  I = 𝒟(Φ_I^T,Φ_I^M).
    
       Stage2: F = 𝒟(Φ^T,Φ^M),

where 𝒟(·) indicates that the Decoder module, F is the final fusion image.



 §.§ Loss Function


The loss of Stage1 is the image reconstruction loss, which is designed to guide the encoder and decoder to learn the basic feature extraction and image reconstruction methods. The specific loss function is as follows:


    ℒ_1 = ℒ_ir + ℒ_vis,

where ℒ_ir and ℒ_vis are reconstruction losses of infrared image and visible image respectively.The reconstruction loss of infrared image can be concretely written as:

    ℒ_ir = ℒ^1_int(I,Î) + ℒ_SSIM(I,Î),

where ℒ^1_int(I,Î) = I-Î^2_2 and ℒ_SSIM(I,Î) = 1-SSIM(I,Î),SSIM(·,·)is the structural similarity index<cit.>.The reconstruction loss of visible image can also be obtained by the same method.

The loss of Stage2 is used to guide the process of image fusion, which mainly consists of intensity loss and gradient loss:

    ℒ_2 = ℒ^2_int + ℒ_grad

where ℒ^2_int = 1/HWI_f - max(I_ir,I_vis)_1 and ℒ_grad = 1/HW = |∇ I_f| - max(|∇ I_ir|,|∇ I_vis|)_1,where ∇ indicates the Sobel gradient operator.



§ EXPERIMENTS




 §.§ Infrared and visible image fusion



  §.§.§ Setup

 IVF experiments use three popular benchmarks to verify our fusion model, i.e., TNO <cit.>, RoadScene<cit.>, and  MSRS <cit.>. We train our model on MSRS training set (1083 pairs) and test it on MSRS test set (361 pairs), RoadScene (50 pairs) and TNO (25 pairs). Note that we did not retrain on the TNO and RoadScene datasets, but directly tested on the test sets of the three datasets using the model trained on the MSRS. We use six metrics to quantitatively measure the fusion results: entropy (EN), standard deviation (SD), spatial frequency (SF), mutual information (MI) ,visual information fidelity (VIF) and QAB/F . Higher metrics indicate that a fusion image is better. The details of these metrics can be found in <cit.>.

 Our experiments are carried out on a machine with one NVIDIA GeForce RTX 3090 GPUs. The training samples are randomly cropped into 128×128 patches in the preprocessing stage. The number of epochs for training is set to 80 with 40 and 40 epochs in the first and second stage respectively. The batch size is set to 4. We adopt the Adam optimizer with the initial learning rate set to 7.5 × 10^-5 and decreasing by 0.5 every 20 epochs. 





  §.§.§ Comparison with advanced methods.


In this section, we test Tmamba on the three test sets and compare the fusion results




the state-of-the-art methods including DIDFuse<cit.>, U2Fusion<cit.>, TarDAL<cit.>, ReCoNet <cit.>,CDDFuse<cit.>,DDFM<cit.>,EMMA<cit.>.

For qualitative comparison, We show the qualitative comparison in Figure 3 and Figure 4. Compared with other existing methods, our Tmamba has two significant advantages. First, the discriminative target from infrared images can be well preserved.The people in Figures 3 and Figure 4 are all highlighted, and our method has a higher contrast and clearer outline compared to other methods. Second, our method can get richer background information.As shown in Figure 4 (the wires and poles in the orange box), our method clearly restores the outline of the wires and poles, while visual inspection shows that other methods do not bring out this background information.

For quantitative comparisons, a total of six metrics are employed to quantitatively compare the above results, which are displayed in Table 1. Our method has excellent performance on almost all metrics.





 §.§ Ablation studies

Ablation experiments are set to verify the rationality of the different modules.We did a total of five experiments to test our model as shown in Table 2 and Table 3.



  §.§.§ Transformer branch.
 We remove the Mamba branch from the Tmamba block, leaving only the Transformer branch, while removing the cross-modality interaction of the Transformer branch. The test results show that all the indexes declined. 



  §.§.§ Cross-modality interaction.
 We add the cross-modality interaction in the Transformer branch.It makes the attention relationship of each channel no longer restricted by single-modal information, which comprehensively improves the effect of image fusion.



  §.§.§ Mamba branch.
 We add the Mamba branch to form a dual-branch structure with Transformer. The test results show that all the indexes improve.  The dual-branch structure composed of Transformer and Mamba can encode image features in different patterns and then fuse them separately, so that the acquired image information is richer and the quality of the fused image is greatly improved.



  §.§.§ Interaction structure.

We added hierarchical interaction structure between Transformer branch and Mamba branch to form a complete Tmamba block. By adding  hierarchical interaction structure, the  information carried by the the features  is continuously exchanged between the two branches, further improving the fusion effect.



  §.§.§ Why mamba is effective.
  We replace the Mamba block with a residual convolution module. The SF index of the fused image has improved, but the overall quality has decreased, as shown in Table 3. We believe that there are two main reasons. On the one hand, in the long sequence image fusion task, Mamba can show feature extraction and feature analysis capabilities that surpass CNN; on the other hand, we use Transformer and Mamba to build the network backbone, which makes our network input-aware and can adjust parameters for different inputs, making it more suitable for multimodal image fusion tasks than the Transformer-CNN hybrid network.


















 §.§ Medical image fusion




  §.§.§ Setup.
 We selected 136 pairs of medical images from the Harvard Medical website for MIF experiments, of which 21 pairs of MRI-CT images, 42 pairs of MRIPET images and 73 pairs of MRI-SPECT images are utilized as the test datasets.








  §.§.§ Comparison with advanced methods.
  In this group of experiments, we selected five networks, U2Fusion <cit.> TarDAL<cit.>, ReCoNet <cit.>, CDDFuse<cit.>, and EMMA<cit.>, to conduct comparative tests.We made quantitative and qualitative comparisons across all networks. Note that none of the models were fine-tuned on the medical data set. 

For qualitative comparison,we visualize the fusion results on the  MRI-CT dataset in Figure 5. The results of quantitative comparison are shown in Table.4.  As you can see, our method achieves leading performance on all datasets.







§ CONCLUSION


In this paper, we proposed a dual-branch Transformer-Mamba network for multi-modality image fusion. By combining the Restormer block and the Vammba block, we extracted features containing different information. We interacted and fused these features and achieved better results. We further proposed cross-modality Interaction at the attention level to break the limitation of single-modality information on attention. Experiments show that our method achieves advanced performance on six datasets across two image fusion tasks.









