fix: better initialization of masked flows

Author: aseyboldt

python/nutpie/normalizing_flow.py

@@ -1054,13 +1054,19 @@ def make_transformer():
     key, key1 = jax.random.split(key)
     embed = eqx.nn.Sequential(
         [
-            eqx.nn.Linear(dim, n_embed, key=key1, dtype=jnp.float32),
+            eqx.nn.Linear(dim, n_embed, key=key1, dtype=jnp.float32, use_bias=True),
             # Activation(_NN_ACTIVATION),
             # eqx.nn.LayerNorm(shape=(n_embed,), dtype=jnp.float32),
         ]
     )
     key, key1 = jax.random.split(key)
-    embed_back = eqx.nn.Linear(n_deembed, size, key=key1, dtype=jnp.float32)
+    embed_back = eqx.nn.Linear(
+        n_deembed, size, key=key1, dtype=jnp.float32, use_bias=True
+    )
+    embed_back = jax.tree_util.tree_map(
+        lambda x: x * 1e-3 if eqx.is_inexact_array(x) else x,
+        embed_back,
+    )
 
     rng = np.random.default_rng(42)  # TODO
     order, counts = _generate_permutations(rng, dim, n_layers)
@@ -1077,20 +1083,25 @@ def make_mvscale(key, n_dim):
     def make_layer(key, mask, embed, embed_back):
         key1, key2, key3, key4, key5 = jax.random.split(key, 5)
         transformer = make_transformer()
-        bias = Add(jax.random.normal(key5, (size,)) * 0.01)
+        bias = Add(jax.random.normal(key5, (size,)) * 0.001)
+        inner = eqx.nn.MLP(
+            n_embed,
+            n_deembed,
+            width_size=nn_width,
+            depth=nn_depth,
+            key=key2,
+            dtype=jnp.float32,
+            activation=_NN_ACTIVATION,
+        )
+        inner = jax.tree_util.tree_map(
+            lambda x: x * 1e-3 if eqx.is_inexact_array(x) else x,
+            inner,
+        )
 
         conditioner = eqx.nn.Sequential(
             [
                 embed,
-                eqx.nn.MLP(
-                    n_embed,
-                    n_deembed,
-                    width_size=nn_width,
-                    depth=nn_depth,
-                    key=key2,
-                    dtype=jnp.float32,
-                    activation=_NN_ACTIVATION,
-                ),
+                inner,
                 eqx.nn.Sequential(
                     [
                         embed_back,
@@ -1110,11 +1121,6 @@ def make_layer(key, mask, embed, embed_back):
             nn_depth=nn_depth,
         )
 
-        coupling = jax.tree_util.tree_map(
-            lambda x: x * 1e-3 if eqx.is_inexact_array(x) else x,
-            coupling,
-        )
-
         if mvscale:
             scale = make_mvscale(key4, dim)
             return bijections.Chain([coupling, scale])

python/nutpie/transform_adapter.py

@@ -112,24 +112,14 @@ def fit_to_data(
 
     for i in loop:
         # Shuffle data
-        start = time.time()
         key, *subkeys = jr.split(key, 3)
         train_data = [jr.permutation(subkeys[0], a) for a in train_data]
         val_data = [jr.permutation(subkeys[1], a) for a in val_data]
-        if verbose and i == 0:
-            print("shuffle timing:", time.time() - start)
-
-        start = time.time()
 
         key, subkey = jr.split(key)
         batches = get_batches(train_data, batch_size)
         batch_losses = []
 
-        if verbose and i == 0:
-            print("batch timing:", time.time() - start)
-
-        start = time.time()
-
         if True:
             for batch in zip(*batches, strict=True):
                 key, subkey = jr.split(key)
@@ -156,10 +146,6 @@ def fit_to_data(
 
         losses["train"].append((sum(batch_losses) / len(batch_losses)).item())
 
-        if verbose and i == 0:
-            print("step timing:", time.time() - start)
-
-        start = time.time()
         # Val epoch
         batch_losses = []
         for batch in zip(*get_batches(val_data, batch_size), strict=True):
@@ -168,9 +154,6 @@ def fit_to_data(
             batch_losses.append(loss_i)
         losses["val"].append(sum(batch_losses) / len(batch_losses))
 
-        if verbose and i == 0:
-            print("val timing:", time.time() - start)
-
         loop.set_postfix({k: v[-1] for k, v in losses.items()})
         if losses["val"][-1] == min(losses["val"]):
             best_params = params
@@ -228,7 +211,7 @@ def inverse_gradient_and_val(bijection, draw, grad, logp):
         )
     elif isinstance(bijection, bijections.Affine):
         draw, logdet = bijection.inverse_and_log_det(draw)
-        grad = grad * bijection.scale
+        grad = grad * unwrap(bijection.scale)
         return (draw, grad, logp - logdet)
     elif isinstance(bijection, bijections.Vmap):
 
@@ -710,12 +693,9 @@ def update(self, seed, positions, gradients, logps):
             )
             params, static = eqx.partition(flow, eqx.is_inexact_array)
 
-            start = time.time()
             new_loss = self._loss_fn(
                 params, static, positions[-128:], gradients[-128:], logps[-128:]
             )
-            if self._verbose:
-                print("new loss function time: ", time.time() - start)
 
             if self._verbose:
                 print(f"Chain {self._chain}: New loss {new_loss}, old loss {old_loss}")
@@ -903,8 +883,8 @@ def make_transform_adapter(
     make_optimizer=None,
     coupling_type="masked",
     mvscale_layer=False,
-    n_embed=None,
-    n_deembed=None,
+    num_project=None,
+    num_embed=None,
 ):
     if extension_windows is None:
         extension_windows = []
@@ -918,8 +898,8 @@ def make_transform_adapter(
             dct_layer=dct_layer,
             nn_depth=nn_depth,
             nn_width=nn_width,
-            n_embed=n_embed,
-            n_deembed=n_deembed,
+            n_embed=num_project,
+            n_deembed=num_embed,
             mvscale=mvscale_layer,
             kind=coupling_type,
         ),