make rwfm better

src/wassersteinflowmatching/riemannian_wasserstein/DefaultConfig.py

@@ -28,20 +28,23 @@ class DefaultConfig:
     :param bias_init: ((Callable) initializer of bias weights (default nn.initializers.zeros_init())
     """ 
     geom: str = 'sphere'
-    monge_map: str = 'row_iter'
+    monge_map: str = 'rounded_matching'
     wasserstein_eps: float = 0.002
     wasserstein_lse: bool = True
     num_sinkhorn_iters: int = 200
     mini_batch_ot_mode: bool = True
-    mini_batch_ot_solver: str = 'chamfer'
+    mini_batch_ot_solver: str = 'frechet'
     minibatch_ot_eps: float = 0.01
     minibatch_ot_lse: bool = True
-    noise_type: str = 'chol_normal'
+    noise_type: str = 'ambient_gaussian'
     scaling: str = 'None'
-    noise_df_scale: float = 2
     factor: float = 1.0
     embedding_dim: int = 512
     num_layers: int = 6
     num_heads: int = 4
     dropout_rate: float = 0.1
     mlp_hidden_dim: int = 512
+    cfg: bool = True
+    p_cfg_null: float = 0.1
+    w_cfg: float = 2.0
+    normalized_condition: bool = False

src/wassersteinflowmatching/riemannian_wasserstein/_utils_Transformer.py

@@ -2,108 +2,217 @@
 import jax.numpy as jnp  # type: ignore
 import jax.random as random  # type: ignore
 from flax import linen as nn  # type: ignore
+from flax.linen.initializers import ones, zeros  # type: ignore
+from typing import Sequence, Optional, Callable
 
-from wassersteinflowmatching.riemannian_wasserstein.DefaultConfig import DefaultConfig
 
-class FeedForward(nn.Module):
-    """Transformer MLP / feed-forward block.
 
-    Attributes:
-    config: DefaultConfig dataclass containing hyperparameters.
-    out_dim: optionally specify out dimension.
-    """
-    config: DefaultConfig
+class SetLayerNorm(nn.Module):
+    epsilon: float = 1e-6
+    scale_init: Callable = ones
+    bias_init: Callable = zeros
 
     @nn.compact
-    def __call__(self, inputs):
-        config = self.config
-        mlp_hidden_dim = config.mlp_hidden_dim
+    def __call__(self, x: jnp.ndarray, mask: Optional[jnp.ndarray] = None) -> jnp.ndarray:
+
+        assert x.ndim == 3, f"Input 'x' must be 3D (B, N, d), but got shape {x.shape}"
+        B, N, d = x.shape
 
-        x = nn.Dense(features = mlp_hidden_dim)(inputs)
-        x = nn.leaky_relu(x)
-        output = nn.Dense(inputs.shape[-1])(x) + inputs
-        return output
+        if mask is None:
+            mask = jnp.ones((B, N), dtype=jnp.int32)
+        else:
+            assert mask.ndim == 2, f"Input 'mask' must be 2D (B, N), but got shape {mask.shape}"
+            assert x.shape[0] == mask.shape[0], \
+                f"Batch dimension mismatch: x {x.shape[0]} vs mask {mask.shape[0]}"
+            assert x.shape[1] == mask.shape[1], \
+                f"N_points dimension mismatch: x {x.shape[1]} vs mask {mask.shape[1]}"
 
-class EncoderBlock(nn.Module):
-    """Transformer encoder layer.
+        mask_expanded = (mask > 0)[..., None].astype(x.dtype)
 
-    Attributes:
-    config: DefaultConfig dataclass containing hyperparameters.
-    """
-    config: DefaultConfig
+        num_valid_elements = jnp.sum(mask, axis=-1) * d
+        num_valid_elements = num_valid_elements[:, None, None]
 
-
-    @nn.compact
-    def __call__(self, inputs, masks, deterministic, dropout_rng = random.key(0)):
+        denominator = jnp.maximum(num_valid_elements, 1.0)
 
-        config = self.config
-        num_heads = config.num_heads
-        dropout_rate = config.dropout_rate
-        # Attention block.
-        # x = nn.LayerNorm(dtype=config.dtype)(inputs)
-        x = nn.MultiHeadDotProductAttention(
-            num_heads = num_heads,
-            dropout_rate=dropout_rate,
-            deterministic=deterministic
-        )(inputs, mask = masks[:, None, None, :],  dropout_rng = dropout_rng) + inputs
-
-        #x = nn.Dropout(rate=config.attention_dropout_rate)(x, deterministic=deterministic)
-        x = nn.LayerNorm()(x)
-        x = FeedForward(config)(x)
-        output = nn.LayerNorm()(x)
-        return output
+        sum_val = jnp.sum(x * mask_expanded, axis=(-2, -1), keepdims=True)
+        mean = sum_val / denominator
 
-class AttentionNN(nn.Module):
+        sq_diff = jnp.square(x - mean) * mask_expanded
+        sum_sq_diff = jnp.sum(sq_diff, axis=(-2, -1), keepdims=True)
+        variance = sum_sq_diff / denominator
 
-    config: DefaultConfig
+        x_norm = (x - mean) / jnp.sqrt(variance + self.epsilon)
 
-    @nn.compact
-    def __call__(self, point_cloud, t, masks = None, labels = None, deterministic = True, dropout_rng=random.key(0)):
+        scale = self.param('scale', self.scale_init, (d,))
+        bias = self.param('bias', self.bias_init, (d,))
 
-        config = self.config
+        y = x_norm * scale + bias
 
-        embedding_dim = config.embedding_dim
-        num_layers = config.num_layers
+        y_masked = y * mask_expanded
 
-        if(labels is not None):
-            embedding_dim =  (3 * config.num_heads) * (embedding_dim // (3 * config.num_heads))
-            concat_dim = (embedding_dim // 3)
-        else:
-            embedding_dim =  (2 * config.num_heads) * (embedding_dim // (2 * config.num_heads))
-            concat_dim = (embedding_dim // 2)
+        return y_masked
+
+
+class FeedForward(nn.Module):
+    """
+    Transformer MLP / feed-forward block.
+    (Pre-norm compatible version)
+    """
+    config: dict
+
+    @nn.compact
+    def __call__(self, inputs, deterministic: bool = True, dropout_rng=None):
+        # Use mlp_hidden_dim from your config
+        mlp_hidden_dim = self.config.mlp_hidden_dim
 
+        x = nn.Dense(features=mlp_hidden_dim)(inputs)
+        x = nn.Dropout(rate=self.config.dropout_rate)(x, deterministic=deterministic, rng=dropout_rng)
+        x = nn.leaky_relu(x) 
+        output = nn.Dense(inputs.shape[-1])(x)
+        return output
 
 
-        space_dim = point_cloud.shape[-1]
+class EncoderBlock(nn.Module):
+    """
+    Transformer encoder layer (optionally conditioned) using 
+    Pre-Normalization for improved stability in deep networks.
+    """
 
-        if masks is None:
-            masks = jnp.ones((point_cloud.shape[0],point_cloud.shape[1]))
-        x_emb = nn.Dense(features = concat_dim)(point_cloud)
-
-        freqs = jnp.arange(embedding_dim) 
-        t_freq = freqs[None, :] * t[:, None]
-        t_four = jnp.concatenate([jnp.cos(t_freq), jnp.sin(t_freq)], axis = -1)
-        t_emb = nn.Dense(features = concat_dim)(t_four)
+    config: dict
 
-        x = jnp.concatenate([x_emb, 
-                             jnp.tile(t_emb[:, None, :], [1, point_cloud.shape[1], 1])], axis = -1)
+    @nn.compact
+    def __call__(
+        self, 
+        inputs: jnp.ndarray, 
+        masks: Optional[jnp.ndarray] = None, 
+        dropout_rng: Optional[jnp.ndarray] = None, 
+        t_emb: Optional[jnp.ndarray] = None, 
+        c_emb: Optional[jnp.ndarray] = None,
+        deterministic: bool = True
+    ) -> jnp.ndarray:
+
+        num_heads = self.config.num_heads
+        dropout_rate = self.config.dropout_rate
 
+        # --- 1. Conditioning ---
+        # Same logic as before: add time and context embeddings
+        # We apply this *before* the first normalization
+
+        conditioning = jnp.zeros_like(inputs)
+        if t_emb is not None:
+            conditioning += t_emb[:, None, :]
+        if c_emb is not None:
+            conditioning += c_emb[:, None, :]
+
+        conditioned_inputs = inputs + conditioning
+        attn_mask = masks[:, None, None, :] if masks is not None else None
+
+        attn_rng, ff_rng = jax.random.split(dropout_rng) if dropout_rng is not None else (None, None)
+        #normed_inputs = SetLayerNorm()(conditioned_inputs, mask=masks)
+        normed_inputs = nn.LayerNorm()(conditioned_inputs)
+        attn_output = nn.MultiHeadDotProductAttention(
+            num_heads=num_heads,
+            dropout_rate=dropout_rate
+        )(
+            normed_inputs, 
+            mask=attn_mask, 
+            deterministic=deterministic,
+            dropout_rng=attn_rng
+        )
+
+        x = inputs + attn_output
+        #normed_x = SetLayerNorm()(x, mask=masks)
+        normed_x = nn.LayerNorm()(x)
+        ff_output = FeedForward(config=self.config)(normed_x, deterministic=deterministic, dropout_rng=ff_rng)
+        output = x + ff_output
+
+        return output
 
-        if(labels is not None):
-            l_emb = nn.Dense(features = concat_dim)(jax.nn.one_hot(labels, config.label_dim))
-            x = jnp.concatenate([x, 
-                                 jnp.tile(l_emb[:, None, :], [1, point_cloud.shape[1], 1])], axis = -1)
 
-
-        for _ in range(num_layers):
-            x = EncoderBlock(config)(inputs = x, masks = masks, deterministic = deterministic, dropout_rng = dropout_rng)   
+class AttentionNN(nn.Module):
+    """
+    Main attention network for WassersteinFlowMatching.
+
+    This version supports selective hypersphere projection for conditioning vectors.
+    """
 
+    config: dict
 
-        x = nn.Dense(features=space_dim, 
-                    kernel_init=nn.initializers.variance_scaling(1e-3, mode='fan_in', distribution='truncated_normal'), 
-                    bias_init=nn.initializers.zeros)(x)
-        return x
+    @nn.compact
+    def __call__(
+        self,
+        point_cloud,
+        t,
+        masks=None,
+        conditioning: Optional[jnp.ndarray] = None,
+        is_null_conditioning: Optional[jnp.ndarray] = None, # <-- New argument
+        deterministic: bool = True,
+        dropout_rng=None,
+    ):
+        config = self.config
+        embedding_dim = config.embedding_dim
+        num_layers = config.num_layers
+        space_dim = point_cloud.shape[-1]
+
+        embedding_dim = config.num_heads * (embedding_dim // config.num_heads)
 
+        if masks is None:
+            masks = jnp.ones((point_cloud.shape[0], point_cloud.shape[1]))
 
+        # 1. Embed point cloud and time (no changes here)
+        x_emb = nn.Dense(features=embedding_dim)(point_cloud)
 
-
+        freqs = jnp.arange(embedding_dim // 2) * (2.0 * jnp.pi / embedding_dim)
+        t_freq = freqs[None, :] * t[:, None]
+        t_four = jnp.concatenate([jnp.cos(t_freq), jnp.sin(t_freq)], axis=-1)
+        t_emb = nn.Dense(features=embedding_dim)(t_four)
+
+        # 2. Process conditioning vector with new logic
+        c_emb = None
+        if conditioning is not None:
+            if is_null_conditioning is None:
+                is_null_conditioning = jnp.zeros(conditioning.shape[0], dtype=bool)
+
+            if config.normalized_condition:
+                # First, pass all conditioning vectors through the dense layer
+                c_emb = nn.Dense(features=embedding_dim)(conditioning)
+
+                # Project the embeddings onto the unit hypersphere
+                norm = jnp.linalg.norm(c_emb, axis=-1, keepdims=True)
+                c_emb = c_emb / (norm + 1e-8)
+
+                # Selectively set the embedding to zero for samples marked as null
+                c_emb = jnp.where(
+                    is_null_conditioning[:, None], 
+                    jnp.zeros_like(c_emb), 
+                    c_emb
+                )
+            else:
+                # Zero out the input conditioning vectors first
+                conditioning = jnp.where(
+                    is_null_conditioning[:, None],
+                    jnp.zeros_like(conditioning),
+                    conditioning
+                )
+                # Then pass through the dense layer
+                c_emb = nn.Dense(features=embedding_dim)(conditioning)
+
+        # 3. Apply Transformer blocks (no changes here)
+        x = x_emb
+        for _ in range(num_layers):
+            dropout_rng, key = jax.random.split(dropout_rng) if dropout_rng is not None else (None, None)
+
+            x = EncoderBlock(
+                config=config
+            )(x, deterministic=deterministic, masks=masks, dropout_rng=dropout_rng, t_emb=t_emb, c_emb=c_emb)
+
+        # 4. Final output layer (no changes here)
+        x = nn.Dense(
+            features=space_dim,
+            kernel_init=nn.initializers.variance_scaling(
+                1e-2, mode="fan_in", distribution="truncated_normal"
+            ),
+            bias_init=nn.initializers.zeros,
+        )(x)
+        #x = nn.Dense(features=space_dim)(x)
+        return x