Source-linked AI summary

TransNormal-2: Geometry-Grounded Rectified Flow with Edge-Aware Decoding for Precise Normal Estimation

Mingwei Li, Yi Yang, Hehe Fan

arXiv:2609.06665v1cs.CV

TL;DR

VAE reconstruction degradation limits pixel-space precision in latent diffusion geometry estimation, especially at object boundaries. TransNormal-2 addresses this with geometry-aware decoded-space supervision and constrained RGB-guided refinement, matching or exceeding MoGe-2 across general-scene metrics while using 1.4% as many annotations and improving transparent-object MAE.

  • Problem

    VAE reconstruction introduces boundary-localized normal-estimation error that latent-space training objectives alone cannot remove.

  • Method

    TransNormal-2 combines geometry-aware pixel-space losses with a lightweight RGB-guided Geometric Refinement Module for constrained post-decoding correction.

  • Results

    TransNormal-2 matches or exceeds MoGe-2 on all eight general-scene metrics using only 1.4% as many task-specific normal annotations, with MAE reductions of 4.2° on ClearGrasp and 3.1° on ClearPose.

  • Takeaways & Limitations

    The approach improves normal estimation most clearly for transparent objects, where refractive appearance and boundary precision make geometry difficult.

  • Takeaways & Limitations

    Model complexity, inference speed, runtime breakdowns, and inference-optimization results are reported only in the Supplementary Material.

Abstract

from arXiv · show

Diffusion-based models enable monocular geometry estimation, yet their pixel-space precision is limited by a shared, under-studied error source: VAE reconstruction degradation. The 8x spatial compression in the VAE encoder-decoder degrades surface normals at object boundaries; even encoding and decoding ground-truth normals introduces 1.3--8.5° of mean angular error (MAE), with edge MAE reaching 2.8x the global MAE. We present TransNormal-2, a FLUX.2-based rectified-flow framework with single-step deterministic inference that addresses this degradation on both sides of the VAE decoder: in how latent predictions are supervised during training, and in how decoded normals are corrected at inference. First, geometry-aware pixel-space losses, including inverse rendering self-consistency, von~Mises-Fisher angular loss, and wavelet edge-aware regularization, complement latent MSE by enforcing spherical normal geometry and diffuse image-formation cues after VAE decoding. Second, a lightweight Geometric Refinement Module (GRM) applies an RGB-guided residual correction to reduce boundary-localized decoding errors without freely rewriting the coarse prediction. On general-scene benchmarks, TransNormal-2 matches or exceeds MoGe-2 on all eight reported metrics while using only 1.4% as many task-specific normal annotations. The gains are clearest for transparent objects, reducing MAE by 4.2° on ClearGrasp and 3.1° on ClearPose over the strongest prior baselines. Code will be released at https://longxiang-ai.github.io/TransNormal-2.

1 INTRODUCTION

TransNormal-2 identifies VAE reconstruction degradation as a shared source of boundary-localized normal-estimation error and addresses it with geometry-aware supervision plus RGB-guided refinement. It matches or exceeds MoGe-2 across general-scene metrics while using far fewer annotations and improves transparent-object results.

  • Motivation: VAE reconstruction degradation is an under-studied error source caused by 8× spatial compression that averages abrupt boundary normal changes.The degradation is inherited by VAE-based latent-diffusion geometry methods.
  • Motivation: Edge-region error reaches up to 2.8× the global average and cannot be removed by improving latent prediction alone.The encode–decode path loses full-resolution boundary evidence.
  • Method: The Geometric Refinement Module applies a lightweight RGB-guided residual correction that targets remaining boundary errors without freely rewriting coarse predictions.Inference remains a single deterministic forward pass.
  • Results: TransNormal-2 matches or exceeds MoGe-2 on all eight general-scene metrics while using only 1.4% as many task-specific normal annotations.The reported gains are especially pronounced for transparent-object scenes.
  • Results: Mean angular error is reduced by 4.2° on ClearGrasp and 3.1° on ClearPose over the strongest prior baselines.These improvements target transparent objects, where refractive appearance and boundary precision make estimation difficult.
  • Method: Geometry-aware pixel-space losses combine inverse-rendering self-consistency, von Mises–Fisher angular loss, and wavelet edge-aware regularization with latent MSE.These objectives impose spherical-normal, image-formation, and boundary-frequency constraints after decoding.

