<jupyter_start><jupyter_text>Neural processesWe'll demonstrate training a neural process and an attentive neural process model using a batch of draws from a Gaussian process. First load some libraries.<jupyter_code>import jax.numpy as np
import jax.random as random
import haiku as hk

%matplotlib inline
import matplotlib.pyplot as plt
import palettes

from ramsey.train import train_neural_process
from ramsey.data import sample_from_gaussian_process

palettes.set_theme()<jupyter_output><empty_output><jupyter_text>Next we sample some data from a Gaussian process:<jupyter_code>key = random.PRNGKey(12)
key, sample_key = random.split(key, 2)

(x_target, y_target), f_target = sample_from_gaussian_process(
    sample_key,
    batch_size=10,
    num_observations=200
)<jupyter_output><empty_output><jupyter_text>Visualize the batch of draws:<jupyter_code>cols = palettes.discrete_sequential_colors(10)

_, ax = plt.subplots(figsize=(10, 6))
for i in range(x_target.shape[0]):
    x = np.squeeze(x_target[i, :, :])
    y = np.squeeze(y_target[i, :, :])
    f = np.squeeze(f_target[i, :, :])
    idxs = np.argsort(x)
    ax.plot(x[idxs], f[idxs], color=cols[i], alpha=0.5)
    ax.scatter(x[idxs], y[idxs], marker="+", color=cols[i], alpha=0.5)
plt.show()<jupyter_output><empty_output><jupyter_text>Neural processHaving sampled a data set, we can define the model. A neural process model takes a `decoder` argument, a `latent_encoder` argument and an optional `deterministic_encoder` argument.- The `decoder` is a neural network, typically an MLP, that maps the latent representation to the data space, or, more specifically, is used to parameterize an exponential family distribution. The decoder needs to respect the dimensionality of the response to be modelled. For instance, if the response is a univariate Gaussian the last layer needs to have two neurons: one do parameterize the mean and one to parameterize the variance of the Gaussian.- The `latent_encoder` is a tuple of two neural networks, typically also two MLPs. The first of which maps the data to a latent representation. This representation gets averaged and then passed through the second MLP to parameterize a latent Gaussian.- The `deterministic_encoder` is typically also an MLP that maps the data two a deterministic univariate representation.Below we create a simple NP model with a decoder and a latent encoder. All networks are MLPs. The decoder has two neurons in its final layer (for mean and variance of the Gaussian that models the data), the encoder's final layer is 2*128 nodes to parameterize a latent 128 dimensional Gaussian.<jupyter_code>from ramsey.models import NP

def neural_process(**kwargs):
    dim = 128
    np = NP(
        decoder=hk.nets.MLP([dim] * 3 + [2]),
        latent_encoder=(hk.nets.MLP([dim] * 3), hk.nets.MLP([dim, dim * 2]))
    )
    return np(**kwargs)

neural_process = hk.transform(neural_process)<jupyter_output><empty_output><jupyter_text>Next we initialize the set of parameters and train the model with a routine provided in `ramsey`:<jupyter_code>key, init_key = random.split(key, 2)
params = neural_process.init(
    init_key, x_context=x_target, y_context=y_target, x_target=x_target
)<jupyter_output><empty_output><jupyter_text>We use 10 context points and 20 target points per iteration of training:<jupyter_code>key, train_key = random.split(key, 2)

n_context, n_target = 10, 20
params, objectives = train_neural_process(
    neural_process,
    params,
    train_key,
    x=x_target,
    y=y_target,
    n_context=n_context,
    n_target=n_target,
    n_iter=10000
)<jupyter_output>ELBO at 0: -82.83424377441406
ELBO at 1000: -33.90901565551758
ELBO at 2000: -16.32452392578125
ELBO at 3000: -3.868741750717163
ELBO at 4000: -6.080962181091309
ELBO at 5000: -12.974806785583496
ELBO at 6000: -11.733842849731445
ELBO at 7000: -2.2121856212615967
ELBO at 8000: -32.64741516113281
ELBO at 9000: -9.536314010620117
ELBO at 9999: -9.654109954833984<jupyter_text>That's it! Finally, we define a set of data points for a batch and make some predictions:<jupyter_code>key, sample_key = random.split(key, 2)

sample_idxs = random.choice(
    sample_key, x_target.shape[1], 
    shape=(n_target + n_context,), replace=False
)
idxs = [0, 2, 5, 7]

