MaDis-Stereo: Enhanced Stereo Matching via Distilled Masked Image Modeling
Jihye Ahn Hyesong Choi Soomin Kim Dongbo Min


Subspace-thermal discrete time crystals from phase transitions between different n-tuple discrete time crystals
    Tzu-Chieh Wei
    Accepted Sep 3 2024 to ApJ Letters

===============================================================================================================



§ ABSTRACT 

In stereo matching, CNNs have traditionally served as the predominant architectures. Although Transformer-based stereo models have been studied recently, their performance still lags behind CNN-based stereo models due to the inherent data scarcity issue in the stereo matching task. In this paper, we propose Masked Image Modeling Distilled Stereo matching model, termed MaDis-Stereo, that enhances locality inductive bias by leveraging Masked Image Modeling (MIM) in training Transformer-based stereo model. Given randomly masked stereo images as inputs, our method attempts to conduct both image reconstruction and depth prediction tasks. While this strategy is beneficial to resolving the data scarcity issue, the dual challenge of reconstructing masked tokens and subsequently performing stereo matching poses significant challenges, particularly in terms of training stability. To address this, we propose to use an auxiliary network (teacher), updated via Exponential Moving Average (EMA), along with the original stereo model (student), where teacher predictions serve as pseudo supervisory signals to effectively distill knowledge into the student model. State-of-the-arts performance is achieved with the proposed method on several stereo matching such as ETH3D and KITTI 2015. Additionally, to demonstrate that our model effectively leverages locality inductive bias, we provide the attention distance measurement.





§ INTRODUCTION


Stereo depth estimation is a critical task in computer vision, focused on predicting disparities between stereo image pairs. Its importance is widely recognized across various applications, including autonomous driving <cit.>, SLAM <cit.>, robotic control <cit.>, drone navigation <cit.>, and beyond. Recently, the performance of the stereo matching has been improved remarkably by adopting deep neural networks (DNNs), but the difficulty of collecting large-scale ground truth training data using costly equipment such as LiDAR <cit.> becomes a factor that hinders the advance of the stereo matching task.

For this reason, the stereo models based on Convolutional Neural Networks (CNNs) <cit.> are still a widely used solution for stereo matching tasks, unlike other vision tasks where the Transformer architectures <cit.> are gradually becoming mainstream. In stereo matching, the CNNs-based models outperform Transformer-based counterparts <cit.> due to the training data efficiency resulting from the local inductive bias inherent in convolutional operations.

This phenomenon contrasts with the prevailing trends observed in most computer vision tasks. Since the introduction of ViT <cit.>, Transformer-based approaches have been recognized for their capacity to effectively learn global representations from large-scale datasets, achieving state-of-the-art results across numerous vision applications such as image classification <cit.>, 2D/3D object detection <cit.>, and semantic segmentation <cit.>. 
Since the Transformer typically requires extensive training data to obtain overwhelming performance, the currently limited ground truth data for stereo matching is an obstacle to achieving competitive performance in the Transformer-based stereo matching networks.



Nonetheless, considering the potential of the ever-evolving Transformer architecture, a few attempts have been made to overcome the above-mentioned limitations of Transformer-based stereo matching networks.
CroCo-Stereo <cit.> partially alleviates the data scarcity issue by leveraging the pre-trained encoder as a backbone. The pre-trained encoder <cit.> is tailored for geometric downstream tasks by performing cross-view completion for massive stereo image pairs. While this approach addresses the data scarcity issue to some extent through the introduction of the pre-training methodology based on cross-view completion, Transformer models—designed primarily for learning global representations with extensive training data—continue to struggle with the complexities of stereo matching tasks.
In the fine-tuning stage for stereo matching, a sufficient amount of training data with ground truth labeled data is still necessary.

To address this data scarcity issue, it is crucial to incorporate locality inductive bias when fine-tuning the whole stereo model, in addition to pre-training the Transformer encoder. Recently, it was reported in <cit.> that Masked Image Modeling (MIM), which tends to concentrate around pixels in attention heads, introduces locality inductive bias in pre-trained Transformer models. The effectiveness of this locality is closely related to the masking ratio and masked patch size during pre-training.

