Does stax.serial support Tanh activation about neural-tangents HOT 3 CLOSED

Hi Wei,

Unfortunately we don't have a Tanh function since we don't know of a closed-form expression to compute the kernel with this activation.

Currently you could try the following workarounds:

1. Use stax.Erf() instead, which is very similar to Tanh and has a closed-form kernel expression.

2. Use the original jax.experimental.stax (accessed as nt.stax.ostax in the snippet below) to construct the network and use NT to compute the Monte Carlo estimate of the kernel, e.g.

from neural_tangents import stax, monte_carlo_kernel_fn
from jax import random

init_fn, apply_fn = stax.ostax.serial(stax.ostax.Dense(512), stax.ostax.Tanh)
kernel_fn = monte_carlo_kernel_fn(init_fn, apply_fn, 
                                  key=random.PRNGKey(1), n_samples=100)

kernel_fn(x1=random.normal(random.PRNGKey(1), (10, 40)), 
          x2=random.normal(random.PRNGKey(2), (20, 40)), 
          get=('nngp', 'ntk'))

3. You could look into approximating the intractable integrals with a lookup table as was done in https://arxiv.org/pdf/1711.00165.pdf (sec 2.5; https://github.com/brain-research/nngp/blob/42a5fb523653c0a213851628bc3ca0b78a2e8d1c/nngp.py#L307), but this would likely require a lot of effort to implement (NT does not support this currently).

Let me know if any of this helps!

from neural-tangents.

Hi Roman,

Many thanks for your reply. The second method works for Tanh function very well.

I just have a suggestion for your consideration.

If erf(x) changes to erf(\sqrt{\pi}/2 x), it is going to look more like the Tanh function, in the context of
the mean field theory (https://arxiv.org/abs/1802.09979). Since the intersection between critical line and \sigma_b = 0 is \sigma_w=1 for both erf(\sqrt{\pi}/2 x) and Tanh.

Best regards,
Wei

from neural-tangents.

Sorry for huge delay, FYI @SiuMath added this option in d7fc2db

from neural-tangents.