Repository: pyro-ppl/numpyro
Pull request #1897 (2024-10-26 14:44:55)

Ensure BlockNeuralAutoregressiveTransform maps real->real
This fixes https://github.com/pyro-ppl/numpyro/issues/1655 by changing BNAF activation to $y=\text{tanh}(x) + cx$, by default (with c=0.01). This ensures the transform maps from real->real, which e.g. 

1) Means ``BlockNeuralAutoregressiveTransform`` will have the correct codomain (currently set as ``real_vector``)
2) Unbounded distributions transformed using the transform will remain unbounded
3) Transforming a distribution with the inverse transformation will now lead to properly normalised densities

Let me know if you think there are any problems with this solution.


fehiepsi: Maybe using LeakyReLU likes in the paper instead?

By the way, is there a reference for LeakyTanh?
danielward27: I believe that all the results from the paper use tanh. From personal experience (on admittedly a very limited set of applications), leaky relu performs far worse than either tanh or leaky tanh.

No, I don't have a reference for it, I just chose it to solve the problem to ensure dy/dx is at least a threshold value, whilst maintaining similar results on a couple toy density estimation tasks.

I completely understand if you'd prefer to keep it to how it is defined in the paper, but IMO the other issues are bound to arise some point e.g. people expecting (sensibly defined) BNAF guides to be supported on the same set as the true posterior, or expecting normalised densities when transforming with the inverse.

