RFM: Riemannian Flow Matching Architectures

PiperOrigin-RevId: 884062709

Author: vdeborto

hackable_diffusion/lib/architecture/__init__.py

@@ -49,6 +49,7 @@
 from hackable_diffusion.lib.architecture.mlp_blocks import MLP
 from hackable_diffusion.lib.architecture.normalization import NormalizationLayer
 from hackable_diffusion.lib.architecture.normalization import NormalizationLayerFactory
+from hackable_diffusion.lib.architecture.riemannian import RiemannianConditionalBackbone
 from hackable_diffusion.lib.architecture.sequence_embedders import RandomFourierSequenceEmbedding
 from hackable_diffusion.lib.architecture.sequence_embedders import RoPESequenceEmbedding
 from hackable_diffusion.lib.architecture.sequence_embedders import SinusoidalSequenceEmbedding

hackable_diffusion/lib/architecture/riemannian.py

@@ -0,0 +1,47 @@
+# Copyright 2026 Hackable Diffusion Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Riemannian Flow Matching architectures."""
+
+import flax.linen as nn
+from hackable_diffusion.lib import manifolds
+from hackable_diffusion.lib.architecture import arch_typing
+
+################################################################################
+# MARK: Riemannian Conditional Backbone
+################################################################################
+
+ConditionalBackbone = arch_typing.ConditionalBackbone
+
+
+class RiemannianConditionalBackbone(ConditionalBackbone):
+  """Velocity model for Riemannian Flow Matching.
+
+  Projects the output of a backbone network to the tangent space of a manifold.
+  """
+
+  backbone: ConditionalBackbone
+  manifold: manifolds.Manifold
+
+  @nn.compact
+  def __call__(self, x, conditioning_embeddings, is_training=True):
+
+    v = self.backbone(x, conditioning_embeddings, is_training=is_training)
+
+    # Project v to tangent space at xt.
+    if isinstance(v, dict) and 'velocity' in v:
+      v = v['velocity']
+
+    v_proj = self.manifold.project(x, v)
+    return v_proj

hackable_diffusion/lib/architecture/riemannian_test.py

@@ -0,0 +1,64 @@
+# Copyright 2026 Hackable Diffusion Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Riemannian Flow Matching architectures."""
+
+from absl.testing import absltest
+from hackable_diffusion.lib import manifolds
+from hackable_diffusion.lib.architecture import arch_typing
+from hackable_diffusion.lib.architecture import mlp
+from hackable_diffusion.lib.architecture import riemannian
+import jax
+import jax.numpy as jnp
+
+
+class RiemannianArchitectureTest(absltest.TestCase):
+
+  def test_riemannian_backbone_projection(self):
+    manifold = manifolds.Sphere()
+    backbone = mlp.ConditionalMLP(
+        hidden_sizes_preprocess=(16,),
+        hidden_sizes_postprocess=(16,),
+        activation='relu',
+        zero_init_output=True,
+        dropout_rate=0.0,
+        conditioning_mechanism=mlp.ConditioningMechanism.CONCATENATE,
+    )
+    model = riemannian.RiemannianConditionalBackbone(
+        backbone=backbone,
+        manifold=manifold,
+    )
+
+    key = jax.random.PRNGKey(0)
+    xt = manifold.random_uniform(key, (4, 3))
+    time_emb = jnp.array([[0.5], [0.5], [0.5], [0.5]])
+
+    # conditioning_embeddings must be a dict keyed by ConditioningMechanism.
+    conditioning_embeddings = {
+        arch_typing.ConditioningMechanism.CONCATENATE: time_emb,
+    }
+
+    variables = model.init(key, xt, conditioning_embeddings, is_training=False)
+    v = model.apply(variables, xt, conditioning_embeddings, is_training=False)
+
+    self.assertEqual(v.shape, (4, 3))
+
+    # Check that v is tangent to xt
+    inner_products = jnp.sum(xt * v, axis=-1)
+    # Project should ensure dot(xt, v) = 0 for sphere
+    self.assertAlmostEqual(jnp.max(jnp.abs(inner_products)), 0.0, places=5)
+
+
+if __name__ == '__main__':
+  absltest.main()

hackable_diffusion/lib/sampling/riemannian_sampling_test.py

@@ -38,7 +38,7 @@ def _make_sampler(manifold):
 class RiemannianFlowSamplerStepTest(absltest.TestCase):
 
   def test_update_sphere(self):
-    """Euler step on S² moves along geodesic."""
+    """Euler step on S2 moves along geodesic."""
     manifold = manifolds.Sphere()
     sampler = _make_sampler(manifold)
     key = jax.random.PRNGKey(0)
@@ -69,7 +69,6 @@ def test_update_so3(self):
     key = jax.random.PRNGKey(1)
 
     xt = jnp.eye(3)[None, ...]  # Identity rotation (1, 3, 3).
-    # Tangent vector at identity is a skew-symmetric matrix.
     v = jnp.array([[[0.0, -0.1, 0.0], [0.1, 0.0, 0.0], [0.0, 0.0, 0.0]]])
 
     current_step = base.DiffusionStep(
@@ -109,7 +108,6 @@ def test_update_torus(self):
     # dt = 1.0, so next_xt = exp(xt, v) = (xt + v) % 1.0.
     expected_xt = jnp.array([[(0.9 + 0.5) % 1.0, (0.1 - 0.5) % 1.0, 0.5]])
     np.testing.assert_allclose(next_step.xt, expected_xt, atol=1e-5)
-    # Result stays in [0, 1).
     self.assertTrue(jnp.all(next_step.xt >= 0.0))
     self.assertTrue(jnp.all(next_step.xt < 1.0))