Inspired by the benefits from MIM <cit.>, we propose a novel approach, termed Masked Image Modeling Distilled Stereo matching model (MaDis-Stereo), which is a Transformer-based model that incorporates MIM <cit.> that aims to impart locality inductive bias to the stereo depth estimation. This is achieved by pre-processing the stereo images through random masking and subsequently reconstructing the masked regions of images within the stereo matching network as described in Fig. <ref>.

However, the straightforward application of the MIM without appropriate measures may not suit fine-tuning tasks such as stereo depth estimation well. This challenge stems from the discrepancy between the pre-training task, where MIM is utilized, and the subsequent fine-tuning tasks. In the MIM based pre-training, the primary objective is to mask and reconstruct the image. In contrast, when applying the MIM strategy directly to the stereo depth estimation model, the training process becomes more complex. In addition to the reconstruction of masked patches, it must predict disparities from the reconstructed patches, posing additional challenges beyond standard MIM pretraining. To mitigate the difficulty of these challenging tasks and facilitate effective and robust training, we propose to use additional supervision from teacher networks. Additionally, the masking ratio is set at 40%, which is lower than the standard used in MIM pre-training. This adjustment was determined experimentally to establish an optimal upper bound, ensuring the model can effectively handle the dual challenge of reconstructing masked tokens and subsequently performing stereo matching.

Experimental validation on KITTI 2015 <cit.> and ETH3D <cit.> datasets confirms that our model exhibits improved performance compared to the existing stereo depth estimation approaches.
Our contributions can be summarized as follows:

    
  * We introduce a novel stereo depth estimation framework that leverages Masked Image Modeling (MIM) applied to stereo images to obtain a locality inductive bias, which is particularly suitable for dense prediction tasks.
    
  * We empirically demonstrate that the effective utilization of local image representation learning techniques derived from self-supervised learning significantly enhances the stereo matching process.
    
  * We present a method to employ pseudo disparity maps generated by an Exponential Moving Average (EMA) teacher as supplementary guidance, thereby improving both stability and performance in disparity estimation.




§ RELATED WORK




  §.§.§ Stereo Depth Estimation

Stereo depth estimation is a fundamental task that involves determining the disparity of objects within a scene using stereo images, simulating the depth perception capabilities of the human visual system  <cit.>. Stereo depth estimation is widely utilized in various fields, including autonomous driving <cit.>, robotics <cit.>, and 3D scene reconstruction <cit.>, where accurate depth information is crucial for navigation, object detection, and environmental interaction. Moreover, it is increasingly being used as ground truth labels in monocular depth estimation tasks <cit.>, further enhancing the training and accuracy of monocular models. With the advancement of transformer frameworks <cit.>, stereo-matching models have evolved significantly, offering enhanced capabilities in stereo depth estimation tasks. Stereo matching models can be broadly categorized into CNN-based and Transformer-based frameworks. Historically, most research has focused on CNN-based models, which construct a 3D cost volume <cit.> to calculate the disparity between corresponding pixels from the stereo images. An alternative approach concatenates features from both images to create a 4D cost volume <cit.>, followed by cost aggregation to refine the volumes. In contrast, Transformer-based models <cit.> employ a different approach. These models utilize a cross-attention layer to facilitate the transfer of information across varying views <cit.>, thereby eliminating the need for generating a correlation volume. Several studies <cit.> demonstrated that the positional encoding inherent in transformers helps establish spatial correspondence between different views. However, these Transformer-based models require massive ground truth training data to achieve competitive performance. Thus, there are fewer Transformer-based models than the latest CNN-based concurrent models <cit.>.



  §.§.§ Self-Supervised Learning

Self-Supervised Learning (SSL) is designed to extract meaningful representations from large amounts of unlabeled data <cit.>. This approach facilitates the fine-tuning of models for various downstream tasks. At the core of SSL are carefully crafted pretext tasks that leverage the intrinsic patterns and relationships within the data. These tasks often involve intentionally removing specific image features, such as color <cit.>, followed by training the model to reconstruct the missing details. Broadly, pretext tasks fall into two categories: (1) augmentation-based <cit.> and (2) reconstruction-based methods <cit.>. Augmentation-based methods generate semantically similar outcomes through various transformations, while reconstruction-based methods focus on reconstructing hidden patches or pixels to discern object structures within the image. Despite the absence of explicit semantic supervision, networks trained on these pretext tasks show their effectiveness in learning rich and meaningful image representations from the data. This capability to use large-scale unlabeled datasets has garnered significant attention, particularly in tasks where the availability of labeled data is constrained.



  §.§.§ Mask Image Modeling