_, axes = plt.subplots(figsize=(10, 6), nrows=2, ncols=2)
for _, (idx, ax) in enumerate(zip(idxs, axes.flatten())):
    x = np.squeeze(x_target[idx, :, :])
    f = np.squeeze(f_target[idx, :, :])
    y = np.squeeze(y_target[idx, :, :])

    srt_idxs = np.argsort(x)
    ax.plot(x[srt_idxs], f[srt_idxs], color="blue", alpha=0.75)
    ax.scatter(
        x[sample_idxs[:n_context]], y[sample_idxs[:n_context]], 
        color="blue", marker="+", alpha=0.75
    )

    for i in range(20):
        key, apply_key = random.split(key, 2)
        y_star = neural_process.apply(
            params=params,
            rng=apply_key,
            x_context=x[np.newaxis, sample_idxs, np.newaxis],
            y_context=y[np.newaxis, sample_idxs, np.newaxis],
            x_target=x_target[[idx], :, :]
        ).mean
        x_star = np.squeeze(x_target[[idx], :, :])
        y_star = np.squeeze(y_star)
        ax.plot(x_star[srt_idxs], y_star[srt_idxs], color="black", alpha=0.1)
plt.show()<jupyter_output><empty_output><jupyter_text>Attentive neural processAn attentive neural process model takes a `decoder` argument, a `latent_encoder` argument and a `deterministic_encoder` argument with an attention module. In comparison to before the`deterministic_encoder` is also a tuple in this case. In addition to the network that generates the representation, one is also requires to provide a mechanism for cross-attention.Below we create an ANP model with a decoder a latent encoder and a deterministic encoder with `MultiHeadAttention`. The encoder itself is again an MLP, the `MultiHeadAttention` module takes also an MLP that embeds the keys and query into a higher dimensional space before attending to the query.<jupyter_code>from ramsey.models import ANP
from ramsey.attention import MultiHeadAttention

def attentive_neural_process(**kwargs):
    dim = 128
    np = ANP(
        decoder=hk.nets.MLP([dim] * 3 + [2]),
        latent_encoder=(hk.nets.MLP([dim] * 3), hk.nets.MLP([dim, dim * 2])),
        deterministic_encoder=(
            hk.nets.MLP([dim] * 3),
            MultiHeadAttention(
                num_heads=8,
                head_size=16,
                embedding=hk.nets.MLP([dim] * 2)
            )
        )
        
    )
    return np(**kwargs)


attentive_neural_process = hk.transform(attentive_neural_process)<jupyter_output><empty_output><jupyter_text>After the setup, the initialization of parameters and training is identical.<jupyter_code>key, init_key, train_key = random.split(key, 3)

params = attentive_neural_process.init(
    init_key, x_context=x_target, y_context=y_target, x_target=x_target
)

params, objectives = train_neural_process(
    attentive_neural_process,
    params,
    train_key,
    x=x_target,
    y=y_target,
    n_context=n_context,
    n_target=n_target,
    n_iter=10000
)<jupyter_output>ELBO at 0: -32.21488952636719
ELBO at 1000: -2.1482131481170654
ELBO at 2000: -11.758901596069336
ELBO at 3000: -4.755011558532715
ELBO at 4000: -3.753476619720459
ELBO at 5000: -1.0437061786651611
ELBO at 6000: 0.0048046112060546875
ELBO at 7000: 0.37055855989456177
ELBO at 8000: 2.766662120819092
ELBO at 9000: -4.042356967926025
ELBO at 9999: -0.4423530101776123<jupyter_text>We again visualize some predictions.<jupyter_code>idxs = [0, 2, 5, 7]

_, axes = plt.subplots(figsize=(10, 6), nrows=2, ncols=2)
for _, (idx, ax) in enumerate(zip(idxs, axes.flatten())):
    x = np.squeeze(x_target[idx, :, :])
    f = np.squeeze(f_target[idx, :, :])
    y = np.squeeze(y_target[idx, :, :])

    srt_idxs = np.argsort(x)
    ax.plot(x[srt_idxs], f[srt_idxs], color="blue", alpha=0.75)
    ax.scatter(
        x[sample_idxs[:n_context]], y[sample_idxs[:n_context]], 
        color="blue", marker="+", alpha=0.75
    )

    for i in range(20):
        key, apply_key = random.split(key, 2)
        y_star = attentive_neural_process.apply(
            params=params,
            rng=apply_key,
            x_context=x[np.newaxis, sample_idxs, np.newaxis],
            y_context=y[np.newaxis, sample_idxs, np.newaxis],
            x_target=x_target[[idx], :, :]
        ).mean
        x_star = np.squeeze(x_target[[idx], :, :])
        y_star = np.squeeze(y_star)
        ax.plot(x_star[srt_idxs], y_star[srt_idxs], color="black", alpha=0.1)
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1084042429187551, 'fraction_numerical': 0.06457629094300035, 'mean_word_length': 4.131578947368421, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '323810', 'license': 'Apache-2.0', 'n_tokens_mistral': 3429, 'n_tokens_neox': 2993, 'n_words': 794, 'path': 'docs/source/notebooks/neural_process.ipynb', 'repo_name': 'dirmeier/ramsey', 'size': '8579'}