2 RELATED WORK

Related work spans classical, discriminative, and generative geometry estimation, with diffusion models increasingly using latent representations and rectified-flow Transformers. TransNormal-2 is positioned around quantifying and correcting the boundary bias introduced by VAE compression, including for transparent objects.

  • Geometry estimation: Surface-normal estimation evolved from physics-based methods through discriminative deep models toward generative diffusion-based approaches.The generative line includes zero-shot geometry, joint depth-normal prediction, reduced diffusion variance, and temporally consistent video normals.
  • Diffusion geometry: Dense geometric prediction with diffusion models progressed from latent diffusion and end-to-end fine-tuning to deterministic single-step prediction and flow matching.Discriminative foundation models remain faster and more metric but lack some generative priors useful for boundary preservation.
  • Architectures: Diffusion backbones evolved from U-Nets to rectified-flow Diffusion Transformers used for geometry and joint appearance–geometry modeling.The related systems include image-editing and video-generative adaptations.
  • VAE bottleneck: VAE-based geometry methods compress spatial resolution by 8×, which can average abrupt boundary changes and produce reconstruction artifacts.Pixel-space fine-tuning and VAE bypasses are alternative responses to this limitation.
  • VAE bottleneck: Prior VAE research studies reconstruction artifacts and compression trade-offs, while post-decoding refinement offers an orthogonal correction path.The paper distinguishes its contribution by quantifying surface-normal VAE degradation as motivation for refinement.
  • Transparent objects: Transparent-object perception is difficult because refraction and reflection corrupt conventional geometric cues, motivating dedicated depth and normal benchmarks.DKT uses video inputs and is therefore not directly comparable to this single-image setting.

3 PRELIMINARIES AND BOTTLENECK ANALYSIS

TransNormal-2 predicts normal latents from RGB latents through a frozen-VAE, single-step FLUX.2 rectified-flow pipeline, then decodes and refines them. Controlled analysis shows that VAE compression introduces substantial angular error concentrated at geometric boundaries.

  • Setup: The frozen VAE encodes RGB and ground-truth normals into latents, while a deterministic FLUX.2 DiT predicts the normal latent in one forward pass.Prediction uses a fixed timestep and empty-prompt conditioning.
  • Pipeline: The LCM-adjusted predicted latent is decoded into a coarse normal map and then refined using RGB-guided geometric correction.The framework includes latent-space continuity processing before frozen-VAE decoding.
  • Compression failure mode: 8× spatial compression preserves smooth color gradients better than geometric discontinuities, where within-block normal changes are blurred or averaged.A one-dimensional boundary model represents the ideal normal jump as a blurred transition of width σ.
  • Boundary model: Boundary error grows with the normal jump and blur width, while away from boundaries the same blur is harmless.The analysis links Euclidean discrepancy to the reported angular error for unit normals.
  • Framework overview: The framework combines direct normal-latent prediction, latent continuity processing, VAE decoding, RGB-edge refinement, and four complementary training losses.The losses are latent MSE, wavelet edge-aware regularization, von Mises–Fisher angular loss, and inverse-rendering self-consistency.
  • Diagnostic: Ground-truth encode–decode reconstruction introduces 1.3°–8.5° MAE across four benchmarks, with Edge/Global MAE reaching 2.82× on ClearGrasp and 1.72× on iBims.The same edge-concentrated pattern appears in other latent-diffusion VAEs.

4 METHOD