Motivated by the successes in natural language processing (NLP) exemplified by BERT <cit.> and the advent of ViT <cit.>, a variety of self-supervised pre-training methodologies using Masked Image Modeling (MIM) have emerged. These approaches draw conceptual parallels to denoising autoencoders and context encoders, with the primary objective of reconstructing masked pixels, discrete tokens, or deep features. Initial efforts, such as iGPT <cit.>, which focused on pixel reconstruction, and ViT <cit.>, which aimed to predict the average color of masked patches, fell short of achieving competitive performance relative to supervised models. However, the introduction of BEiT <cit.>, which predicts visual tokens derived from a pre-trained Variational Autoencoder (VAE) <cit.>, marked a milestone in the evolution of self-supervised learning. Further notable advancements include the Masked Autoencoder (MAE) <cit.>, which employs raw pixel prediction for pre-training and emphasizes the importance of a high mask ratio (, 75%) due to spatial redundancy inherent in images. The MultiMAE <cit.> framework has been developed to extend its applicability to multi-modal or multitask scenarios. Furthermore, lightweight approaches <cit.> are being developed to mitigate the substantial training resource requirements of MIM models. Recently, MTO <cit.> optimizes masked tokens to significantly enhance pre-training efficiency, while SBAM <cit.> introduces a saliency-based adaptive masking strategy that further refines the process by dynamically adjusting masking ratios based on token salience. 





§ PROPOSED METHOD



 §.§ Motivation and Overview

Although Transformer-based models exhibit promising performance, they typically require a large scale of training data compared to CNN-based models. The Masked Image Modeling (MIM) approach offers a potential solution to this data scarcity issue by providing the locality inductive bias into the model. Motivated by the strengths of MIM, we present a novel stereo depth estimation framework, referred to as MaDis-Stereo. However, the straightforward integration of MIM complicates the dual tasks of image reconstruction and disparity prediction, which can introduce model instability. To circumvent these issues, we introduce a novel strategy that effectively integrates MIM into the model while maintaining its stability.

Fig. <ref> illustrates the overall architecture of the proposed method. The MaDis-Stereo comprises a teacher network and a student network. The student network is constructed with a ViT-Base <cit.> encoder to extract features from visible tokens in given images. The ViT-Base decoder consists of plain Transformer decoder blocks including self-attention among token features from the left image, cross-attention with token features from the right image, and a linear layer. A predicted disparity is generated by passing features gathered from different intermediate decoder blocks into the head module <cit.>. In addition, the reconstructed image is produced by a prediction layer. The teacher network shares the same architecture as the student network except a prediction layer for reconstruction, since it takes unmasked images as inputs. During training, the parameters of the teacher network are updated via exponential moving average (EMA) from the student network. The student network is trained using both the ground truth disparities and pseudo disparity maps generated by the teacher network.



 §.§ Masking-and-Reconstruction of Stereo Images
 
MaDis-Stereo's key architectural innovation lies in its capability to learn local representations through a masking-and-reconstruction methodology. For that, we first divide given left and right images I_l and I_r into N non-overlapping patches represented as I_l={I_l,1,..., I_l, N} and I_r={I_r,1,..., I_r,N}, and then apply random masking to them. A masking ratio β∈ [ 0,1  ] is provided to the model as a hyperparameter, which determines the proportion of masked areas within input data. Simply speaking, n patches are masked as follows; n = β· N.

