Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

HybridMambaDepth

Uncertainty-Aware Mamba–Transformer Fusion for Metric Depth Estimation

Python PyTorch Framework Status

A research prototype that combines learned metric-depth priors, optical focus cues, uncertainty, Mamba state-space modeling, and self-attention in one dense prediction system.

Selected NYU Depth V2 predictions

At a glance

  • Core contribution: uncertainty-aware fusion of two domain-specialized metric-depth priors, depth from focus, and DFF uncertainty, refined by a Mamba–Transformer U-Net.
  • Best available NYU result: 0.3282 m RMSE, 0.0911 AbsRel, and 0.9301 δ1 across the archived 654-image evaluation artifact.
  • Training evidence: completed 60k–75k optimizer steps with finite recorded losses and finite final model/Adam states.
  • Engineering scope: 60.37M trainable parameters inside a 274.84M-parameter hybrid system, trained with PyTorch Lightning and CUDA selective scan.
  • Portfolio scope: a data-free, license-conscious research repository with transparent evaluation limitations and no SOTA claim.

Why this project

Single-image depth priors capture strong scene semantics, while depth from focus (DFF) provides a complementary optical cue and an explicit uncertainty estimate. My idea was to make the model hybrid at two levels:

  1. Hybrid cues: fuse indoor and outdoor Depth Anything V2 metric priors with focal-stack DFF depth and uncertainty.
  2. Hybrid architecture: refine those cues with a U-shaped network combining Mamba state-space blocks and self-attention.

My architecture hypothesis is that dense metric depth needs both global scene context and spatially precise reconstruction: Mamba-style state-space mixing propagates long-range context, while self-attention retains content-dependent nonlocal interactions. This is a design hypothesis, not a measured efficiency claim; the archive does not contain a controlled Mamba-only versus attention-only benchmark.

The result is not a repackaged monocular backbone. It is a learned, uncertainty-aware fusion pipeline designed to turn heterogeneous depth evidence into one positive metric-depth prediction.

Best available results

NYU Depth V2 — 654-image evaluation artifact

RMSE ↓ AbsRel ↓ SqRel ↓ Log RMSE ↓ δ1 ↑ δ2 ↑ δ3 ↑
0.3282 m 0.0911 0.0463 0.1140 0.9301 0.9896 0.9974

These values are the recomputed mean of all 654 per-image rows in the best available evaluation artifact. The experiment uses the project's synthetic focal-stack protocol, where focal observations are generated from NYU RGB-D samples. It is therefore an internal hybrid-depth evaluation, not a blind monocular benchmark or a SOTA claim. The artifact predates strict checkpoint hashing, so it is reported separately from checkpoint validation loss.

Optimization evidence

Model/run Steps Best val/total_loss What the record supports
Hybrid fusion, surviving checkpoint 60k 1.1091 Best currently available fusion weights; epoch 188 / step 57,267
Hybrid fusion, historical callback record 50k 1.1025 Lower historical objective, but its best-weight file is missing
Hybrid fusion continuation 75k 1.3322 Only 0.14% better than its 60k parent run
Pure RGB HMT variant 75k 2.8566 Stable comparison run; not a controlled ablation because optimizer settings differ

val/total_loss is the training objective—SILog + 0.5 × multi-scale gradient loss + masked L1—and must not be read as RMSE.

The loss terms serve different roles: SILog emphasizes relative depth structure, multi-scale gradient loss penalizes inconsistent spatial transitions, and masked L1 anchors the prediction to absolute metric scale.

Validation-loss milestones

Every inspected train/validation loss scalar was finite, and the formal runs reached 60k–75k optimizer steps. The curves also expose a real limitation: the fusion run reached a validation plateau while training loss kept falling, indicating overfitting rather than unlimited benefit from longer training.

Training stability: from failed runs to completed training

The early report called the failure “gradient vanishing,” but the source audit separates two questions: was the log-depth objective numerically valid, and did the deep dense refiner provide short enough gradient paths? I addressed them at different levels:

  • Gradient-flow design: added dense multi-scale reconstruction, pre-normalized residual Mamba/attention blocks, and encoder–decoder skip paths.
  • Numerical stability: replaced a temporary hard clamp with differentiable Softplus + ε positive-depth output and evaluated SILog in full precision.
  • Update stability: froze pretrained priors, used larger effective batches through accumulation, and enabled 1.0-norm gradient clipping.

The strongest evidence is operational: 60k/75k checkpoints contain finite final model and Adam states, and the resumed 15k-step continuation changed every FusionCNN parameter element and nearly every HMT parameter element, while all frozen priors remained unchanged.

That still does not prove that gradient vanishing was eliminated. The historical runs did not log per-layer gradient norms, and parameter changes can also include Adam momentum and weight decay. I therefore describe the result as empirical training stabilization. See the training-stability audit and rerun protocol.