fehiepsi: LGTM! Could you double check if the change still works for https://num.pyro.ai/en/stable/examples/covtype.html and https://num.pyro.ai/en/stable/examples/neutra.html tutorials?
danielward27: Neutra HMC example:
```
Start vanilla HMC...
sample: 100%|█████████████████████████████████████████████████| 5000/5000 [00:07<00:00, 694.39it/s, 15 steps of size 1.97e-01. acc. prob=0.93]

                mean       std    median      5.0%     95.0%     n_eff     r_hat
      x[0]     -0.71      1.67     -1.56     -2.40      2.08     14.57      1.00
      x[1]      0.04      1.16      0.01     -1.82      1.84    776.44      1.01

Number of divergences: 0
Start training guide...
100%|█████████████████████████████████████████████| 10000/10000 [00:06<00:00, 1584.22it/s, init loss: 1.3300, avg. loss [9501-10000]: -1.7834]
Finish training guide. Extract samples...

Start NeuTra HMC...
sample: 100%|█████████████████████████████████████████████████| 5000/5000 [00:15<00:00, 319.74it/s, 13 steps of size 1.25e-01. acc. prob=0.80]

                           mean       std    median      5.0%     95.0%     n_eff     r_hat
auto_shared_latent[0]     -0.15      1.00     -0.08     -1.82      1.38    750.76      1.01
auto_shared_latent[1]      0.12      1.00      0.12     -1.59      1.52    803.72      1.00

Number of divergences: 14
Transform samples into unwarped space...

                mean       std    median      5.0%     95.0%     n_eff     r_hat
      x[0]     -0.05      1.82     -1.00     -2.31      2.23    594.42      1.01
      x[1]      0.05      1.16      0.10     -1.81      1.81    833.70      1.00
```
[neutra.pdf](https://github.com/user-attachments/files/18021045/neutra.pdf)

Covtype results:
```
sample: 100%|██████████████████████████████████████████████████| 2000/2000 [05:03<00:00,  6.58it/s, 15 steps of size 4.09e-02. acc. prob=0.90]
Mean accept prob: 0.4889448

                            mean       std    median      5.0%     95.0%     n_eff     r_hat
 auto_shared_latent[0]     -0.29      0.09     -0.29     -0.44     -0.14    298.38      1.04
 auto_shared_latent[1]     -0.04      0.15     -0.04     -0.28      0.21    546.84      1.00
 auto_shared_latent[2]     -0.15      0.17     -0.15     -0.42      0.11     46.35      1.03
 auto_shared_latent[3]      0.01      0.15      0.01     -0.22      0.28    725.17      1.00
 auto_shared_latent[4]     -0.22      0.11     -0.23     -0.40     -0.04   1191.72      1.00
 auto_shared_latent[5]      0.28      0.16      0.28      0.04      0.57     90.04      1.06
 auto_shared_latent[6]     -0.05      0.17     -0.04     -0.31      0.22   1423.81      1.00
 auto_shared_latent[7]     -0.01      0.16     -0.01     -0.27      0.25    350.08      1.00
 auto_shared_latent[8]     -0.57      0.15     -0.57     -0.80     -0.32    128.09      1.04
 auto_shared_latent[9]     -0.24      0.19     -0.25     -0.55      0.06     41.80      1.09
auto_shared_latent[10]     -0.11      0.22     -0.11     -0.46      0.25    217.61      1.02
auto_shared_latent[11]     -0.33      0.18     -0.33     -0.60     -0.02   1224.29      1.01
auto_shared_latent[12]     -0.20      0.17     -0.19     -0.48      0.07    105.26      1.00
auto_shared_latent[13]      0.26      0.18      0.25     -0.04      0.55    205.09      1.03
auto_shared_latent[14]     -0.50      0.23     -0.50     -0.86     -0.12    201.40      1.01
auto_shared_latent[15]      0.55      0.21      0.55      0.21      0.89    569.94      1.00
auto_shared_latent[16]      0.14      0.17      0.14     -0.17      0.41    664.81      1.00
auto_shared_latent[17]      0.02      0.20      0.03     -0.28      0.37   1319.47      1.00
auto_shared_latent[18]      0.21      0.19      0.21     -0.09      0.53   1071.26      1.00
auto_shared_latent[19]      0.42      0.19      0.42      0.11      0.72     85.36      1.02
auto_shared_latent[20]      0.57      0.31      0.57      0.07      1.10     11.20      1.24
auto_shared_latent[21]     -0.61      0.26     -0.62     -1.06     -0.21     79.74      1.01
auto_shared_latent[22]     -0.04      0.26     -0.04     -0.47      0.36     59.62      1.07
auto_shared_latent[23]     -0.06      0.27     -0.05     -0.51      0.37    801.37      1.00
auto_shared_latent[24]     -0.13      0.29     -0.12     -0.61      0.34    129.50      1.01
auto_shared_latent[25]     -0.10      0.33     -0.10     -0.60      0.45    879.86      1.00
auto_shared_latent[26]      0.91      0.25      0.92      0.52      1.30     13.22      1.18
auto_shared_latent[27]      0.07      0.27      0.07     -0.36      0.49   1167.03      1.00
auto_shared_latent[28]     -0.03      0.25     -0.04     -0.47      0.35    394.86      1.02
auto_shared_latent[29]      0.68      0.31      0.69      0.17      1.16    483.95      1.00
auto_shared_latent[30]      0.29      0.28      0.29     -0.12      0.79    595.76      1.03
auto_shared_latent[31]     -0.12      0.30     -0.12     -0.56      0.42   1034.56      1.00
auto_shared_latent[32]     -0.43      0.31     -0.43     -0.93      0.08    291.05      1.02
auto_shared_latent[33]      0.12      0.26      0.11     -0.34      0.53    232.03      1.02
auto_shared_latent[34]      0.27      0.31      0.27     -0.21      0.82    134.19      1.05
auto_shared_latent[35]     -1.09      0.32     -1.10     -1.55     -0.49    438.95      1.00
auto_shared_latent[36]     -0.19      0.26     -0.18     -0.58      0.24   1692.49      1.00
auto_shared_latent[37]     -0.37      0.25     -0.38     -0.78      0.01    101.38      1.04
auto_shared_latent[38]     -0.57      0.39     -0.57     -1.23      0.05    528.35      1.01
auto_shared_latent[39]     -0.00      0.35     -0.00     -0.56      0.56    864.91      1.00
auto_shared_latent[40]      0.82      0.30      0.82      0.36      1.34   1286.39      1.00
auto_shared_latent[41]      0.55      0.27      0.56      0.11      0.98    793.97      1.02
auto_shared_latent[42]     -0.27      0.19     -0.27     -0.58      0.01    146.59      1.05
auto_shared_latent[43]     -0.31      0.19     -0.31     -0.60      0.03    411.32      1.02
auto_shared_latent[44]     -0.05      0.25     -0.04     -0.45      0.40   1385.11      1.00
auto_shared_latent[45]     -0.36      0.33     -0.37     -0.90      0.15    472.51      1.00
auto_shared_latent[46]      0.17      0.36      0.18     -0.51      0.69    789.30      1.00
auto_shared_latent[47]     -0.13      0.36     -0.13     -0.75      0.44   1040.68      1.00
auto_shared_latent[48]      0.21      0.36      0.22     -0.42      0.77    881.87      1.00
auto_shared_latent[49]     -0.88      0.51     -0.86     -1.64      0.02    436.51      1.01
auto_shared_latent[50]      0.32      0.46      0.32     -0.48      1.00    707.01      1.00
auto_shared_latent[51]      0.76      0.57      0.76     -0.19      1.65    383.05      1.02
auto_shared_latent[52]     -0.14      0.53     -0.14     -0.93      0.81    169.44      1.00
auto_shared_latent[53]     -1.39      0.37     -1.39     -1.99     -0.78    778.57      1.00
auto_shared_latent[54]      1.40      0.28      1.40      0.93      1.85    690.11      1.01
              coefs[0]      1.97      0.01      1.97      1.96      1.98    264.58      1.04
              coefs[1]     -0.04      0.00     -0.04     -0.05     -0.04    628.71      1.00
              coefs[2]     -0.10      0.00     -0.10     -0.10     -0.09     43.07      1.04
              coefs[3]     -0.30      0.00     -0.30     -0.30     -0.30    447.55      1.00
              coefs[4]     -0.09      0.00     -0.09     -0.10     -0.09   1061.69      1.00
              coefs[5]     -0.14      0.00     -0.14     -0.15     -0.14    395.67      1.02
              coefs[6]      0.08      0.01      0.08      0.07      0.09    454.57      1.00
              coefs[7]     -0.56      0.00     -0.56     -0.56     -0.55   1211.38      1.01
              coefs[8]      0.39      0.01      0.39      0.38      0.41   1080.88      1.00
              coefs[9]     -0.02      0.00     -0.02     -0.02     -0.01     61.51      1.07
             coefs[10]      0.35      0.01      0.35      0.32      0.36   1246.26      1.00
             coefs[11]     -0.08      0.01     -0.08     -0.09     -0.07   1128.76      1.01
             coefs[12]      0.02      0.01      0.02     -0.00      0.04    960.22      1.00
             coefs[13]     -0.89      0.03     -0.89     -0.94     -0.84    206.75      1.04
             coefs[14]     -0.04      0.01     -0.04     -0.06     -0.03    470.03      1.00
             coefs[15]     -0.84      0.07     -0.84     -0.93     -0.71    991.30      1.00
             coefs[16]     -0.49      0.03     -0.48     -0.53     -0.45    454.90      1.00
             coefs[17]     -0.10      0.01     -0.10     -0.11     -0.08   1168.92      1.01
             coefs[18]     -0.46      0.05     -0.46     -0.54     -0.39   1085.81      1.00
             coefs[19]     -0.18      0.03     -0.18     -0.22     -0.14     41.67      1.07
             coefs[20]     -1.05      0.05     -1.05     -1.13     -0.96     18.74      1.16
             coefs[21]     -0.00      0.00     -0.00     -0.00      0.00     87.85      1.00
             coefs[22]      0.05      0.00      0.05      0.05      0.05     45.99      1.07
             coefs[23]      0.02      0.00      0.02      0.01      0.03    421.21      1.02
             coefs[24]     -0.01      0.00     -0.01     -0.01     -0.00    223.80      1.01
             coefs[25]     -0.15      0.00     -0.15     -0.16     -0.15     73.17      1.04
             coefs[26]     -0.03      0.00     -0.03     -0.03     -0.02     12.90      1.19
             coefs[27]     -0.48      0.05     -0.48     -0.56     -0.40   1246.89      1.00
             coefs[28]     -0.46      0.05     -0.46     -0.53     -0.38     37.09      1.09
             coefs[29]      0.03      0.00      0.03      0.03      0.04    130.23      1.01
             coefs[30]      0.02      0.00      0.02      0.01      0.02    851.93      1.02
             coefs[31]     -0.02      0.00     -0.02     -0.03     -0.02   1182.32      1.00
             coefs[32]      0.05      0.00      0.05      0.05      0.06     35.03      1.08
             coefs[33]      0.10      0.00      0.10      0.09      0.10     95.00      1.04
             coefs[34]      0.13      0.00      0.13      0.12      0.13    130.54      1.05
             coefs[35]      0.24      0.00      0.24      0.23      0.25     27.56      1.09
             coefs[36]      0.22      0.00      0.22      0.21      0.23     28.31      1.09
             coefs[37]      0.10      0.00      0.10      0.10      0.11     21.46      1.11
             coefs[38]     -0.03      0.00     -0.03     -0.03     -0.02    362.09      1.02
             coefs[39]     -0.01      0.00     -0.01     -0.02     -0.01    782.24      1.00
             coefs[40]      0.03      0.00      0.03      0.03      0.04   1447.84      1.01
             coefs[41]     -0.02      0.00     -0.02     -0.02     -0.02    300.47      1.03
             coefs[42]     -0.11      0.01     -0.11     -0.12     -0.10     49.22      1.05
             coefs[43]     -0.09      0.00     -0.09     -0.10     -0.08    432.89      1.00
             coefs[44]      0.09      0.00      0.09      0.09      0.10     64.63      1.05
             coefs[45]     -0.00      0.00     -0.00     -0.01      0.01     24.46      1.11
             coefs[46]      0.12      0.00      0.12      0.11      0.12     60.11      1.06
             coefs[47]     -0.09      0.00     -0.09     -0.10     -0.08    785.12      1.00
             coefs[48]     -0.09      0.00     -0.09     -0.10     -0.09    668.53      1.00
             coefs[49]     -0.03      0.00     -0.03     -0.04     -0.03    453.09      1.02
             coefs[50]     -0.48      0.04     -0.48     -0.55     -0.41    822.76      1.00
             coefs[51]     -0.18      0.00     -0.18     -0.18     -0.17     18.01      1.16
             coefs[52]     -0.17      0.00     -0.17     -0.17     -0.16     23.29      1.11
             coefs[53]     -0.20      0.00     -0.20     -0.21     -0.20    275.47      1.04
             coefs[54]     -1.48      0.01     -1.48     -1.50     -1.46    425.51      1.00
```

covtype svi losses:
![covtype_losses](https://github.com/user-attachments/assets/0e8a5c24-cf58-4ada-9f35-f018d93ed761)



Patch:
diff --git a/numpyro/nn/block_neural_arn.py b/numpyro/nn/block_neural_arn.py
index 452ae55f..1d2e00c7 100644
--- a/numpyro/nn/block_neural_arn.py
+++ b/numpyro/nn/block_neural_arn.py
@@ -114,6 +114,32 @@ def Tanh():
     return init_fun, apply_fun
 
 
+def LeakyTanh(min_grad: float = 0.01):
+    """
+    Leaky Tanh nonlinearity :math:`y=\text{tanh}(x) + cx` with its log-Jacobian.
+
+    This choice when used in ``BlockNeuralAutoregressiveNN`` ensures the image of the
+    transformation is the set of real values (unlike ``Tanh``).
+
+    :param float min_grad: The minimum gradient value (:math:`c` above). Defaults to 0.01.
+    :return: an (`init_fn`, `apply_fn`) pair.
+    """
+
+    def init_fun(rng, input_shape):
+        return input_shape, ()
+
+    def apply_fun(params, inputs, **kwargs):
+        x, logdet = inputs
+        out = jnp.tanh(x) + min_grad * x  # ensure grad at least 0.01.
+        tanh_logdet = -2 * (x + softplus(-2 * x) - jnp.log(2.0))
+        act_logdet = jnp.logaddexp(tanh_logdet, jnp.log(min_grad))
+        # Reshape to match logdet shape
+        act_logdet = act_logdet.reshape(logdet.shape[:-2] + (1, logdet.shape[-1]))
+        return out, logdet + act_logdet
+
+    return init_fun, apply_fun
+
+
 def FanInResidualNormal():
     """
     Similar to stax.FanInSum but also keeps track of log determinant of Jacobian.
@@ -156,10 +182,19 @@ def FanInResidualGated(gate_init=normal(1.0)):
     return init_fun, apply_fun
 
 
-def BlockNeuralAutoregressiveNN(input_dim, hidden_factors=[8, 8], residual=None):
+def BlockNeuralAutoregressiveNN(
+    input_dim,
+    hidden_factors=[8, 8],
+    residual=None,
+    activation=None,
+):
     """
     An implementation of Block Neural Autoregressive neural network.
 
+    In contrast to the original paper, by default, we use ``LeakyTanh`` as the
+    activation, defined as :math:`y=Tanh(x) + cx` with :math:`c` being a small constant
+    (default to 0.01), which ensures the transform maps from real -> real.
+
     **References**
 
     1. *Block Neural Autoregressive Flow*,
@@ -170,13 +205,15 @@ def BlockNeuralAutoregressiveNN(input_dim, hidden_factors=[8, 8], residual=None)
         input dimension. This corresponds to both :math:`a` and :math:`b` in reference [1].
         The elements of hidden_factors must be integers.
     :param str residual: Type of residual connections to use. One of `None`, `"normal"`, `"gated"`.
+    :param tuple activation: An (`init_fn`, `update_fn`) pair. Defaults to ``LeakyTanh``.
     :return: an (`init_fn`, `update_fn`) pair.
     """
     layers = []
     in_factor = 1
+    activation = LeakyTanh() if activation is None else activation
     for hidden_factor in hidden_factors:
         layers.append(BlockMaskedDense(input_dim, in_factor, hidden_factor))
-        layers.append(Tanh())
+        layers.append(activation)
         in_factor = hidden_factor
     layers.append(BlockMaskedDense(input_dim, in_factor, 1))
     arn = stax.serial(*layers)


Test patch:
diff --git a/test/test_flows.py b/test/test_flows.py
index 9e65cffb..ce285269 100644
--- a/test/test_flows.py
+++ b/test/test_flows.py
@@ -9,13 +9,17 @@ import pytest
 
 from jax import jacfwd, random
 from jax.example_libraries import stax
+import jax.numpy as jnp
+import jax.random as jr
 
+from numpyro.distributions import Normal, TransformedDistribution
 from numpyro.distributions.flows import (
     BlockNeuralAutoregressiveTransform,
     InverseAutoregressiveTransform,
 )
 from numpyro.distributions.util import matrix_to_tril_vec
 from numpyro.nn import AutoregressiveNN, BlockNeuralAutoregressiveNN
+from numpyro.nn.block_neural_arn import LeakyTanh, Tanh
 
 
 def _make_iaf_args(input_dim, hidden_dims):
@@ -34,8 +38,12 @@ def _make_iaf_args(input_dim, hidden_dims):
     return (partial(arn, init_params),)
 
 
-def _make_bnaf_args(input_dim, hidden_factors):
-    arn_init, arn = BlockNeuralAutoregressiveNN(input_dim, hidden_factors)
+def _make_bnaf_args(input_dim, hidden_factors, activation):
+    arn_init, arn = BlockNeuralAutoregressiveNN(
+        input_dim,
+        hidden_factors,
+        activation=activation,
+    )
     _, rng_key_perm = random.split(random.PRNGKey(0))
     _, init_params = arn_init(random.PRNGKey(0), (input_dim,))
     return (partial(arn, init_params),)
@@ -46,10 +54,24 @@ def _make_bnaf_args(input_dim, hidden_factors):
     [
         (InverseAutoregressiveTransform, _make_iaf_args(5, hidden_dims=[10]), 5),
         (InverseAutoregressiveTransform, _make_iaf_args(7, hidden_dims=[8, 9]), 7),
-        (BlockNeuralAutoregressiveTransform, _make_bnaf_args(7, hidden_factors=[4]), 7),
         (
             BlockNeuralAutoregressiveTransform,
-            _make_bnaf_args(7, hidden_factors=[2, 3]),
+            _make_bnaf_args(7, hidden_factors=[4], activation=LeakyTanh()),
+            7,
+        ),
+        (
+            BlockNeuralAutoregressiveTransform,
+            _make_bnaf_args(7, hidden_factors=[2, 3], activation=LeakyTanh()),
+            7,
+        ),
+        (
+            BlockNeuralAutoregressiveTransform,
+            _make_bnaf_args(7, hidden_factors=[4], activation=Tanh()),
+            7,
+        ),
+        (
+            BlockNeuralAutoregressiveTransform,
+            _make_bnaf_args(7, hidden_factors=[2, 3], activation=Tanh()),
             7,
         ),
     ],
@@ -91,3 +113,21 @@ def test_flows(flow_class, flow_args, input_dim, batch_shape):
 
         assert np.sum(np.abs(np.triu(jac, 1))) == 0.00
         assert np.all(np.abs(matrix_to_tril_vec(jac)) > 0)
+
+
+def test_bnaf_normalization():
+    dim = (1,)
+    x = jnp.linspace(-1000, 1000, 5000)[:, None]
+
+    init_fn, apply_fn = BlockNeuralAutoregressiveNN(dim[0], activation=LeakyTanh(0.1))
+    params = init_fn(jr.PRNGKey(0), (1,))[1]
+    arn = partial(apply_fn, params)
+    bnaf = BlockNeuralAutoregressiveTransform(arn)
+    dist = TransformedDistribution(Normal(jnp.zeros(dim), 0.5), bnaf.inv)
+    probs = jnp.exp(dist.log_prob(x))
+    probs, x = jnp.squeeze(probs), jnp.squeeze(x)
+
+    # Rough integral
+    integral = jnp.trapezoid(probs, x)
+    assert integral > 0.9
+    assert integral < 1.1