RemEdit: Efficient Diffusion Editing with Riemannian Geometry

AI summary

Overview

• Research area: Diffusion-based controllable image editing, with an emphasis on Riemannian geometry of latent spaces and inference-time acceleration.
• Technical level: Advanced. The paper assumes familiarity with diffusion models (DDIM, U-Net h-space, inversion), differential geometry (manifolds, Christoffel symbols, exponential maps, geodesic ODEs), ODE solvers, and attention token pruning.
• Scope: RemEdit is a diffusion editing framework that replaces Euclidean latent offsets with learned geodesic paths in h-space and adds a task-aware attention pruner, evaluated on CelebA-HQ, LSUN-Church, and AFHQ-Dog at 256×256.

What This Paper Is About

Diffusion models produce high-quality images, but editing them precisely is slow and often breaks the identity of unedited regions. Existing methods either treat the latent h-space as flat (so edits fall off the data manifold and cause artifacts) or accelerate inference with content-agnostic token pruning that does not know which regions the edit actually needs. RemEdit's goal is to deliver both high semantic fidelity and fast inference by modeling h-space as a Riemannian manifold and pruning attention tokens in an edit-aware way.

Key Contributions

1. Geodesic navigation of h-space. A lightweight Mamba-based module learns the manifold's Christoffel symbols, defining a learnable exponential map. The geodesic ODE is solved from t=0 to t=1 with a high-precision adaptive ODE solver (Dopri5) to produce the edit offset Δh = exp_h(v₀) − h, avoiding post-hoc Jacobian estimation while keeping O(N) complexity.
2. Dual-SLERP blending for fidelity control. An inner SLERP on the Riemannian feature manifold blends the original feature h with the geodesically shifted feature exp_h(v₀) via α_inner to control edit strength. An outer orthogonal SLERP in noise space projects the semantic prediction x₀^sem onto the fidelity latent x₀^fid, extracts the orthogonal component o, and interpolates with α_outer to preserve identity.
3. Goal-aware prompt enrichment. A single pass of a frozen Qwen2-VL model generates a detailed, instance-specific caption for the source image, forming a broader CLIP + Qwen2 joint embedding space that constrains underspecified prompts (e.g., "face" → "face with makeup").
4. Task-specific attention pruning. A lightweight PruningHead, conditioned on both the token sequence T and the edit vector d_edit, predicts per-token importance scores; top-k tokens with k = ⌊N·(1−ρ)⌋ are kept for attention. Training uses L = L_fidelity + λ_sparsity·L_sparsity, balancing output matching against savings.

Main Findings

• Editing quality on CelebA-HQ (Table 1): Across Smiling, Sad, and Tanned, RemEdit reports S_dir of 0.1982 / 0.1792 / 0.1948 and segmentation consistency of 92.41% / 89.72% / 92.18%, versus the strongest listed baselines — Asyrp at 0.190 / 0.159 / 0.177 and 87.90% / 88.90% / 89.31%, and Diffusion-CLIP at 0.170 / 0.163 / 0.174.
• Inference speed on CelebA-HQ: RemEdit reports 2.8 seconds, against 145.0 s for Prompt-to-Prompt, 45.2 s for Diffusion-CLIP, 38.7 s for BoundaryDiffusion, 28.9 s for Asyrp, 20.1 s for LEdits++, 12.3 s for StyleGAN-NADA, and 8.5 s for StyleCLIP.
• Editing quality on LSUN-Church (Table 2): RemEdit leads on all three tasks — Department Store 0.1959 S_dir / 59.03% seg. consistency, Ancient 0.2193 / 63.91%, Red Brick 0.2200 / 68.20% — ahead of Asyrp (0.1932 / 57.62%, 0.2087 / 62.65%, 0.2170 / 67.42%).
• Reconstruction fidelity (Table 3): Deeper inversion and more generative steps improve reconstruction. The best reported cell (t₀=300, S_for=500, S_gen=500) reaches 0.019 MAE / 0.065 LPIPS / 0.923 SSIM; the weakest reported cell (t₀=600, S_for=6, S_gen=6) falls to 0.084 / 0.283 / 0.501.
• Ablation of components (Table 4): The Euclidean h-space baseline scores 0.190 S_dir / 87.9% seg. consistency at 1.82 s. Adding geodesic navigation raises this to 0.196 / 89.5% (2.76 s). Replacing SLERP with linear interpolation drops it to 0.188 / 86.7%. Full unpruned RemEdit reaches 0.198 / 92.4% but costs 2.89 s.
• Pruning recovers the cost: At ρ = 0.2, RemEdit reaches 0.192 S_dir / 90.1% at 2.38 s; at ρ = 0.5 it reaches 0.184 / 89.5% at 2.31 s — still competitive with the 0.190 baseline while being faster than the unpruned model.
• Aggressive pruning still works: Figure 1 reports that 10% pruning is virtually indistinguishable from the unpruned edit while still being ~10% faster, and that 90% pruning is ~18% faster while remaining visually acceptable.
• Task-aware token retention: The importance heatmaps (Fig. 8) show the pruner concentrating on semantically relevant regions — the mouth for "Smiling," the entire face for "Sleep" — and assigning lower weight to background.
• Attention is the bottleneck: Self-attention layers consume over 80% of total GFLOPs in a single forward pass, exceeding convolution and ResNet blocks by over 4× at 256×256 and over 8× at 512×512.
• Few-shot training: RemEdit trains on only 500 images with a two-shot setup (n=2), versus Asyrp's 1000+ images. One training iteration suffices for the Riemannian block to capture high-level semantics; roughly three iterations refine details such as eyebrows and hair texture.
• Step efficiency on hard cases (Fig. 10): RemEdit reaches the edit in ~40 steps (~3.1 s), against Asyrp at ~1000 steps (~65.4 s) and NT-P2P at ~200 steps (~110.9 s).
• Off-manifold diagnosis: The paper attributes artifacts in prior h-space methods to adding a vector to a latent code, which pushes it off the manifold of natural images.
• Zero-shot comparison: Figure 7 shows qualitative results on the RealEdit benchmark under identical prompts, inversion depth, and seeds; quantitative numbers for RealEdit are not reported in the provided content.
• AFHQ-Dog: Listed as an evaluation benchmark at 256×256, but no corresponding results appear in the provided content.