TransNormal-2 addresses VAE-induced normal degradation through geometry-aware pixel-space supervision during training and RGB-guided residual correction after decoding. Its two-phase design combines decoded-normal constraints with a gated GRM that targets boundary errors while preserving the coarse prediction.

  • Overview: TransNormal-2 divides VAE degradation control between pixel-space training supervision and post-decode correction.The first improves what the latent can encode, while the second reduces errors that remain after decoding.
  • Core prediction: The core predictor decodes LCM-adjusted normal latents into coarse normal maps before refinement.FLUX.2 produces raw normal tokens, the LCM repairs latent-level seams, and the frozen VAE decoder produces the coarse prediction.
  • Geometric refinement: The GRM uses RGB-guided anchors, a gated residual, and L2 normalization to correct boundary errors without freely rewriting the normal field.Opaque images use RGB-guided filtering, while transparent images use the coarse normal because RGB edges are unreliable on refractive surfaces.
  • Phase 1 objective: The first phase combines latent MSE with wavelet, von Mises-Fisher, and inverse-rendering losses on decoded normals.Wavelet supervision targets ground-truth normal boundaries; vMF measures directional error on S2; inverse rendering uses Lambertian shading and detached lighting estimation.
  • Phase 2 objective: The second phase freezes the core predictor and trains the GRM in base and calibration substages.The calibration substage supervises the confidence gate on pixels whose anchor error exceeds a threshold while penalizing mean activation.
  • Reported outcome: TransNormal-2 achieves the best result on every reported general-scene metric using 122K labeled samples.The reported comparison covers eight metrics, while gains are largest on ClearGrasp and ClearPose and TN-Syn is near-saturated.

5 EXPERIMENTS

The experiments evaluate TransNormal-2 with a frozen FLUX.2 VAE, LoRA adaptation, two-phase training, and single-pass inference across general-scene and transparent-object settings. Qualitative results emphasize fine geometric structure and robustness to refracted background textures.

  • Setup: TransNormal-2 fine-tunes FLUX.2[klein] 9B with LoRA while keeping the VAE encoder and decoder frozen.The model uses empty-prompt conditioning and a rectified-flow DiT backbone.
  • Setup: Training uses a staged Phase 1 schedule followed by Phase 2 GRM training and single-forward-pass inference.Phase 1 progresses from general scenes to transparent-object data and then enables inverse rendering; Phase 2 trains the residual and confidence gate.
  • Baselines: The baseline comparison covers discriminative feed-forward models and diffusion-based dense-prediction methods.The listed baselines include Omnidata, DSINE, MoGe-2, GeoWizard, StableNormal, Marigold, Lotus, Lotus-2, and E2E-FT.
  • General scenes: Qualitative general-scene examples test preservation of object-scale shape alongside fine structures such as gear teeth and watch mechanisms.The examples use unlabeled in-the-wild objects, so they show predicted normal maps without error rows.
  • Transparent objects: On transparent-object benchmarks, TransNormal-2 produces lower-error geometry under refractive conditions than existing methods.The figure compares predicted normal maps and masked error maps, where refracted background textures commonly distort competing predictions.

5.3 Quantitative Results

TransNormal-2 matches or exceeds strong baselines across general and transparent-object evaluations, with improvements over its predecessor and gains from the staged training and refinement recipe.

  • General-scene results: TransNormal-2 matches or exceeds MoGe-2 on all eight general-scene metrics while using about 73× fewer labeled samples.It attains Avg. Rank 1.4 versus MoGe-2’s 2.3.
  • Comparison with TransNormal: TransNormal-2 improves every head-to-head benchmark over TransNormal, with the largest gain on zero-shot ClearPose at −6.4°.The redesign replaces the earlier U-Net pipeline with FLUX.2 LoRA adaptation and adds geometry-aware losses and GRM refinement.
  • Ablation results: Adding wavelet and vMF losses improves mean angular error by 0.3° on both NYUv2 and ScanNet over the MSE-only baseline.Enabling inverse-rendering self-consistency yields a further 0.6° gain on each benchmark.
  • Ablation results: The decoupled GRM achieves the best ablation results on NYUv2 at 14.7° and 62.5%, and ScanNet at 12.7° and 69.2%.Jointly training the GRM with the DiT underperforms the decoupled setup.

5.5 GRM Design Ablation