System design

flowchart LR
    RGB[Single RGB frame] --> DIN[Depth Anything V2<br/>Indoor prior · frozen]
    RGB --> DOUT[Depth Anything V2<br/>Outdoor prior · frozen]
    FS[10-frame focal stack<br/>+ focus distances] --> DFF[DFFNet · frozen]
    DFF --> DDFF[DFF metric depth]
    DFF --> UNC[Uncertainty map]
    DIN --> CAT[Concatenate 3 depth cues]
    DOUT --> CAT
    DDFF --> CAT
    CAT --> FUSION[FusionCNN<br/>3 → 32 → 32 → 2]
    FUSION --> RCAT[Append uncertainty<br/>2 + 1 channels]
    UNC --> RCAT
    RCAT --> HMT[HMT U-Net refiner<br/>Mamba + self-attention]
    HMT --> POS[Softplus + ε]
    POS --> DEPTH[Positive metric depth]
Loading

Model composition

Component Role Training status Parameters
Depth Anything V2 Indoor Base Indoor-domain metric prior Frozen 97.47M
Depth Anything V2 Outdoor Base Outdoor-domain metric prior Frozen 97.47M
DFFNet Focal-stack depth and uncertainty Frozen 19.52M
FusionCNN Learned cross-prior fusion Trainable 10.7K
HMT refiner Mamba–attention dense refinement Trainable 60.36M
Total 60.37M trainable / 274.84M total 274.84M

What I implemented

  • Fusion architecture: designed the dual-domain depth-prior + DFF + uncertainty fusion path.
  • Dense prediction: adapted a Mamba–Transformer U-Net from segmentation to one-channel continuous metric-depth regression.
  • Gradient flow: reworked dense refinement with residual and multi-scale skip paths.
  • Numerical stability: added a learned fusion stem and positive-depth Softplus + ε parameterization.
  • Objective design: combined scale-invariant log loss, multi-scale gradient consistency, and masked L1 loss.
  • Diagnostics: added module-level gradient health instrumentation for controlled future reproductions.
  • Experiment execution: ran and audited 60k–75k-step PyTorch Lightning experiments on one and two RTX 3090/4090-class GPUs.
  • Experiment forensics: designed the hybrid-fusion vs. pure-RGB HMT sequence and recovered actual executed hyperparameters from checkpoints.
  • Systems integration: integrated CUDA selective-scan dependencies, synthetic focal-stack generation, checkpoint resume, and NYU metric evaluation.
  • Engineering postmortem: diagnosed a resume-time optimizer-state issue and a TensorBoard image-logging path that generated more than 400 GB of events.

The project started as a MambaVision-to-dense-prediction adaptation with multi-scale feature taps and a DPT-style decoder, then evolved into the current confidence-aware hybrid fusion model. See Design evolution, the full Experiment log, and the Training stability audit.

Repository layout

HybridMambaDepth/
├── src/hybrid_mamba_depth/   # Fusion, losses, metrics, and gradient diagnostics
├── configs/                  # Reconstructed experiment configurations
├── results/                  # Lightweight, auditable result summaries
├── scripts/                  # Figure-generation utilities
├── tests/                    # Shape, gradient, loss, and metric tests
└── docs/                     # Experiments, stability audit, design, attribution

Quick start

This repository intentionally injects third-party priors/refiners as PyTorch modules instead of redistributing their source code.

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"
pytest

The clean reference model expects four components: an indoor prior, an outdoor prior, a DFF module returning (depth, uncertainty), and a dense refiner.

from hybrid_mamba_depth import HybridMambaDepth

model = HybridMambaDepth(
    indoor_prior=indoor_depth_model,
    outdoor_prior=outdoor_depth_model,
    dff_prior=dff_model,
    refiner=hmt_refiner,
)

metric_depth = model(rgb, focal_stack, focus_dist)

See the model implementation for the precise tensor contract.

Scope and reproducibility

This is a data-free research portfolio repository. It contains no NYU images, raw TensorBoard events, pretrained third-party code, or model checkpoints.

  • The 654-image NYU split was used for validation during development and for the available final metric artifact, so it is not an untouched test set.
  • The available metric CSV was not written with a checkpoint hash; a clean release should re-run evaluation against the surviving 1.1091 fusion checkpoint.
  • Reported runtime from the original script is omitted because it did not synchronize CUDA and excluded parts of preprocessing.
  • The project is a prototype and does not claim state of the art.

Foundations and attribution

This work builds on ideas and interfaces from HybridDepth, Depth Anything V2, MambaVision, HMT-UNet, and DPT. Their architectures and code are not claimed as my inventions.

The clean integration code in this repository is original portfolio code. Third-party source is deliberately not vendored; see Attribution and licensing before attempting a full reproduction.

About

Uncertainty-aware Mamba-Transformer fusion for metric depth estimation

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages