UniTT-Stereo: Unified Training of Transformer for Enhanced Stereo Matching
Soomin Kim Hyesong Choi Jihye Ahn Dongbo Min


Low-characteristic-impedance superconducting tadpole resonators in the sub-gigahertz regime
    Mikko Möttönen
    September 9, 2024
===========================================================================================



§ ABSTRACT 
Unlike other vision tasks where Transformer-based approaches are becoming increasingly common, stereo depth estimation is still dominated by convolution-based approaches. This is mainly due to the limited availability of real-world ground truth for stereo matching, which is a limiting factor in improving the performance of Transformer-based stereo approaches.
In this paper, we propose UniTT-Stereo, a method to maximize the potential of Transformer-based stereo architectures by unifying self-supervised learning used for pre-training with stereo matching framework based on supervised learning. 
To be specific, we explore the effectiveness of reconstructing features of masked portions in an input image and at the same time predicting corresponding points in another image from the perspective of locality inductive bias, which is crucial in training models with limited training data. Moreover, to address these challenging tasks of reconstruction-and-prediction, we present a new strategy to vary a masking ratio when training the stereo model with stereo-tailored losses. 
State-of-the-art performance of UniTT-Stereo is validated on various benchmarks such as ETH3D, KITTI 2012, and KITTI 2015 datasets. 
Lastly, to investigate the advantages of the proposed approach, we provide a frequency analysis of feature maps and the analysis of locality inductive bias based on attention maps. 




















§ INTRODUCTION


Stereo matching remains fundamental for various computer vision applications, including autonomous driving, 3D reconstruction, and the recognition of objects <cit.>. The goal is to estimate a pixel-wise disparity map from two (or more) images capturing the same scene from distinct viewpoints, typically achieved by computing disparity from corresponding pixels.
The process of stereo matching is divided into two main parts: (1) feature matching and (2) disparity refinement. 
The key is to calculate the matching cost from two image pairs for feature matching and refine it accurately to obtain a reliable disparity map, considering challenges such as low-texture areas and occlusions.



While most approaches <cit.> adopt convolutional neural networks (CNNs) for extracting stereo feature and aggregating cost volume, recent studies <cit.> have attempted to utilize the Transformer architecture, which is known to have superior representation capabilities and larger receptive fields compared to traditional CNNs.

It is reported that attention mechanisms within the Transformer framework can effectively replace the traditional cost volume approaches. This enables for the dense computation of correlation between two high-resolution features without being constrained by pre-defined disparity search range unlike cost volume approaches.


Nevertheless, the performance of Transformer-based stereo approaches is at best comparable or even inferior to that of convolution-based approaches, which means that  the Transformer architecture is not yet fully utilized in the context of stereo matching. 
To maximize the advantages of the Transformer while addressing its under-utilization in the stereo matching, the characteristics of the Transformer and stereo task need to be examined thoroughly. 
Recent research on the Transformer <cit.> suggests that it demands more training data for ensuring convergence due to the lack of inductive bias compared to CNNs, which benefit from structural characteristics like local receptive fields. 
In contrast, a deep-seated challenge of the stereo task is the limited availability of real-world ground truth, primarily due to the requirements of specialized equipment such as active range sensors (e.g., LiDAR), leading to increased cost and complexity of collecting large-scale labeled training data. 
Thus, it is crucial to resolve the inductive bias deficit and effectively use stereo information from limited stereo training data when utilizing the Transformer in the stereo matching task.





In this context, we propose a novel approach, UniTT-Stereo, which stands for Unified Training of Transformer for enhanced stereo matching. Our model unifies self-supervised learning methods <cit.>, traditionally used for pre-training, with stereo matching framework based on supervised learning, enabling an effective learning tailored to Transformer based stereo matching. Our approach partially masks a left image and uses its remaining portions along with a right image to simultaneously reconstruct the original left image and estimate depth values for all pixels.

It is important to note that our experiments revealed that simply introducing the masking-and-reconstruction methodology does not necessarily improve performance. Therefore, our approach leverages a Variable Masking Ratio within the unified network, enabling the model to learn richer and more diverse information. Higher masking ratios facilitate the learning process during reconstruction, while lower masking ratios are advantageous for precise and detailed depth prediction. This balance ensures that our model can effectively capture both broad and fine-grained details, further enhancing its performance. 
We also adapt stereo-tailored losses to fully utilize the limited stereo information as mentioned earlier. We employ three synergistic loss functions at the final output level, RGB level, and feature level. Remarkably, we achieve improvements without the need for additional parameters or frameworks, relying solely on these well-designed approaches and stereo-tailored losses.



















