pith. sign in

arxiv: 2603.10718 · v2 · pith:3JDTSMZInew · submitted 2026-03-11 · 💻 cs.LG

Riemannian MeanFlow for One-Step Generation on Manifolds

Zichen Zhong , Haoliang Sun , Yukun Zhao , Yongshun Gong , Yilong Yin This is my paper

Pith reviewed 2026-05-21 11:02 UTC · model grok-4.3

classification 💻 cs.LG
keywords Riemannian manifoldflow matchingone-step generationparallel transportgenerative modelmanifold samplingMeanFlow
0
0 comments X

The pith

Riemannian MeanFlow identity allows one-step sampling on manifolds without ODE integration or trajectory simulation.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper extends MeanFlow to Riemannian manifolds so that generative models can be trained and sampled directly on curved spaces. It defines an average-velocity field with parallel transport and derives a Riemannian MeanFlow identity that relates this average velocity to the instantaneous velocity at any point. The identity supplies intrinsic supervision that removes the need to simulate full trajectories during training. Implementation uses a log-map representation in tangent spaces together with a decomposed objective and conflict-aware multi-task learning to keep optimization stable. The resulting model supports conditional generation and produces competitive one-step samples on spheres, tori, SO(3), and SE(3).

Core claim

RMF defines an average-velocity field via parallel transport and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision, enabling practical one-step sampling on manifolds without trajectory simulation.

What carries the argument

The Riemannian MeanFlow identity, which uses parallel transport to connect average and instantaneous velocities in location-dependent tangent spaces and thereby provides direct supervision.

If this is right

  • Generative training on manifold data becomes free of any ODE or SDE simulation steps.
  • Sampling cost drops from many numerical integration steps to a single function evaluation.
  • The same training procedure works across standard manifolds including spheres, tori, SO(3), and SE(3).
  • Classifier-free guidance can be added directly to produce conditional samples on manifolds.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same identity construction could be tested on other curved domains such as hyperbolic space or Grassmann manifolds.
  • The conflict-aware multi-task technique might stabilize training in other geometric learning settings that combine multiple loss terms.
  • If the log-map approximation proves robust, libraries that already support Riemannian operations could adopt one-step manifold generators with minimal extra code.

Load-bearing premise

The Riemannian MeanFlow identity remains accurate and tractable when realized through log-map tangent representations, and the two-term objective plus conflict-aware multi-task learning resolves gradient interference without adding new biases.

What would settle it

Numerical checks on a simple manifold such as the sphere showing whether one-step samples drawn from the trained model match the quality and distribution of samples obtained by integrating the probability-flow ODE to high accuracy.

read the original abstract

Flow Matching enables simulation-free training of generative models on Riemannian manifolds, yet sampling typically still relies on numerically integrating a probability-flow ODE. We propose Riemannian MeanFlow (RMF), extending MeanFlow to manifold-valued generation where velocities lie in location-dependent tangent spaces. RMF defines an average-velocity field via parallel transport and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision. We make this identity practical in a log-map tangent representation, avoiding trajectory simulation and heavy geometric computations. For stable optimization, we decompose the RMF objective into two terms and apply conflict-aware multi-task learning to mitigate gradient interference. RMF also supports conditional generation via classifier-free guidance. Experiments on spheres, tori, SO(3), and SE(3) demonstrate competitive one-step sampling with improved quality-efficiency trade-offs and substantially reduced sampling cost.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The manuscript proposes Riemannian MeanFlow (RMF) as an extension of MeanFlow to Riemannian manifolds for one-step generative modeling. It defines an average-velocity field via parallel transport, derives a Riemannian MeanFlow identity linking average and instantaneous velocities to enable intrinsic supervision without trajectory simulation, implements the identity in log-map tangent representation, decomposes the objective into two terms, and applies conflict-aware multi-task learning to address gradient interference. The method supports classifier-free guidance for conditional generation and is evaluated on spheres, tori, SO(3), and SE(3), claiming competitive quality with substantially reduced sampling cost compared to ODE integration in flow matching.

Significance. If the derived identity holds exactly under the log-map implementation and the optimization stabilizes without introducing new biases, RMF would represent a meaningful advance in simulation-free training and one-step sampling for manifold-valued generative models. The approach addresses a practical bottleneck in Riemannian flow matching by avoiding numerical integration at inference time, and the experimental results on standard manifolds like SO(3) and SE(3) suggest tangible efficiency gains. The conflict-aware multi-task learning component is a pragmatic addition for handling the decomposed objective.