A set of visible left patches is indicated by Ĩ_̃l̃={I_l,i|m_i=0 }, where m_i=0 signifies that the patch I_l,i is not masked, and m_i=1 means otherwise. The same random masking strategy is also applied to the right image. MaDis-Stereo proceeds with the reconstruction process using the masked view as input thereafter. Ĩ_̃l̃ and Ĩ_̃r̃ are processed individually by the ViT <cit.> encoder E_θ. Our model reconstructs the image Î_̂l̂ using a linear layer after decoding the image features with D_ϕ. Here, θ and ϕ are the weight parameters of the Transformer encoder and decoder, respectively. In the decoder, MaDis-Stereo uses cross-attention to reconstruct the image. For reconstructing the left image, the left encoded feature E_θ (Ĩ_̃l̃) is used as a query, while the right encoded feature E_θ (Ĩ_̃r̃) is employed as a key and a value for allowing information exchange between the two views. The process of reconstructing the right image is conducted similarly, except for reversing the query E_θ(Ĩ_̃r̃) and the key and value E_θ(Ĩ_̃l̃). Formally, the left image reconstruction can be written as

    Î_̂l̂= 𝙻𝚒𝚗𝚎𝚊𝚛( D_ϕ ( E_θ  ( Ĩ_̃l̃ );E_θ ( Ĩ_̃r̃ ) )).
 

Simultaneously reconstructing masked images and estimating depth is a challenging task and can destabilize the network. Therefore, we use a relatively lower value of r than the masking ratio typically used during the pre-training phase, 60% default masking ratio for SimMIM <cit.>. Experimentally, we found an appropriate masking ratio for MaDis-Stereo to be 40%. 

Additionally, Fig. <ref> demonstrates that our approach enhances the model's capability to focus more effectively on local patterns within the images. The attention distances were computed by obtaining attention weights from the cross-attention layer after processing the KITTI 2015 <cit.> dataset through the model. The comparison is made between our model, which employs a masking-and-reconstruction strategy, and a standard Transformer-based stereo matching model <cit.>. A small attention distance indicates that the model's attention heads focus more on nearby pixels, demonstrating stronger locality inductive bias.





 §.§ Network Architecture and Training Procedure


As illustrated in Fig. <ref>, the architecture comprises two networks: (1) the student network and (2) the teacher network updated via Exponential Moving Average (EMA). Both networks are built upon the ViT <cit.> encoder for feature extraction, the Transformer decoder with multi-head cross-attention layers, and the RefineNet-based feature fusion block <cit.> as a head module for disparity prediction. Additionally, the student network includes a prediction module composed of a linear layer for reconstructing masked image patches, following the methodology outlined in SimMIM <cit.>.

The overall training procedure of MaDis-Stereo is described as follows. We provide the masked stereo images obtained by applying random masking and original stereo images without masking as inputs. The two masked views are fed into the student network and are then reconstructed to impose the locality inductive bias. Note that this learning process is only conducted within the student network. In the student network, the feature of the masked image is extracted through the ViT <cit.> encoder by processing visible patches, and then the ViT <cit.> decoder uses the multi-head cross-attention layers to facilitate information exchange between the left and right features. The decoded features are subsequently passed through a linear layer for image reconstruction, while also being input into a RefineNet-based feature fusion block <cit.> to predict the output disparity. The RefineNet-based fusion block generates the disparity map by reshaping and merging four features from several transformer decoder blocks using convolutional layers.

On the other hand, the teacher network does not undergo a process of reconstructing masked regions since intact stereo images are used as inputs and only perform disparity prediction. While the student network continues its learning process, the teacher network maintains a stop-gradient state. The encoder of the teacher network is EMA-updated using the encoder of the student network, ensuring smooth training. The overall structure follows the framework of self-supervised learning methods <cit.>.
In our model, the EMA update is applied up to the teacher ViT <cit.> encoder, and the Transformer decoder and feature fusion block <cit.> are built in a Siamese manner. Formally, the teacher encoder parameter θ^T_t+1 at iteration t+1 is EMA-updated using the student encoder parameter θ_t+1 as 

    θ^T_t+1←αθ^T_t+(1-α)θ_t+1,

where α is a hyperparameter of EMA update, set to 0.9999 in our work.



 §.§ Distilling Teacher Knowledge to Student Model
 
Performing image reconstruction and depth estimation simultaneously can be a challenging task. To alleviate the challenges of the task, we utilize disparity maps generated by the teacher network as supplementary pseudo-labels to guide the student network. To be specific, the ground truth disparity maps are typically sparse due to the inherent limitation of active depth sensors <cit.> used to collect the training data.
Employing pseudo labels facilitates effective model training in scenarios of sparse ground truth disparity maps. These dense pseudo disparity maps act as supervisory signals for regions where ground truth disparity values obtained by LiDAR are unavailable.