We also provided detailed analysis on the proposed approach at various aspects. First, the locality inductive bias of our approach using the reconstruction-and-prediction task is examined and compared with existing Transformer based stereo matching methods by computing attention distances (Fig. <ref>). This will be further validated by visualizing the attention maps of cross attention module between left and right features. Additionally, the ability to capture local patterns effectively may be related to exploiting high-frequency spatial information, which offers advantages in stereo tasks by improving the accuracy of depth estimation, particularly at object boundaries and fine details. To investigate how our method effectively amplifies high-frequency information, we perform Fourier analysis on the decoder's feature map used in the disparity head, in Fig. <ref>. Detailed explanations are provided in Analysis section.


We achieve state-of-the-art results on ETH3D <cit.>, KITTI 2012 <cit.> and KITTI 2015 <cit.>
 
 datasets, demonstrating the effectiveness of our method. Our key contributions are as follows:


    
  * We examine the impact of the reconstruction-and-prediction approach on stereo depth estimation and propose a unified training approach based on these insights. 
    
  * We enhanced performance by introducing a stereo-tailored combination of loss functions from multiple perspectives: feature, RGB, and disparity.
    
  * Through extensive analyses and experiments, we validate that our model effectively leverages Transformer for stereo matching.
    







§ RELATED WORKS



 §.§ Dense prediction with Transformer

Fully convolutional networks <cit.> serve as the backbone for dense prediction, with various adaptations proposed over time. These architectures commonly depend on convolution and downsampling as fundamental components for acquiring multiscale representations, enabling the incorporation of a substantial contextual understanding. However, the low resolution in the deeper layer causes difficulty in dense prediction, so there have been many researches to maintain high resolution. 
Transformers <cit.>, based on the self-attention mechanism, demonstrate success with high-capacity architectures trained on extensive datasets. Since the Vision Transformer <cit.> adapts this mechanism to the image domain successfully but not in dense prediction, two main approaches have appeared. One is to design a specialized Transformer fitted to the dense prediction task <cit.>, and the other is to use a plain Vision Transformer and the customized decoder for dense prediction. Dense Prediction Transformer <cit.> used the latter method and achieved state-of-the-art performance in 2021. We propose an approach to fully utilize transformer-based architecture for stereo depth estimation.




 §.§ Masked image modeling

Masked image modeling (MIM) is a technique for self-supervised representation learning <cit.> using images that have masked parts. In this approach, some of the tokenized input sequence is replaced with trainable mask tokens, and the model is trained to predict the missing context based solely on the visible context. This approach, which does not require labels, is widely used for pre-training. SimMIM <cit.> and MAE <cit.> suggest that random masking with a higher mask ratio (e.g. 90%) or size can perform well for self-supervised pretraining from image data. Recently, MTO <cit.> has improved pre-training efficiency by optimizing masked tokens, while SBAM <cit.> has introduced a dynamic approach to the process with a saliency-based adaptive masking strategy that adjusts masking ratios according to the salience of the tokens. CroCo <cit.> and CroCo v2 <cit.> introduced a novel self-supervised pretraining approach exclusively designed for 3D tasks, reconstructing the masked image using the reference image. One of the advantages of Transformers is the abundance of these well-pretrained models available for use. Several studies <cit.> have investigated the effects and what the model learns from MIM as a pretraining method compared to other approaches like contrastive learning. Meanwhile, through experimentation, we have identified how MIM can impact stereo tasks and the strategies to actively leverage MIM for the specific task of stereo depth estimation.



 §.§ Stereo depth estimation


Stereo depth estimation is extensively used in fields such as autonomous driving <cit.>, robotics <cit.>, where accurate depth data is essential for navigation and object detection, and it is also increasingly employed as ground truth labels in monocular depth estimation tasks <cit.>. Stereo depth estimation requires predicting a pixel-wise dense disparity map, capturing detailed and fine information, especially for boundary regions. In traditional deep stereo matching methods, the primary steps involve four components: feature extraction, cost volume creation, feature matching, and disparity regression. To enhance either accuracy or speed, researchers have proposed several strategies to improve these four components. 3D correlation cost volume <cit.> or 4D concatenation cost volume <cit.> can be constructed to measure the similarity between two views. Several studies <cit.> have adopted iterative methods to construct disparity maps, and concurrent work <cit.> has also improved performance using this approach. Recent studies <cit.> have utilized cross-attention mechanisms to enable the exchange of information between different images instead of cost volume. We improve the performance by applying optimized approaches from an analytical perspective on the compatibility between Transformer architecture and stereo depth estimation task. 





