Illustration of a basic bi-vae implementation which is essentially 2 VAEs glued together to share the same latent space

bivae.py

import equinox as eqx
import equinox.nn as nn
import jax
import jax.numpy as jnp
import optax
import matplotlib.pyplot as plt

#VAE parameters
H = 400 #Hidden network size
L = 2 #Latent space size

#Data parameters
MEANS = (jnp.array([-2,2]),jnp.array([2,-2]))
VARIANCES = [jnp.eye(2)]*2
WEIGHTS = [0.5, 0.5]

A_N = 2000
B_N = 2000
PAIR_N = 0

#Training Params
LR = 0.003

def gaussian_mixture(key, means, covariances, weights, n_samples): 
    index_key, gaussain_key =  jax.random.split(key)
    
    samples = []
    for i in range(len(means)):
        gaussain_key,sub_key = jax.random.split(gaussain_key)
        m = means[i]
        c = covariances[i]
        samples.append(jax.random.multivariate_normal(sub_key, m, c, (n_samples,)))

    indicies = jax.random.randint(index_key, shape=(n_samples,), minval=0, maxval=len(means))
    output = jnp.zeros((n_samples,len(means[0])))
    for i in range(n_samples):
        output = output.at[i].set(samples[indicies[i]][i])
        
    return output, indicies


class VAEComponent(eqx.Module): 
    #Maps from a continuous input variable to a gaussian distribution
    input_layer: eqx.nn.Linear
    mean_output: eqx.nn.Linear 
    log_var_output: eqx.nn.Linear
    def __init__(self, n_input, n_hidden, n_output, key):
        a,b,c = jax.random.split(key, 3) 
        self.input_layer = eqx.nn.Linear(n_input, n_hidden,key=a)
        self.mean_output = eqx.nn.Linear(n_hidden, n_output, key=b)
        self.log_var_output = eqx.nn.Linear(n_hidden, n_output, key=c)

    def __call__(self,x):
        h = jax.nn.leaky_relu(self.input_layer(x))
        mean = self.mean_output(h)
        log_var = self.log_var_output(h)
        return mean, log_var

#Gaussian VAE primitives
def gaussian_kl_divergence(p, q):
    p_mean, p_log_var = p
    q_mean, q_log_var = q

    kl_div = (q_log_var-p_log_var + (jnp.exp(p_log_var)+(p_mean-q_mean)**2)/jnp.exp(q_log_var)-1)/2
    return kl_div

def gaussian_log_probabilty(p, x):
    p_mean, p_log_var = p
    log_p = (-1/2)*((x-p_mean)**2/jnp.exp(p_log_var))-p_log_var/2-jnp.log(jnp.sqrt(2*jnp.pi))
    return log_p

def sample_gaussian(p, key):
    p_mean, p_log_var = p
    samples = jax.random.normal(key,shape=p_mean.shape)*jnp.exp(p_log_var/2)+p_mean
    return samples

def concat_probabilties(p_a, p_b):
    mean = jnp.concatenate([p_a[0],p_b[0]], axis=1)
    log_var = jnp.concatenate([p_a[1],p_b[1]], axis=1)
    return (mean,log_var)

@jax.jit
def bivae_loss(bivae, data, key, prt=False):
    #In an actual implentation you wouldn't do all of this using full batch updates, 
    #you would compute something with the same expectation using SGD.
    decoders, encoders = bivae
    data_a, data_b, data_pair = data
    decoder_a, decoder_b = decoders 
    encoder_a, encoder_b, encoder_pair = encoders

    #Generate latent q distributions in z space
    q_a = jax.vmap(encoder_a)(data_a)
    q_b = jax.vmap(encoder_b)(data_b)
    q_pair = jax.vmap(encoder_pair)(data_pair)

    key_a, key_b, key_pair = jax.random.split(key, 3)

    #Sample Z values
    z_a = sample_gaussian(q_a,key_a)
    z_b = sample_gaussian(q_b,key_b)
    z_pair = sample_gaussian(q_pair,key_pair)

    #Compute kl_loss terms
    z_prior = (0,0)
    kl_a = gaussian_kl_divergence(q_a,z_prior)
    kl_b = gaussian_kl_divergence(q_b,z_prior)
    kl_pair = gaussian_kl_divergence(q_pair,z_prior)

    #Ground truth predictions
    p_a = jax.vmap(decoder_a)(z_a)
    p_b = jax.vmap(decoder_b)(z_b)
    p_pair = concat_probabilties(jax.vmap(decoder_a)(z_pair),jax.vmap(decoder_b)(z_pair))

    #Compute the probablity of the data given the latent sample
    log_p_a = gaussian_log_probabilty(p_a,data_a)
    log_p_b = gaussian_log_probabilty(p_b,data_b)
    log_p_pair = gaussian_log_probabilty(p_pair,data_pair)

    #Maximise p assigned to data, minimize KL div
    loss = sum(map(jnp.sum,[-log_p_a, -log_p_b, -log_p_pair, kl_a, kl_b, kl_pair]))
    if prt:
        #print("KL_DIV:",sum(map(jnp.sum,[kl_a, kl_b, kl_pair])))
        pass

    return loss