From another perspective, using the pseudo labels of the teacher network facilitates effective knowledge transfer to the student network, ultimately enhancing the performance of the student network. Similar approaches have also been employed in weakly/unsupervised domain adaptation <cit.>, where the model is trained using labeled source data and is then applied to the target domain to generate pseudo labels using the teacher network. It was also reported that these pseudo labels help train the student model.



 §.§ Loss Function




  §.§.§ Disparity loss.

MaDis-Stereo parameterizes the network output using a Laplacian distribution <cit.>, similar to CroCo-Stereo <cit.>. Given a stereo image input, the model predicts a disparity map d and a scale parameter σ. The training objective involves minimizing the negative log-likelihood function of the ground-truth target disparity map d^gt and pseudo-disparity map d^pgt computed from the EMA-updated teacher model as follows:

    L_disp   =    L^gt_disp + L^pgt_disp
       =  1/|Ω|[ ∑_i∈ G( | d_i-d^gt_i|/σ_i - 2logσ_i) . 
          + . ∑_j∈Ω∖ G( | d_j-d^pgt_j|/σ_j - 2logσ_j) ],


where L^gt_disp and L^pgt_disp are loss functions defined using d^gt and d^pgt, respectively. 
The scale parameter σ serves as an indication of predictive uncertainty: higher values of σ result in less stringent penalties for larger errors, while lower values of σ lead to greater rewards for accurate predictions.

Ω and G indicate a set of all pixels in the image and a set of pixels where ground truth depth labels are available, respectively. 
Since the ground truth depth labels d^gt are typically sparse, we use the dense pseudo labels d^pgt generated by the teacher network as an additional guidance.
Due to the potential inaccuracies of the pseudo labels d^pgt, we exclude d^pgt for pixels where the d^gt is available. Specifically, for pixel i where ground truth exists, the training process relies solely on the d^gt_i, avoiding any guidance from the pseudo labels d^pgt_i. For the location of pixel j where ground truth d^gt_j is unavailable (, where j ∈Ω∖ G), the pseudo labels d^pgt_j can serve as a guidance. 

Similar to the semi-supervised learning paradigms <cit.>, the disparity maps generated from the EMA-teacher model are used as the pseudo depth labels. Specifically, the student network takes masked images with incomplete information as inputs and predicts disparity maps with lower accuracy than the teacher. In contrast, the teacher network predicts pseudo disparity maps from unmasked (original) images, which are more precise compared to those from the student network.



  §.§.§ Image Reconstruction loss.

The reconstruction loss is computed by predicting the pixels within masked regions of the input and evaluating L_1 loss, L_img, between the reconstructed images and the original ones.
L_img is computed for the masked patches only as

    L_img(Î,I) = 1/N∑_im_I(i)·||Î(i)-I(i)||^2

where m_I(i) indicates whether the i-th pixel is masked or not (0 if not masked, 1 otherwise). N=∑_im_I(i) denotes the total number of masked pixels from the input image.



  §.§.§ Total loss.

A final loss is calculated as the sum of the disparity loss L_disp and the image reconstruction loss L_img.
As random masking is applied to both the left and right images in the student network, L_img is computed for both left and right images.
The total loss is defined as L_total = L_disp + L_img(Î_̂l̂,I_l)+L_img(Î_̂r̂,I_r).



§ EXPERIMENTS






 §.§ Implementation Details


We fine-tuned with pre-trained weights from CroCo-Stereo <cit.>, adhering to its implementation settings. We evaluated our method on KITTI 2015 <cit.>, and ETH3D <cit.>. The final performance on the KITTI 2015 benchmark in Table <ref> and the ETH3D benchmark in Table <ref> are shown. Further details regarding the training datasets and the specific splits can be found in the supplementary material.



  §.§.§ Pre-training.