major comments (2)
  1. [Derivation of Riemannian MeanFlow identity (likely §3)] The central claim that the Riemannian MeanFlow identity enables accurate intrinsic supervision in log-map tangent representation is load-bearing for the one-step generation result. On positively curved manifolds, parallel transport of velocities from multiple points does not necessarily commute with averaging; the manuscript should explicitly derive or bound any higher-order curvature terms that arise when the identity is instantiated via log-maps (rather than stating it is made 'practical' without further qualification).
  2. [Optimization and objective (likely §4)] The assumption that the two-term objective decomposition plus conflict-aware multi-task learning fully resolves gradient interference without new biases is not sufficiently validated. The manuscript should report ablation results isolating the effect of the conflict-aware weights on the one-step sampling metrics (e.g., MMD or negative log-likelihood) across the tested manifolds.
minor comments (2)
  1. [Preliminaries] Notation for the log-map, exponential map, and parallel transport operators should be defined explicitly with a table or dedicated subsection early in the paper to improve readability for readers less familiar with Riemannian geometry.
  2. [Experiments] Figure captions for the qualitative samples on SO(3) and SE(3) should include the exact number of function evaluations used for the baseline ODE integrator to allow direct comparison of sampling cost.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive and detailed review. The comments help clarify the presentation of the Riemannian MeanFlow identity and the optimization procedure. We respond to each major comment below and indicate the revisions we will make to the manuscript.

read point-by-point responses

  1. Referee: [Derivation of Riemannian MeanFlow identity (likely §3)] The central claim that the Riemannian MeanFlow identity enables accurate intrinsic supervision in log-map tangent representation is load-bearing for the one-step generation result. On positively curved manifolds, parallel transport of velocities from multiple points does not necessarily commute with averaging; the manuscript should explicitly derive or bound any higher-order curvature terms that arise when the identity is instantiated via log-maps (rather than stating it is made 'practical' without further qualification).

    Authors: We appreciate the referee's emphasis on the role of curvature in the derivation. The Riemannian MeanFlow identity follows directly from integrating the velocity field and applying the parallel transport operator to define the average velocity in a common tangent space. When the identity is realized via the log-map, the parallel transport is performed to the tangent space at the base point, and the resulting expression is exact for the chosen reference. On positively curved manifolds the parallel transport operator admits a curvature-dependent expansion; we will revise §3 to include this expansion explicitly, derive the leading higher-order terms involving the Riemann curvature tensor, and provide a bound on the remainder that depends on sectional curvature and the integration interval. We will also report the magnitude of these terms on the sphere (constant positive curvature) and SO(3) for the time discretizations used in our experiments, thereby qualifying the accuracy of the log-map implementation. revision: yes

  2. Referee: [Optimization and objective (likely §4)] The assumption that the two-term objective decomposition plus conflict-aware multi-task learning fully resolves gradient interference without new biases is not sufficiently validated. The manuscript should report ablation results isolating the effect of the conflict-aware weights on the one-step sampling metrics (e.g., MMD or negative log-likelihood) across the tested manifolds.

    Authors: We agree that isolating the contribution of the conflict-aware weighting is necessary for a complete validation. The objective decomposition separates supervision of the average-velocity and instantaneous-velocity fields, while the conflict-aware weights are computed from the cosine similarity of the task gradients to reduce interference. In the revised manuscript we will add a dedicated ablation section that compares (i) the full conflict-aware scheme, (ii) uniform weighting of the two terms, and (iii) a single-term baseline. For each variant we will report MMD and negative log-likelihood on the sphere, torus, SO(3), and SE(3) datasets, together with training stability metrics. These results will quantify the improvement attributable to the conflict-aware component and confirm the absence of introduced bias in the one-step sampling performance. revision: yes

Circularity Check

1 steps flagged

Riemannian MeanFlow identity links velocities by construction of the parallel-transport average field

specific steps
  1. self definitional [Abstract]
    "RMF defines an average-velocity field via parallel transport and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision."

    The identity is derived after defining the average-velocity field via parallel transport; the link between average and instantaneous velocities is therefore tautological to the construction of the field itself rather than an independent result.

full rationale

The abstract states that RMF defines an average-velocity field via parallel transport and then derives the identity linking average and instantaneous velocities for supervision. This creates moderate circularity because the supervision signal is built directly from the same parallel-transport operation used to construct the field, without independent external validation shown in the provided text. No other patterns like self-citation chains or fitted predictions are evident from the abstract. The log-map implementation is presented as making the identity practical, but the core link reduces to the definitional step.

Axiom & Free-Parameter Ledger

1 free parameters · 1 axioms · 0 invented entities

The central claim rests on the validity of the Riemannian MeanFlow identity and the effectiveness of the conflict-aware decomposition; these are introduced without independent external benchmarks in the provided abstract.

free parameters (1)
  • conflict-aware multi-task weights
    Weights used to balance the two terms of the decomposed RMF objective during optimization.
axioms (1)
  • domain assumption The Riemannian MeanFlow identity holds when velocities are represented via parallel transport in tangent spaces.
    Invoked to link average and instantaneous velocities for supervision without trajectory simulation.

pith-pipeline@v0.9.0 · 5679 in / 1313 out tokens · 45984 ms · 2026-05-21T11:02:03.809171+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.