§ PROPOSED METHOD



We introduce MIM for effective learning by utilizing pairs of a masked left image and an unmasked right image. Unlike pre-training that focuses solely on reconstruction, our goal is to improve the performance of specific downstream tasks, and thus we consider the need for a more suitable masking method. To this end, we introduce variable ratio masking through a truncated normal distribution. After both image tokens pass through the Transformer encoder, we use cross-attention modules for inter-image information exchange. Finally, the model outputs a disparity map and a reconstructed image through respective heads. Our training process involves three losses: feature consistency loss, image reconstruction loss, and disparity loss. Fig. <ref> shows the overall architecture of the proposed method. Additionally, we provide an analysis of how our approach impacts the stereo task and enhances performance using attention distance, attention map, and Fourier Transform.




 §.§ Architecture

Given left and right images I_l and I_r, each of which captures the same scene from different viewpoints, both are split into N non-overlapping patches, denoted as p_l = {p_l^1, ..., p_l^N} and p_r = {p_r^1, ..., p_r^N}. n= ⌊ rN ⌋ tokens are randomly masked only in the left image, where r∈ [0,1] is a selected masking ratio. Siamese encoders deal with visible tokens from a left image and whole tokens from a right image independently. The encoder consists of 12 blocks with a dimension of 768 for ViT-Base and 24 blocks with a dimension of 1024 for ViT-Large. 

The left image tokens from the encoder are padded with masked tokens, resulting in F_l with N tokens, which is the same number as the tokens from the right image feature F_r. The encoded left feature is then utilized by a decoder, conditional on the encoded right feature. The model constructs the query, key, and value in a self-attention block from the left token sequence in order to compute attention scores and identify relationships between tokens in the same sequence. In contrast, the model generates a cross-attention block by using the left token sequence to set up the query and the right token sequence to generate the key and value in order to find correspondences between the two images. It is composed of 12 blocks, each with a dimension of 768. 

For generating pixel-wise depth predictions, RefineNet-based fusion module is adapted to reshape and merge four features from different transformer decoder blocks. We utilize features from {2, 5, 8, 12}^th layers in the decoder. A linear head is used to get the reconstructed image output. 




 §.§ Variable Ratio Masking


Inspired from MIM pre-trained models, we leverage masked input to capture local and high-frequency patterns effectively. However, masking too much information can make it excessively difficult for the model to directly predict disparity maps, potentially hindering the model's ability to learn from raw RGB images. Conversely, when masking with a low ratio, there is no significant change in performance. To address this, we introduce a variable mask ratio, as shown in Algorithm <ref>, to ensure the model learns effectively across a range of information scales. 






















    







We use random masking with mask size 16, similar to SimMIM or MAE <cit.> with variable ratio. The masking ratio is determined from a truncated normal distribution with specified upper and lower bounds, which is generated by the given mean and standard deviation. A new masking ratio is then randomly sampled from this distribution for each batch and it is rounded to one decimal place. For example, 0.32 is rounded to 0.3, resulting in 30% masking. We confirmed the effectiveness of variable ratio masking through experiments in various settings. By default, we use a truncated normal probability distribution with a lower bound of 0.0, an upper bound of 0.5, a mean of 0.25, and a standard deviation of 1.0.






 §.§ Losses




  §.§.§ Feature Consistency Loss
 To the output feature maps from each encoder, we introduce the consistency loss which aims to enhance the alignment between two corresponding features from a stereo pair. This makes the model improve matching performance at a feature level. This is achieved by warping the feature from the right image to the feature from the left image, using the ground truth disparity information. F̃_̃l̃ means reconstructed left feature from the right feature with disparity. 

    L_consist=∑_i|F_l,i-F̃_̃l̃,̃ĩ|





  §.§.§ Disparity Loss
 The disparity loss is common and plain, but the most powerful matching loss which can be supervised by ground truth. 
We minimize negative log-likelihood with a Laplacian distribution to train the proposed model, following <cit.>: 

    L_disp=∑_i[|D _i-D̅_̅i̅|/s_i-2log s_i]