Methodology in Plain English

The framework starts from a real image and inverts it into the diffusion model's latent space. Instead of nudging the U-Net bottleneck features by a straight-line offset, it treats that space as a curved surface. The researchers build an initial point by combining the h-space feature with a temporal embedding, predict an initial tangent velocity using a linear layer plus a norm-preserving tanh retraction, and then let a small Mamba network estimate the local curvature (the Christoffel symbols). A numerical ODE solver integrates the geodesic equation from t=0 to t=1, and the endpoint minus the starting point becomes the edit vector. This keeps the edit on the manifold.

Two interpolation stages then tune the result. The first blends between the original and geodesically shifted features inside the network, using a scalar to set how strong the edit should be. The second happens at the end of a denoising step, where two candidate clean-image predictions exist — one faithful to the original, one semantically edited. The semantic one is projected onto the fidelity one, the orthogonal remainder is extracted, and a spherical interpolation between the fidelity prediction and that remainder produces the final output, controlling how much identity is preserved.

Because plain prompts are ambiguous, a frozen Qwen2-VL model captions the source image in one pass, enriching the text embedding so the edit direction is grounded in the specific instance. Finally, a learned PruningHead scores every attention token using both the features and the edit direction, keeps only the top-k fraction, computes attention on that reduced set, and scatters the result back to the full token grid. It is trained with two competing terms: one that matches the unpruned model's output and one that pushes toward more pruning. All experiments use a DDIM architecture at 256×256 unless noted, and report wall-clock end-to-end inference including VLM captioning, inversion, geodesic ODE solving, and decoding.

Why This Matters

Impact on research. The paper reframes h-space stability — previously explained by linearity — as emerging from an underlying Riemannian structure of diffusion dynamics. It offers an alternative to post-hoc Jacobian or score-metric estimation by learning the connection directly with an O(N) Mamba module. It also argues that content-agnostic token pruning is misaligned with editing, because it cannot distinguish globally redundant tokens (a patch of sky) from locally critical ones (an unedited object), and introduces task conditioning of the pruner as the fix.

Real-world applications.

• Photo retouching and portrait editing, where attribute changes such as "Smiling," "Sad," "Makeup," and "Tanned" must not alter hair or background.
• Creative and stylized transformations such as "Zombie" and "Pixar" looks while retaining subject identity.
• Architectural and scene image editing, as demonstrated on LSUN-Church store, ancient, and red-brick styles.
• Interactive or real-time editing tools, where the reported 2.8 s CelebA-HQ inference and sub-50% pruning budgets make per-edit latency practical.

Industry relevance. Inference cost is a direct driver of serving economics for generative editing products. Attention layers dominating GFLOPs by more than 4× at 256×256 and more than 8× at 512×512 makes them the natural target, and RemEdit's reported speed-ups come without a large fidelity penalty in the reported metrics. The few-shot training setup — 500 images, two-shot, one to three iterations — lowers the barrier for adapting the framework to new domains. Source code is released at github.com/eashanadhikarla/RemEdit under a CC BY-NC-SA 4.0 license.

Future Directions

• Reducing dependence on pretrained models. The authors note RemEdit relies on pretrained Qwen2-VL and CLIP, and propose exploring learned, domain-specific guidance for out-of-distribution tasks.
• Extending evaluation coverage. AFHQ-Dog is listed as an evaluation benchmark and RealEdit is used qualitatively, but quantitative results for both are not reported in the provided content; filling these gaps would test generalization.
• Pushing pruning further. Figure 1 reports ~18% speed-up at 90% pruning and ~10% at 10% pruning, while the abstract cites real-time performance under 50% pruning; characterizing the quality/speed frontier beyond these points is an open question.
• Broadening the geometry. The Riemannian block captures high-level semantics within one to three training iterations; whether the learned connection transfers across resolutions (e.g., the 512×512 regime discussed for GFLOPs) or across model architectures is not established.

Target Audience

Researchers and graduate students working on diffusion-based image editing, generative model controllability, and inference acceleration; practitioners building interactive or production image-editing systems who need to balance fidelity against latency; and readers with a background in differential geometry or manifold learning who want to see Riemannian tools applied to latent-space manipulation in diffusion models.

Authors’ 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.