def update_state(state,data):
    model, opt_state = state

def make_bivae(a, b, h, l, key):
    keys = jax.random.split(key,5)
    decoder_a = VAEComponent(l, h, a, keys[0])
    decoder_b = VAEComponent(l, h, b, keys[1])
    
    encoder_a = VAEComponent(a, h, l, keys[2])
    encoder_b = VAEComponent(b, h, l, keys[3])
    encoder_pair = VAEComponent(a+b, h, l, keys[4])
    

    decoders = decoder_a, decoder_b
    encoders = encoder_a, encoder_b, encoder_pair
    bivae = decoders, encoders
    return bivae

def make_data(means, variances, weights, a_n, b_n, pair_n, key):
    a_key, b_key, pair_key = jax.random.split(key,3)
    #plt.scatter(data[:,0],data[:,1], c=idx)
    #plt.show()
    a_data,a_idx = gaussian_mixture(a_key, means, variances, weights, a_n)
    b_data,b_idx = gaussian_mixture(b_key, means, variances, weights, b_n)
    pair_data,pair_idx = gaussian_mixture(pair_key, means, variances, weights, pair_n)

    data = a_data[:,0:1], b_data[:,1:2], pair_data
    idx = a_idx, b_idx, pair_idx
    #plt.hist(data[0][:,0])
    #plt.show()
    #plt.hist(data[1][:,0])
    #plt.show()
    return data, idx

def show_data(data, idx):
    plt.scatter(data[:,0],data[:,1], c=idx)
    plt.show()

def update_state(state, data, optimizer, loss_fn):
    model, opt_state, key = state
    new_key, subkey = jax.random.split(key)


    #loss_fn(model, data, subkey, prt=True)
    
    loss, grads = jax.value_and_grad(loss_fn)(model, data, subkey)

    updates, new_opt_state = optimizer.update(grads, opt_state)
    new_model = eqx.apply_updates(model, updates)
    new_state = new_model, new_opt_state, new_key
    
    return loss,new_state

def sample_bivae(l, n_samples, bivae, key):
    z_key, x_key = jax.random.split(key)
    decoders = bivae[0]
    p_z = (jnp.zeros((n_samples,l)),)*2
    z = sample_gaussian(p_z, z_key)
    p_x = concat_probabilties(jax.vmap(decoders[0])(z),jax.vmap(decoders[1])(z))
    x = sample_gaussian(p_x, x_key)
    return x


def main():
    key = jax.random.PRNGKey(42)
    data_key, init_key, state_key, sample_key = jax.random.split(key,4)

    data, idx = make_data(MEANS, VARIANCES, WEIGHTS, A_N, B_N, PAIR_N, data_key)
    bivae = make_bivae( 1, 1, H, L, init_key)
    plt.scatter(data[0][:,0],data[1][:,0])
    plt.show()


    optimizer = optax.adam(LR)
    opt_state = optimizer.init(bivae)

    state = bivae, opt_state, state_key

    for i in range(1500):
        loss,state = update_state(state, data, optimizer, bivae_loss)
        print(i,loss)

    trained_bivae = state[0]

    n_samples = 1500
    samples = sample_bivae(L, n_samples,  trained_bivae, sample_key)
    plt.scatter(samples[:,0],samples[:,1])
    plt.show()
    plt.hist(samples[:,0])
    plt.show()
    plt.hist(samples[:,1])
    plt.show()


if __name__ == "__main__":
    main()