RBE-Flow: Recurrent Bayesian Estimation on Feature Manifolds for Cross-Modal Registration¶
Conference: ECCV 2026
arXiv: 2606.30492
Code: https://github.com/NEU-Liuxuecong/RBE-Flow
Area: Image Registration / Multi-Modal Perception
Keywords: Cross-modal registration, Bayesian estimation, optical flow estimation, feature manifolds, uncertainty modeling
TL;DR¶
Reformulates cross-modal dense optical flow estimation as a closed-loop recurrent Bayesian state estimation problem on feature manifolds. Relying on Recurrent Manifold Optimization (RMO) to generate flow observations with uncertainty, and Uncertainty-Adaptive Probabilistic Update (UAPU) to perform MMSE-optimal Bayesian fusion via Sigma-point projections while feeding back the posterior covariance to regulate the next optimization step. It consistently achieves SOTA on three cross-modal datasets (OSdataset, WHU-OPT-SAR, and RoadScene), reducing the AEPE by up to 45.4% compared to previous best methods.
Background & Motivation¶
Cross-modal image registration (e.g., optical-SAR, visible-infrared) is a fundamental capability in multi-sensor perception, supporting downstream applications such as remote sensing analysis, 3D scene understanding, and autonomous navigation. However, the inherent non-linear radiation variations (e.g., intensity inversion, speckle noise, spectral gaps) and geometric deformations (viewpoint changes, scale differences, local deformations) of heterogeneous sensors severely violate standard photometric alignment assumptions, making the matching landscape highly non-convex and full of local minima. Existing deterministic matching methods lack uncertainty awareness and tend to accumulate erroneous updates in ambiguous regions, leading to registration failures—uncertainty is typically used only as a passive output for post-hoc analysis rather than an active component of the inference process.
Key Challenge: Cross-modal matching is inherently a highly uncertain inference problem, but existing methods treat it as a deterministic feed-forward regression. Goal: Reframe optical flow estimation as a dynamic state estimation problem, letting the system adjust its optimization behavior at each step based on its predictive confidence. Key Insight: Perform local non-linear optimization on the feature manifold to generate likelihood evidence, followed by error self-correction via recurrent Bayesian updates. Core Idea: Tightly couple RMO (solving damped non-linear least squares on feature manifolds to yield flow observations and uncertainty) and UAPU (fusing prior and likelihood using Sigma-point projections to output posterior covariance) as a closed loop. The posterior covariance is fed back to regulate the damping intensity of the next RMO, enabling the system to automatically transition to conservative gradient descent in high-uncertainty regions while maintaining efficient Gauss-Newton steps in low-uncertainty regions.
Method¶
Overall Architecture¶
RBE-Flow formulates cross-modal optical flow estimation as a closed-loop recurrent system of "prior initialization - likelihood observation generation - Bayesian posterior update - covariance feedback". The input is a pair of cross-modal images (e.g., optical-SAR or visible-infrared), and the output is a dense pixel-level optical flow field. The entire pipeline is divided into an initialization phase and a recurrent iteration phase: First, a shared-weight CNN extracts features at three scales, 1/2, 1/4, and 1/8 (MFE), followed by merging multi-scale features via self-attention and cross-attention (LFI). Simultaneously, a 4D correlation volume is computed at the 1/8 scale to obtain the initial optical flow field (GFI), establishing the prior mean for the recurrent Bayesian estimation. Once entering the iteration loop, the RMO module solves local non-linear least squares on the feature manifold using Spectral-Entropy Guided Correlation Fusion (SGCF) and Covariance-Adaptive Damped Optimization (CDO) to generate flow observations and their uncertainties. The UAPU module fuses these observations with the prior via deterministic Sigma-point projections, calculating the MMSE-optimal Bayesian gain along with the updated posterior mean and covariance. The posterior covariance is fed back to the RMO to adjust the damping coefficient for the next iteration—transitioning to conservative gradient descent when uncertainty is high, and efficient Gauss-Newton steps when uncertainty is low.
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["Input Image Pair<br/>Optical-SAR / Visible-IR"] --> B["MFE: Multi-scale Feature Extraction<br/>Shared-weight CNN → 1/2, 1/4, 1/8 Features"]
B --> C["LFI: Local Feature Initialization<br/>SA + CA Multi-scale Fusion"]
B --> D["GFI: Global Flow Initialization<br/>1/8 Feature 4D Correlation → Initial Flow"]
C --> E["RMO: Recurrent Manifold Optimization<br/>SGCF Correlation Fusion + CDO Covariance-Adaptive Damping"]
D --> E
E --> F["UAPU: Uncertainty-Adaptive Probabilistic Update<br/>Sigma-point Projection → MMSE Optimal Gain → Posterior Covariance"]
F -->|"Posterior Covariance Feedback Refers to Damping for Next RMO"| E
F --> G["Output: Dense Optical Flow Field"]
Key Designs¶
1. RMO: Spectral-Entropy Guided Correlation Fusion and Covariance-Adaptive Damped Optimization
RMO is the "observation generator" in Bayesian estimation. It assumes that high-dimensional image features do not scatter randomly in the ambient space but lie on a low-dimensional non-linear manifold; thus, optical flow estimation is essentially finding the shortest path connecting corresponding features on this manifold. RMO must simultaneously produce flow observations and their associated observation noise covariance.
RMO consists of two key mechanisms. The first is Spectral-Entropy Guided Correlation Fusion (SGCF): To capture the multi-scale geometric structure of the manifold, RMO constructs two layers of correlation volumes—the local window correlation \(C_{L0}\) captures high-frequency details, while the pooled feature correlation \(C_{L1}\) provides low-frequency global consistency. For each correlation volume, a probability distribution \(\mathbf{p}\) is obtained via softmax, and the spectral entropy is calculated as \(\mathcal{H}(\mathbf{u}) = -\sum_i p_i \log p_i\). Then, the two layers of correlation volumes and their respective entropy values are concatenated and fed into a lightweight network \(\Phi_{\text{Spec}}\), which outputs adaptive fusion weights \([\omega_{L0}, \omega_{L1}]\) via softmax. The final fused correlation is obtained by \(C_{\text{fused}} = \omega_{L0} \cdot C_{L0} + \omega_{L1} \cdot C_{L1}\). This mechanism allows the network to dynamically adjust its reliance on different scales of information based on the matching ambiguity in the current region—relying more on local details in texture-rich regions and more on global context in low-texture regions.
The second is Covariance-Adaptive Damped Optimization (CDO): RMO uses a multi-head ConvGRU to directly predict the Jacobian \(\boldsymbol{J}_k\), residual gradient \(\mathbf{g}_k\), and damping metric \(\lambda_k\) required for subsequent Levenberg-Marquardt (LM) computations. The key innovation is introducing the posterior covariance \(\boldsymbol{P}_{k-1|k-1}\) from the previous step (output by UAPU) into the damping adjustment:
When the prior uncertainty is high (large \(\operatorname{tr}(\boldsymbol{P}_{k-1|k-1})\)), \(\lambda'_k\) increases, and the LM update naturally transitions from a Gauss-Newton step to a more conservative gradient descent step. The flow increment is obtained by solving \((\boldsymbol{J}_k^{\top}\boldsymbol{J}_k + \lambda'_k\boldsymbol{I})\Delta\boldsymbol{Flow}_k = -\boldsymbol{J}_k^{\top}\mathbf{g}_k\), yielding the final flow observation \(\boldsymbol{Z}_k = \hat{\boldsymbol{Flow}}_{k-1} + \Delta\boldsymbol{Flow}_k\). The observation noise covariance \(\boldsymbol{R}_k\) is estimated by a lightweight ConvNet, jointly constituting the parameters of the likelihood.
2. UAPU: Uncertainty-Adaptive Bayesian Posterior Update
RMO's observations are accurate on the local cost surface but may suffer from errors due to occlusions or motion boundaries from a global perspective. The core idea of UAPU is to model flow refinement as a non-linear belief propagation problem: at each step, the posterior distribution from the previous step is taken as the prior belief, which is then fused with the likelihood evidence provided by RMO using Bayes' rule to compute the current posterior distribution, thus achieving error self-correction.
The prior propagation is modeled as a random walk: \(\hat{\boldsymbol{Flow}}_{k|k-1} = \hat{\boldsymbol{Flow}}_{k-1|k-1}\), and QNet \(\Psi_Q\) adaptively predicts the spatially heterogeneous process noise \(\boldsymbol{Q}_k\), yielding the prior covariance \(\boldsymbol{P}_{k|k-1} = \boldsymbol{P}_{k-1|k-1} + \boldsymbol{Q}_k\). Cholesky decomposition \(\boldsymbol{S}\) is introduced to ensure numerical stability and positive semi-definiteness.
Due to the strong non-linearity of the feature manifold, performing local linearization on the mapping from the state space to the observation space introduces non-negligible truncation errors. UAPU employs a strategy similar to the Unscented Transform: generating \(2L+1=5\) deterministic Sigma sample points around the prior mean \(\hat{\boldsymbol{Flow}}_{k|k-1}\) (1 center point + 4 symmetric perturbation points along the principal axes of the manifold), projecting them to the observation space as \(\mathbf{Z}_k^{(i)}\) using a learnable observation network \(\text{ObsNet}(\cdot)\) (a lightweight 1x1 convolution MLP), and performing a weighted average to obtain the predicted observation mean \(\hat{\boldsymbol{Z}}_k = \sum_{i=0}^{4} W_m^{(i)}\mathbf{Z}_k^{(i)}\) with weight \(W_m^{(0)} = \frac{\lambda_{\text{scale}}}{L+\lambda_{\text{scale}}}\).
During the fusion stage, the innovation covariance is computed as \(\boldsymbol{P}_{zz} = \sum_{i=0}^{4} W_c^{(i)}(\boldsymbol{Z}_k^{(i)} - \hat{\boldsymbol{Z}}_k)(\boldsymbol{Z}_k^{(i)} - \hat{\boldsymbol{Z}}_k)^{\top} + \boldsymbol{R}_k\) and the cross-covariance as \(\boldsymbol{P}_{xz} = \sum_{i=0}^{4} W_c^{(i)}(\boldsymbol{X}_k^{(i)} - \hat{\boldsymbol{Flow}}_{k|k-1})(\boldsymbol{Z}_k^{(i)} - \hat{\boldsymbol{Z}}_k)^{\top}\). Under the Gaussian approximation, minimizing the mean squared error of the posterior estimate is equivalent to solving for the optimal Bayesian fusion gain \(\mathcal{K}_k = \boldsymbol{P}_{xz}\boldsymbol{P}_{zz}^{-1}\). The posterior is updated as \(\hat{\boldsymbol{Flow}}_{k|k} = \hat{\boldsymbol{Flow}}_{k|k-1} + \mathcal{K}_k(\boldsymbol{Z}_k - \hat{\boldsymbol{Z}}_k)\), and the posterior covariance as \(\boldsymbol{P}_{k|k} = \boldsymbol{P}_{k|k-1} - \mathcal{K}_k\boldsymbol{P}_{zz}\mathcal{K}_k^{\top}\). This update restricts information gain by shrinking the prior search space, automatically downweighting observational evidence in unreliable regions.
Loss & Training¶
A multi-stage hybrid supervision strategy is adopted to bridge deterministic global alignment and stochastic local refinement.
Initialization Loss: For the initial optical flow \(\boldsymbol{Flow}_{\text{init}}\) generated by GFI, a simple L1 supervision is used. After downsampling the ground-truth optical flow to the same resolution as \(\boldsymbol{Flow}_{\text{init}}\), the L1 distance is computed: \(\mathcal{L}_{\text{init}} = \frac{1}{N}\sum_{\mathbf{p}}|\boldsymbol{Flow}_{\text{init}} - \boldsymbol{Flow}_{\text{gt}}|\).
Geometry-Aware Rectified NLL Loss: Standard NLL objectives often lead to variance collapse—the model minimizes penalties by predicting variance values that approach zero, causing gradient explosions. To address this, a geometry-aware rectification mechanism is introduced to impose a dynamic lower bound on the predicted variance that scales with the geometric error:
where \(\mathbf{S}\) is the Cholesky factor output by UAPU. The NLL loss for the \(k\)-th iteration is \(\mathcal{L}_{\text{NLL}}^{(k)} = \frac{1}{2}\sum_{j\in\{x,y\}}\left(\log(\hat{\sigma}_k^{(j)})^2 + \frac{(\hat{\boldsymbol{Flow}}_k^{(j)} - \boldsymbol{Flow}_{\text{gt}}^{(j)})^2}{(\hat{\sigma}_k^{(j)})^2}\right)\). An exponentially decaying deep supervision with \(\gamma=0.9\) is applied across all \(N_{\text{iter}}\) iterations, resulting in the total loss \(\mathcal{L}_{\text{total}} = \lambda_{\text{init}}\mathcal{L}_{\text{init}} + \lambda_{\text{NLL}}\mathcal{L}_{\text{NLL}}\). This design forces high-error regions to exhibit high uncertainty, preventing "overconfident yet incorrect" pathological configurations.
Training is initialized with the publicly available pre-trained weights of XoFTR, and end-to-end fine-tuning is performed on RoadScene, WHU-OPT-SAR, and OSdataset. The learning rate is set to \(3\times10^{-4}\) with a batch size of 8. Inputs are cropped to \(64\times64\) patches, and data augmentations include random scaling [0.9, 1.1], rotation [-30, 30], and translations of up to 15 pixels.
Key Experimental Results¶
Main Results¶
Evaluated against 9 representative methods spanning handcrafted features, sparse matching, semi-dense matching, and dense matching across three datasets: OSdataset (optical-SAR, forests/farmlands/rivers), WHU-OPT-SAR (complex urban optical-SAR), and RoadScene (RGB-IR street views).
| Method | OSdataset AEPE | OSdataset CMR@1px | WHU-OPT-SAR AEPE | WHU-OPT-SAR CMR@1px | RoadScene AEPE | RoadScene CMR@1px |
|---|---|---|---|---|---|---|
| HOWP (ISPRS'23) | 20.39 | 0.3 | 157.14 | 0.0 | 16.67 | 0.2 |
| LNIFT (TGRS'22) | 46.98 | 0.3 | 167.00 | 0.0 | 28.89 | 0.1 |
| MSG (TGRS'25) | 44.50 | 0.3 | 120.97 | 0.1 | 43.59 | 0.0 |
| RIFT2 (TIP'20) | 29.83 | 0.7 | 214.08 | 0.0 | 29.45 | 0.0 |
| GMFlow (TPAMI'23) | 5.13 | 2.1 | 3.57 | 2.7 | 4.42 | 0.5 |
| XoFTR (CVPR'24) | 5.41 | 0.0 | 18.05 | 0.0 | 6.00 | 0.0 |
| RAFT (ECCV'20) | 2.89 | 14.0 | 2.07 | 25.4 | 1.79 | 31.3 |
| ADRNet (TGRS'24) | 1.41 | 50.5 | 1.43 | 52.9 | 0.98 | 71.3 |
| GDROS (TGRS'25) | 1.51 | 48.8 | 2.87 | 28.7 | 0.68 | 79.0 |
| RBE-Flow | 0.77 | 78.5 | 0.53 | 91.0 | 0.49 | 90.8 |
RBE-Flow achieves the best AEPE on all three datasets. On OSdataset, the AEPE is reduced from ADRNet's 1.41px to 0.77px (a relative reduction of 45.4%); on WHU-OPT-SAR, the AEPE is only 0.53px, which is a nearly threefold improvement compared to the second-best ADRNet's 1.43px; on RoadScene, the AEPE sits at 0.49px, and the [email protected] reaches 37.7%, substantially outperforming other methods. Crucially, whereas existing methods suffer from steep performance degradation under strict thresholds (\(t<1\text{px}\)), RBE-Flow retains robust matching rates, highlighting its unique advantage in sub-pixel high-precision alignment.
Ablation Study¶
Starting from the baseline (CNN multi-scale encoding + LFI + L1 loss), each component is incrementally added.
| Configuration | OSdataset AEPE | OSdataset CMR@1px | WHU-OPT-SAR AEPE | WHU-OPT-SAR CMR@1px | RoadScene AEPE | RoadScene CMR@1px |
|---|---|---|---|---|---|---|
| (1) A+L1 (Baseline) | 1.32 | 22.9 | 5.52 | 2.5 | 0.69 | 87.4 |
| (2) A+L1+GFI | 0.98 | 61.8 | 1.10 | 47.3 | 0.64 | 87.3 |
| (3) A+GFI+loss (Rectified NLL) | 0.91 | 66.7 | 0.76 | 82.9 | 0.59 | 89.0 |
| (4) A+GFI+loss+RMO | 0.84 | 71.9 | 0.61 | 89.3 | 0.54 | 89.8 |
| (5) A+GFI+loss+RMO+UAPU (Full) | 0.77 | 78.5 | 0.53 | 91.0 | 0.49 | 90.8 |
Component-wise analysis: (2) GFI provides a reliable initial flow field by computing dense feature correlations, reducing the AEPE on WHU-OPT-SAR from 5.52 to 1.10. This indicates that global correlation is crucial for handling large displacements. (3) Replacing the standard L1 loss with the geometry-aware rectified NLL improves the CMR@1px on WHU-OPT-SAR from 47.3% to 82.9%, validating that probabilistic supervision outperforms deterministic regression. (4) RMO introduces local non-linear optimization on the manifold along with covariance-adaptive damping, elevating the [email protected] on RoadScene from 12.6% to 26.9%. (5) Finally, incorporating UAPU to close the loop with uncertainty feedback increases the CMR on OSdataset and WHU-OPT-SAR from 3.1% to 11.6% and from 11.9% to 24.0%, respectively, under the safest 0.3px threshold.
Key Findings¶
- The largest contribution comes from the UAPU closed-loop: The synergistic effect of RMO+UAPU is particularly prominent at strict sub-pixel thresholds; removing UAPU causes the [email protected] to plummet by over 60%.
- Rectified NLL is critical for stability: The variance of standard NLL training often collapses to zero. The lower bound mechanism in geometry-aware rectification serves as a key enabler for training stability.
- Excellent efficiency: With a parameter count of 14.3M, inference on a single RTX 4090 takes an average of 1.48ms per pair, with each iteration requiring only 0.13ms. The complete framework (RMO+UAPU) incurs only a 12.9% time overhead per step (+29.42ms) compared to the deterministic RMO-only baseline, and the UAPU memory overhead is extremely low at +0.002GB, keeping the overall peak GPU memory below 10.16GB.
- Strong generalization: Traditional methods and general optical flow estimators (e.g., RAFT, GMFlow) fail almost entirely on datasets with dramatic modal disparities like WHU-OPT-SAR, whereas RBE-Flow maintains a robust 91.0% CMR@1px.
Highlights & Insights¶
- Uncertainty transitions from a passive output to an active control signal: Most existing methods treat uncertainty as a post-hoc analysis tool. This work is the first in dense cross-modal optical flow to feed back the posterior covariance into the optimizer's damping parameter, establishing a true "confidence-driven optimization" loop. This design philosophy can extend to any iterative refinement task—as long as the refinement process contains adjustable step size parameters, the Bayesian posterior variance can be utilized for adaptive tuning.
- Sigma-point projections bridge non-linearity and analytical optimality: Instead of using linearization, UAPU employs 5 deterministic sampling points to handle the strong non-linearity of the feature manifold. This avoids the truncation errors of the Extended Kalman Filter (EKF) while preserving closed-form solutions for the MMSE-optimal gain. This represents a clever compromise between accuracy and differentiability, offering broad reference value for probabilistic state estimation tasks requiring end-to-end training.
- Geometry-aware rectified NLL prevents variance collapse: By explicitly incorporating a geometric-error-driven lower bound \((\hat{\boldsymbol{Flow}} - \boldsymbol{Flow}_{\text{gt}})^2\) into the variance prediction, the model is kept from falsely predicting zero variance in high-error regions. This is more sensible than adding a fixed lower bound \(\epsilon\) because the lower bound scales adaptively with the error, avoiding collapse without penalizing fine-grained calibration in converged regions.
Limitations & Future Work¶
- Expansion from 2D to 3D is a major direction: The authors explicitly mention planning to extend the probabilistic manifold optimization strategy from 2D dense matching to 3D multi-modal scene reconstruction. The current framework relies on 2D ground-truth optical flow for supervision; obtaining reliable supervision signals in 3D space is an open challenge.
- Limited modalities: The experiments only cover two types of modal pairs: optical-SAR and visible-infrared. The generalization ability to more extreme modal combinations (such as optical-thermal infrared, CT-MRI medical images) has not yet been validated. The geometric error lower bound of the Rectified NLL relies on ground-truth optical flow, and how to generalize this to real-world scenarios where ground truth is unavailable (such as unlabeled medical images) requires further exploration.
- Fixed number of closed-loop iterations: Currently, \(N_{\text{iter}}\) is a predefined hyperparameter (set to 7 in the paper) rather than a data-adaptive early stopping mechanism. In theory, the posterior covariance can be used to assess convergence—stopping the iteration when \(\operatorname{tr}(\boldsymbol{P}_{k|k})\) ceases to decrease—which is a natural direction for improvement.
- Lightweight feature encoder might be a bottleneck: MFE uses only a simple shared-weight ResNet CNN rather than a transformer or a stronger backbone. Under extreme cross-modal scenarios, the quality of feature manifolds encoded by CNNs directly determines the upper limit of RMO optimization. Exploring stronger feature extractors (such as cross-attention pre-aligned encoding) could further elevate performance.
Related Work & Insights¶
- vs DROID-SLAM: DROID-SLAM embeds a differentiable LM solver into the BA optimization of SLAM, but it is tailored for mono-modal scenes and lacks uncertainty feedback. RBE-Flow introduces a differentiable second-order solver to dense cross-modal optical flow for the first time, introducing a closed-loop Bayesian update—where posterior covariance feedback regulates damping, representing a probabilistic inference dimension absent in DROID-SLAM.
- vs PDC-Net / Probabilistic Registration: PDC-Net learns to predict confidence maps to identify unreliable regions, but the uncertainty is a passive output of inference. RBE-Flow uses uncertainty as an active driving signal of the inference process—representing a general paradigm upgrade for any iterative refinement task.
- vs VBReg (CVPR 2023): VBReg utilizes variational Bayes for outlier rejection in 3D point clouds. RBE-Flow extends Bayesian filtering principles from 3D outlier rejection to recurrent state estimation on 2D feature manifolds, tackling the intense non-convexity of image matching. While they share the philosophy of Bayesian updating, their application scenarios and specific mechanisms (variational vs. Sigma-point projection) are entirely different.
- vs RAFT: RAFT uses ConvGRU for recurrent correlation lookup and flow updates, where the updates are learned implicitly without an explicit optimization process. RBE-Flow's RMO explicitly solves the LM equations to obtain flow increments, giving the updates clear geometric meaning and an interpretable damping mechanism, with the Bayesian fusion in UAPU providing probabilistic calibration.
Rating¶
- Novelty: ⭐⭐⭐⭐⭐ Tightly coupling closed-loop recurrent Bayesian estimation with feature manifold optimization for cross-modal dense optical flow represents a paradigm-level innovation in the field of registration, reshaping uncertainty from a passive output to an active control signal.
- Experimental Thoroughness: ⭐⭐⭐⭐⭐ Covers three cross-modal datasets (optical-SAR and RGB-IR) and compares against 9 representative methods spanning four paradigms: handcrafted, sparse, semi-dense, and dense. Ablation experiments incrementally validate each component's contribution, and a dedicated efficiency analysis (time/memory overhead) is provided.
- Writing Quality: ⭐⭐⭐⭐ The methodology is clearly derived, offering a complete chain from MFE to GFI, RMO, UAPU, and Loss, backed by rigorous mathematical definitions. The design motivation and mechanism of SGCF are slightly condensed, which may require additional context for readers unfamiliar with spectral entropy.
- Value: ⭐⭐⭐⭐⭐ The closed-loop uncertainty feedback paradigm is not only applicable to cross-modal registration but also holds potential for transfer to any vision tasks requiring iterative refinement that are prone to errors in ambiguous regions (e.g., stereo matching, multi-frame fusion, video object tracking). The variance lower bound design of Rectified NLL is also a practical technique to prevent variance collapse during probabilistic model training.