where D_i and D̅_̅i̅ are an estimated disparity and the ground truth disparity at pixel i, respectively. The scale parameter s_i is also outputted from a model. It can be understood as an uncertainty score for predictions. 



  §.§.§ Image Reconstruction Loss
 The reconstruction loss evaluates reconstruction accuracy by the Mean Squared Error (MSE) only for masked patches p_l∖p̃_̃l̃ where p_l denotes a set of patches from the first image, p̃_̃l̃ is a set of visible patches from the first image, and p̂_̂l̂ is the reconstructed first image. 
Notably, the left image undergoes reconstruction through image completion, utilizing corresponding information from the right image. Consequently, this loss can be considered as a form of matching loss from an RGB perspective.

    L_recon=1/|p_l∖p̃_̃l̃|∑_p_l,i∈ p_l∖p̃_̃l̃p̂_̂l̂,̂î-p_l,i^2













  §.§.§ Total Loss
 We supervise the model with a linear combination of three synergistic losses which are disparity loss from final output, reconstruction loss from reconstructed image, and consistency loss from feature map as L_total=λ*L_disp+L_recon+L_consist
where λ is set to 3. Since our objective is to attain improved performance on stereo depth estimation, the disparity loss takes on more weight than other losses. 



Moreover, to mitigate the potential learning of erroneous information during the early stages of training due to masked images, we calculate an uncertainty score using the reconstruction loss, allowing us to assess how effectively the model has adjusted to masked images. We used a fixed value τ =0.4 as a threshold for generalized reconstruction error ϕ (L_recon) = tanh(L_recon) to decide whether to impose the disparity loss or not. 
As estimated disparity from unsteady reconstructed feature makes the model unstable, if ϕ>τ, only L_recon + L_consist is used. 




 §.§ Analysis



Locality Inductive Bias: Fig. <ref> (a) illustrates the attention distances calculated from the attention scores obtained after passing the entire KITTI 2015 <cit.> test dataset through the cross attention module. Each point represents the average attention distance across 12 heads for each layer. Our method encourages the model to focus on these local patterns to reconstruct the masked parts using locality inductive bias via MIM. This harmonizes well with Transformers, which adept at learning global information. To further investigate the effect of the locality inductive bias in our method, we visualize the attention map in Figure. <ref> (b). It visualizes cross attention maps for an example query patch from an left image, divided into a set of 16×16 patches, in KITTI 2015 training dataset. 
Ideally, the attention score should be highest at the location of the corresponding patch in the left image, as identified by the ground truth disparity. Our approach, which tends to focus more on local information, helps prevent incorrect attention values, when compared with the plain method that has a higher risk of occasionally concentrating attention on patches located far away, as shown in the right part of the example map.




Fourier Analysis: In Fig. <ref> (a), we additionally conducted Fourier analysis on feature maps of blocks in the decoder of our model in Fig. <ref>, following <cit.>. This experiment conducted on ETH3D dataset <cit.> reveals that our approach tends to generate feature maps with a higher proportion of high-frequency components and indicates that the model is capturing detailed and sharp features, which can be advantageous for tasks like stereo depth estimation.
High frequency components of the decoder features used in the depth estimation head enable for more precise and detailed representations of object edges and fine structures in the disparity map. 
As illustrated in Fig. <ref> (b), by effectively utilizing high-frequency information, the proposed method achieved sharper boundaries on the disparity maps.




Attention Map by Varying Losses: To demonstrate the synergy when the losses are used together, Fig. <ref> visualizes the attention maps when the model is separately trained using each loss. Masking was applied only when the model was trained with the reconstruction loss or the total loss. We averaged the attention scores from every head and self-attention layer in the left encoder for self-attention visualization, and applied the same approach using cross-attention layer in the decoder to cross-attention visualization. 
In the self-attention map, the score should be highest at the location of the query patch, whereas the score should peak at the location of the corresponding patch in the cross-attention map.
As shown in Fig. <ref> (b), when trained with each loss individually, incorrect attention values appear at various locations in each case. Even the disparity loss case, which is supervised with ground truth, has its limitations. However, when these losses are combined (L_total), they can correct each other's errors and emphasize common attention patterns, working synergistically to guide the model towards more accurate attention placement. This mutual reinforcement allows the model to learn more effectively, facilitating improved stereo matching performance, as also validated in the ablation study of Table <ref>.





