RemEdit: Efficient Diffusion Editing with Riemannian Geometry

TL;DR

RemEdit tackles the fidelity-speed trade-off in diffusion-based image editing by modeling h-space as a Riemannian manifold and solving for geodesic edits via a learnable exponential map. It introduces a dual-SLERP blending scheme for controllable edits and a goal-aware prompt enrichment using Qwen2-VL to ground edits in the input context, together with a task-aware attention pruner for acceleration. The approach yields on-manifold edits with improved semantic alignment and identity preservation, while achieving practical speedups through pruning and efficient geodesic computation. Empirical results on CelebA-HQ, LSUN-Church, and AFHQ-Dog demonstrate superior performance over prior editors and notable speedups, underscoring the method's real-world applicability in high-fidelity, real-time diffusion editing.

Abstract

Controllable image generation is fundamental to the success of modern generative AI, yet it faces a critical trade-off between semantic fidelity and inference speed. The RemEdit diffusion-based framework addresses this trade-off with two synergistic innovations. First, for editing fidelity, we navigate the latent space as a Riemannian manifold. A mamba-based module efficiently learns the manifold's structure, enabling direct and accurate geodesic path computation for smooth semantic edits. This control is further refined by a dual-SLERP blending technique and a goal-aware prompt enrichment pass from a Vision-Language Model. Second, for additional acceleration, we introduce a novel task-specific attention pruning mechanism. A lightweight pruning head learns to retain tokens essential to the edit, enabling effective optimization without the semantic degradation common in content-agnostic approaches. RemEdit surpasses prior state-of-the-art editing frameworks while maintaining real-time performance under 50% pruning. Consequently, RemEdit establishes a new benchmark for practical and powerful image editing. Source code: https://www.github.com/eashanadhikarla/RemEdit.

Figures (12)

• Figure 1: RemEdit maintains semantic fidelity under aggressive token pruning; 90% pruning is $\sim$18% faster yet remains visually acceptable, while 10% pruning is virtually indistinguishable from the unpruned edit and still $\sim$10% faster.
• Figure 2: RemEdit Architecture Diagram. Overview of our diffusion editing pipeline integrating exponential map for geodesic navigation, dual-SLERP interpolation for fidelity control, Qwen2-VL for prompt enrichment. The individual modules flow from top to bottom.
• Figure 3: Dual‐SLERP two‐stage interpolation: [Left] Inner SLERP on the Riemannian feature manifold, blending the original feature $h$ and the geodesically shifted feature $\exp_{h}(v_{0})$ via the interpolation parameter $\alpha_{\mathrm{inner}}$. [Right] Outer orthogonal SLERP on the hyperspherical noise latent space, projecting the semantic prediction $x_{0}^{\mathrm{sem}}$ onto the fidelity latent $x_{0}^{\mathrm{fid}}$, extracting the orthogonal component $o$, and interpolating between $x_{0}^{\mathrm{fid}}$ and $o$ using $\alpha_{\mathrm{outer}}$, thus disentangling attribute edits from identity preservation.
• Figure 4: GFLOPs breakdown of the DDPM U-Net at 2562 and 5122.
• Figure 5: Editing results of RemEdit on CelebA-HQ dataset where attributes {"Sad", "Smiling", "Makeup", "Young"} are Human in-distribution and {"Zombie", "Pixar"} are Human out-of-distribution.