The GRM remains stable across capacity choices, while RGB guidance is especially valuable for general-scene refinement and less necessary on synthetic or transparent benchmarks.

  • Capacity ablation: RGB-guided GRM variants differ by only 0.03° in average mean angular error across tested depths and widths.Increasing capacity to K=8 or hidden width 128 provides no systematic gain, while K=2 nearly matches the default.
  • Qualitative ablation: The qualitative ablation attributes progressively better local geometry to wavelet and vMF supervision, inverse-rendering consistency, and full GRM refinement.The full model reduces residual boundary artifacts in in-the-wild objects.
  • Input-design ablation: Removing RGB guidance increases mean error by about 1.0° on NYUv2 and 0.7° on ScanNet.The ablation removes the edge encoder, RGB channels, and RGB-guided anchor.
  • Input-design ablation: Without RGB guidance, refinement remains competitive on synthetic and transparent benchmarks, matching or slightly improving CG-Syn, TN-Syn, and ClearPose.This supports RGB guidance as a cross-domain image-edge prior for general scenes.

5.6 LoRA Rank Ablation

LoRA rank 256 provides the strongest capacity-efficiency trade-off: larger rank does not improve the core predictor, and GRM comparisons use the same coarse predictions where specified.

  • LoRA rank: Rank 256 achieves the best average error at 15.4°, while ranks 64, 128, and 512 converge near 16.5°.Rank 256 matches or leads the best rounded value on every benchmark.
  • LoRA rank: Doubling trainable parameters from 700M to 1.4B at rank 512 provides no improvement.The result suggests rank 256 already captures the FLUX.2 backbone’s useful representational capacity.
  • Experimental setup: The rank ablation evaluates the core predictor with GRM refinement disabled, isolating DiT LoRA capacity.The accompanying table identifies rank 256 as best on average with no further gains at rank 512.

5.7 GRM vs. Classical Image-Guided Filters

The learned GRM outperforms classical image-guided filters across the evaluated benchmarks, with its largest improvements concentrated in edge regions.

  • Overall comparison: The GRM gives the lowest mean angular error on all four benchmarks and the best average, 10.6° versus 11.0° for the best classical filter.No classical filter improves all four benchmarks.
  • Frequency analysis: High-frequency angular-error energy decreases by 53.2% on NYUv2 and 42.9% on ScanNet after GRM refinement.The reported ratios compare angular-error power spectral density before and after refinement.
  • Edge analysis: After GRM refinement, edge pixels improve by about 2.6°, compared with 0.6°–0.8° for non-edge pixels.Before refinement, edge MAE is 1.9–2.1× non-edge MAE.
  • Edge analysis: The absolute edge-region gain is 3.3–4.4× the non-edge gain, matching the VAE’s spatial error profile.The diagnostic separates regions using Canny edges on RGB, dilated by 3×3.

5.9 Frequency-Domain Error Analysis

Frequency-domain analysis finds that GRM refinement reduces angular-error energy across spatial frequencies, especially for localized boundary errors; runtime details are deferred to supplementary material.

  • The analysis computes radially averaged PSDs on mean-subtracted angular-error maps, comparing predictions with and without GRM.The PSD is computed on angular-error maps rather than predicted normal maps.
  • The post-refinement-to-pre-refinement PSD ratio remains below 1 across the spectrum on both general-scene benchmarks.
  • GRM primarily suppresses spatially localized boundary errors rather than uniformly shifting all pixels.
  • Model complexity, inference speed, runtime breakdowns, and optimization results are reported in the Supplementary Material.

6 CONCLUSION

TransNormal-2 identifies VAE reconstruction as a boundary-focused source of normal-estimation error and addresses it with geometry-aware supervision plus RGB-guided refinement. It matches or exceeds MoGe-2 broadly, with larger gains on transparent-object benchmarks, while retaining stated scope limitations.

  • VAE reconstruction introduces up to 8.5° angular error even on ground-truth normals, with edge MAE reaching up to 2.8× the global MAE.
  • TransNormal-2 combines vMF, wavelet, and inverse-rendering losses with an RGB-guided GRM that reduces boundary-localized errors.
  • TransNormal-2 matches or exceeds MoGe-2 on all eight general-scene metrics.
  • MAE decreases by 4.2° on ClearGrasp and 3.1° on zero-shot ClearPose over the strongest prior baselines.
  • The inverse-rendering loss assumes diffuse Lambertian reflectance, while GRM cannot compensate for large-region errors from the core predictor.
Loading 2609.06665v1…