Pre-trained weights from CroCo-Stereo were fine-tuned for our experimentation purposes. CroCo-Stereo pre-trained on various stereo datasets such as CREStereo <cit.>, SceneFlow <cit.>, ETH3D <cit.>, Booster <cit.>, and Middlebury <cit.>. The pre-trained weights were trained with a batch size of 6 pairs of image crops (704×352) for 32 epochs, utilizing the AdamW optimizer <cit.> with a weight decay rate of 0.05 and a cosine learning rate schedule with a single warm-up epoch, with the learning rate set to 3.10^-5.





  §.§.§ Fine-tuning.

CroCo-Stereo was fine-tuned the pre-trained model on 1216×352 crops from non-masked KITTI 2012 <cit.> and KITTI 2015 <cit.> datasets for 20 epochs. Following the most settings of CroCo-Stereo, we used crops of size 1216×352 for fine-tuning, with a learning rate of 3.10^-5. Our MaDis-Stereo adopted the ViT-Base model as its backbone, in contrast to CroCo-Stereo, used the ViT-Large backbone for fine-tuning. Additionally, unlike the original CroCo-Stereo approach, we extended the training of MaDis-Stereo to 100 epochs to ensure effective image reconstruction. For a fair comparison, we compared our results with fine-tuning the CroCo-Stereo for also 100 epochs.















 §.§ Evaluation




We evaluated our model on KITTI 2015 <cit.> and ETH3D <cit.> two-view stereo. We experimentally confirmed that our model is most effective when conducting image reconstruction concurrently, particularly at 100 epochs. Similarly, we compared the results of CroCo-Stereo trained for up to 100 epochs with ViT-Base backbone. Table <ref> for MaDis-Stereo results on KITTI 2015 leaderboard. MaDis-Stereo achieves the best result on the main D1-all metrics with 1.57, with the best value also on D1-fg pixels with 2.31. Furthermore, Table <ref> shows the results on the ETH3D leaderboard, where our model demonstrates state-of-the-art performance. 









 §.§ Ablation Study
 
Our ablations were performed on the KITTI 2015 <cit.> dataset  for stereo matching. Here, we conducted three ablation studies in this section. Firstly, we analyzed the impact of the masking ratio r on MaDis-Stereo. Secondly, we compared the results with/without the EMA structure to assess its influence on the model's stability. Lastly, we explored the effect of utilizing pseudo disparity maps generated by the EMA teacher network to guide the student network.



  §.§.§ Masking ratio

We measure the impact of variation in the masking ratio in Table <ref>. The masking ratio is one of the crucial factors influencing the performance of the MIM network. The model learns a locality inductive bias for the missing parts in the process of predicting the masked regions of the image, which is determined by the masking ratio. However, a higher masking ratio reduces the visible image tokens available for reconstruction and makes it a challenging task. To mitigate the difficulty of the applied MIM strategy, we use a lower masking ratio than what is typically used in the MIM pre-training process. We observed that masking ratio r = 0.4 is optimized for MaDis-Stereo.



  §.§.§ EMA structure


Fig <ref> illustrates the outcomes achieved with/without the utilization of the EMA structure and each validation error results are described in Table <ref>. To ensure more stable learning, we build our model with an EMA structure. Following the standard EMA model, our MaDis-Stereo consists of the teacher and student network. The parameters of the student network are EMA to the parameters of the teacher network. The model without an EMA structure implies that solely the student network is employed for training. Our observations indicate that the MaDis-Stereo with the EMA structure exhibits much more stable learning than the model without it.



  §.§.§ Effects of using pseudo disparity map as supplement guidance

Performing image reconstruction and depth estimation simultaneously presents challenges. To tackle this issue, we utilize pseudo-disparity maps generated by an EMA teacher network as guidance to the student network during training. Table <ref> shows the impact of varying the weight of L_disp, which utilizes the disparity map as guidance.



§ CONCLUSION

We investigate the limitations of Transformer-based fine-tuning models for stereo depth estimation tasks and introduce a novel framework that integrates supervised and self-supervised learning approaches, diverging from conventional supervised methods. Our focus lies in leveraging the masking-and-reconstruction approach to enhance the inductive locality bias essential for scenarios with limited data, thereby addressing the bias deficiency in Transformers. Experimentally, we demonstrate the beneficial impact of learning local image representations during fine-tuning on a stereo depth estimation network.