§ EXPERIMENTS









  §.§.§ Implementation Details

We train our UniTT-Stereo for 32 epochs using batches of 6 pairs initializing the encoder and decoder from the pre-trained weights by CroCov2 <cit.>. For optimization, we employ the AdamW optimizer <cit.> with a weight decay of 0.05. The learning rate of 3×10^-5 follows a cosine schedule with a single warm-up epoch. We utilize SceneFlow <cit.>, CREStereo <cit.>, ETH3D <cit.>, Booster <cit.>, Middlebury (2005,
2006, 2014, 2021 and v3) <cit.> with crop size of 704 × 352 to train UniTT-Stereo.
Afterward, we trained our model on KITTI 2012 <cit.> and KITTI 2015 <cit.> with crop size of 1216 × 352 for 100 epochs using effective batches of 6 pairs. We use a learning rate of 3×10^-5. For inference, we use a tiling-based strategy in which we sample overlapping tiles with the same size as the training crops, following <cit.>. 



 §.§ Stereo Depth Estimation Performance

We evaluate UniTT-Stereo on representative stereo datasets with their metrics and compare with the published state-of-the-art methods. 



  §.§.§ ETH3D

UniTT-Stereo sets a new state-of-the-art on ETH3D. Table <ref> compares UniTT-Stereo with HITNet <cit.>, RAFT-Stereo <cit.>, GMStereo <cit.>, IGEV-Stereo <cit.>, CREStereo <cit.>, and CroCo-Stereo <cit.>. 


 
















  §.§.§ KITTI 2012 & 2015


We also achieve state-of-the-art performance compared to other published methods, aside from concurrent work, on both KITTI 2012 and 2015. Table <ref> compares UniTT-Stereo with HITNet <cit.>, PCWNet <cit.>, ACVNet <cit.>, LEAStereo <cit.>, CREStereo <cit.>, IGEV-Stereo <cit.>, and CroCo-Stereo <cit.>. The qualitative comparison is shown in Fig. <ref>. 






  §.§.§ Middlebury


We also conducted performance evaluations on the Middlebury evaluation dataset. Table <ref> compares UniTT-Stereo with LeaStereo <cit.>, HITNet <cit.>, RAFT-Stereo <cit.>, CREStereo <cit.>, GMStereo <cit.>, and CroCo-Stereo <cit.>. As shown in Fig. <ref>, our method delivered high performance on data that requires precise estimation.



  §.§.§ Limitations

While our method achieved comparable results to other methods in Middlebury, there was a limitation in handling certain large maximum disparities leading to bad performance on several sequences. This is likely due to the constraints of tiling-based inference, which can restrict the ability to capture long-range correspondences. To address this, it may be necessary to use larger tile sizes or increase the overlap ratio. 



 






 §.§ Zero-shot Generalization

Generalizing from synthetic to real data is crucial due to the challenge of gathering real-world datasets. The result suggests that our approach help the model learn invariant features across different domains. Table <ref> compares UniTT-Stereo with GANet <cit.>, RAFT-Stereo <cit.>, and DSMNet <cit.>. 


 






 §.§ Ablation Study




  §.§.§ Key Components
 We conducted an ablation study on the effectiveness of each key component, including masking and three loss functions. As listed in Table <ref>, introducing each additional loss function led to performance improvement in a sequential order. Optimal performance was achieved when employing every key component.



  §.§.§ Masking Ratio
 We also evaluate the effectiveness of the variable ratio masking. As shown in Fig. <ref>, high ratio fixed masking and variable ratio masking effectively amplifies the high frequency information. But interestingly, the experimental results suggest that a high ratio may hinder learning, as performance actually deteriorated, while the use of variable ratio masking resulted in a significant improvement. The low ratio did not lead to any dramatic changes in performance. Table <ref> shows that masking with a modest level of r_max proved beneficial for performance by imparting inductive bias without compromising depth information.



 





 




§ CONCLUSION

We proposed UniTT-Stereo to maximize the strengths of Transformer-based architecture, which have traditionally lagged behind in stereo matching task. We enhance performance in a simple yet effective manner by employing reconstruction-and-prediction strategy and a combination of losses specifically designed to learn stereo information. Our approach achieves state-of-the-art performance on prominent stereo datasets and demonstrates strong zero-shot generalization capabilities. Throughout this process, we have analyzed the specific advantages our approach brings to stereo depth estimation.
















