<jupyter_start><jupyter_text>Rotational Invariance of L2HMCThis notebook examines how well an l2hmc sampler trained on one strongly correlated gaussian distribution performs on rotated versions of that same distribution   Setup<jupyter_code>from google.colab import drive
drive.mount('/content/drive', force_remount=True)
import sys
sys.path.append('/content/drive/MyDrive/repos/l2hmc')
!cat /proc/meminfo
!cat /proc/cpuinfo
gpu_info = !nvidia-smi
gpu_info = '\n'.join(gpu_info)
if gpu_info.find('failed') >= 0:
  print('Select the Runtime > "Change runtime type" menu to enable a GPU accelerator, ')
  print('and then re-execute this cell.')
else:
  print(gpu_info)
%tensorflow_version 1.x
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import time
%matplotlib inline
%load_ext autoreload
%autoreload 2
from utils.func_utils import accept, jacobian, autocovariance, get_log_likelihood, \
    get_data, binarize, normal_kl, acl_spectrum, ESS
from utils.distributions import Gaussian, GMM, GaussianFunnel, gen_ring
from utils.layers import Linear, Sequential, Zip, Parallel, ScaleTanh
from utils.dynamics import Dynamics
from utils.sampler import propose
from utils.notebook_utils import get_hmc_samples<jupyter_output><empty_output><jupyter_text>Network architectureWe define the network architecture for our L2HMC network. We first embed the first two variables ($\{x, \partial_x U\}$ or $\{v, x_m\}$) as well as the time, formatted as $\tau(t) = \left(\cos(\frac{2\pi t}{M}), \sin(\frac{2\pi t}{M})\right)$. They are forwarded through an MLP and then produce $S, T$ and $Q$.<jupyter_code>def network(x_dim, width, scope, factor):
    with tf.variable_scope(scope):
        net = Sequential([
            Zip([
                Linear(x_dim, width, scope='embed_1', factor=1.0 / 3),
                Linear(x_dim, width, scope='embed_2', factor=factor * 1.0 / 3),
                Linear(2, width, scope='embed_3', factor=1.0 / 3),
                lambda _: 0.,
            ]),
            sum,
            tf.nn.relu,
            Linear(width, width, scope='linear_1'),
            tf.nn.relu,
            Parallel([
                Sequential([
                    Linear(width, x_dim, scope='linear_s', factor=0.001), 
                    ScaleTanh(x_dim, scope='scale_s')
                ]),
                Linear(width, x_dim, scope='linear_t', factor=0.001),
                Sequential([
                    Linear(width, x_dim, scope='linear_f', factor=0.001),
                    ScaleTanh(x_dim, scope='scale_f'),
                ])
            ])  
        ])
        
    return net<jupyter_output><empty_output><jupyter_text>DistributionWe define our energy function. It is a Gaussian distribution with zero mean. The covariance is a $\pi/4$ rotation of the eigenvalues $[100, 10^{-1}]$. We set up our dynamics which take as input our energy function, the number of time step of our operator, the (learnable) step-size and our architecture.<jupyter_code>x_dim = 2
network_width = 10
degree_rotation = 2
num_rotations = 5

mu = np.zeros(2,)
cov = np.array([[50.05, -49.95], [-49.95, 50.05]])
theta = np.radians(degree_rotation)
R = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])

# Applying rotations to the original covariance matrix
covs = []
for i in range(num_rotations):
  covs.append(cov)
  cov = R@cov@R.T

distribution = Gaussian(mu,covs[0])
distribution2 = Gaussian(mu,covs[1])
distribution3 = Gaussian(mu,covs[2])
distribution4 = Gaussian(mu,covs[3])
distribution5 = Gaussian(mu,covs[4])
dynamics = Dynamics(x_dim, distribution.get_energy_function(), T=10, eps=0.1, width=network_width, net_factory=network)<jupyter_output>WARNING:tensorflow:From /content/drive/MyDrive/repos/l2hmc/utils/distributions.py:45: The name tf.placeholder is deprecated. Please use tf.compat.v1.placeholder instead.

WARNING:tensorflow:From /content/drive/MyDrive/repos/l2hmc/utils/dynamics.py:52: The name tf.get_variable is deprecated. Please use tf.compat.v1.get_variable instead.

WARNING:tensorflow:From /content/drive/MyDrive/repos/l2hmc/utils/dynamics.py:54: The name tf.log is deprecated. Please use tf.math.log instead.

WARNING:tensorflow:
The TensorFlow contrib module will not be included in TensorFlow 2.0.
For more information, please see:
  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md
  * https://github.com/tensorflow/addons
  * https://github.com/tensorflow/io (for I/O related ops)
If you depend on functionality not listed there, please file an issue.

WARNING:tensorflow:From /content/drive/MyDrive/repos/l2hmc/utils/layers.py:34: calling Constant.__init__ (from tensorflow.python.ops[...]<jupyter_text>We can directly sample from this distribution and plot it for sanity-check.<jupyter_code>S1 = distribution.get_samples(200)
S2 = distribution2.get_samples(200)
S3 = distribution3.get_samples(200)
S4 = distribution4.get_samples(200)
S5 = distribution5.get_samples(200)
plt.title('Strongly Correlated Gaussians (SCG)')
plt.scatter(S1[:, 0], S1[:, 1])
plt.scatter(S2[:, 0], S2[:, 1])
plt.scatter(S3[:, 0], S3[:, 1])
plt.scatter(S4[:, 0], S4[:, 1])
plt.scatter(S5[:, 0], S5[:, 1])
plt.show()<jupyter_output><empty_output><jupyter_text>We set up the loss on both $p(\xi)$ (here `x`) and $q(\xi)$ (here `z`). We then train with Adam with a learning rate of $10^{-3}$.<jupyter_code>x = tf.placeholder(tf.float32, shape=(None, x_dim))
z = tf.random_normal(tf.shape(x))

Lx, _, px, output = propose(x, dynamics, do_mh_step=True)
Lz, _, pz, _ = propose(z, dynamics, do_mh_step=False)

loss = 0.

v1 = (tf.reduce_sum(tf.square(x - Lx), axis=1) * px) + 1e-4
v2 = (tf.reduce_sum(tf.square(z - Lz), axis=1) * pz) + 1e-4
scale = 0.1

loss += scale * (tf.reduce_mean(1.0 / v1) + tf.reduce_mean(1.0 / v2))
loss += (- tf.reduce_mean(v1) - tf.reduce_mean(v2)) / scale
print(loss)
global_step = tf.Variable(0., name='global_step', trainable=False)
learning_rate = tf.train.exponential_decay(1e-3, global_step, 1000, 0.96, staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss, global_step=global_step)<jupyter_output><empty_output><jupyter_text>The training loop described in Algorithm $1$.<jupyter_code>all_placeholders = [placeholder for op in tf.get_default_graph().get_operations() if op.type=='Placeholder' for placeholder in op.values()]
print(all_placeholders)
n_steps = 5000
n_samples = 200

loss_hist = []
acceptance_hist = []

samples = np.random.randn(n_samples, x_dim)
curr_cov = covs[0]

sess = tf.Session()
sess.run(tf.global_variables_initializer())

start_time = time.time()
for t in range(n_steps):
    _, loss_, samples, px_, lr_ = sess.run([
        train_op,
        loss,
        output[0],
        px,
        learning_rate,
    ], feed_dict={x: samples, 'sigma_var:0' : curr_cov})

    if t % 100 == 0:
        print('Step: %d / %d, Loss: %.2e, Acceptance sample: %.2f, LR: %.5f' % (t, n_steps, loss_, np.mean(px_), lr_))
        loss_hist.append(loss_)
        acceptance_hist.append(np.mean(px_))


training_time = time.time() - start_time<jupyter_output>Step: 0 / 5000, Loss: -6.78e+01, Acceptance sample: 0.85, LR: 0.00100
Step: 100 / 5000, Loss: -3.07e+02, Acceptance sample: 0.57, LR: 0.00100
Step: 200 / 5000, Loss: -3.31e+02, Acceptance sample: 0.46, LR: 0.00100
Step: 300 / 5000, Loss: -5.30e+02, Acceptance sample: 0.42, LR: 0.00100
Step: 400 / 5000, Loss: -1.41e+03, Acceptance sample: 0.27, LR: 0.00100
Step: 500 / 5000, Loss: -1.54e+03, Acceptance sample: 0.19, LR: 0.00100
Step: 600 / 5000, Loss: -3.30e+03, Acceptance sample: 0.23, LR: 0.00100
Step: 700 / 5000, Loss: -3.47e+03, Acceptance sample: 0.27, LR: 0.00100
Step: 800 / 5000, Loss: -3.39e+03, Acceptance sample: 0.26, LR: 0.00100
Step: 900 / 5000, Loss: -3.80e+03, Acceptance sample: 0.26, LR: 0.00100
Step: 1000 / 5000, Loss: -3.86e+03, Acceptance sample: 0.24, LR: 0.00096
Step: 1100 / 5000, Loss: -3.88e+03, Acceptance sample: 0.25, LR: 0.00096
Step: 1200 / 5000, Loss: -4.83e+03, Acceptance sample: 0.26, LR: 0.00096
Step: 1300 / 5000, Loss: -4.47e+03, Acceptance sample: 0.30, LR[...]<jupyter_text>After training, we generate $200$ chains for $2000$ steps for evaluation purposes. Testing on 0, 2, 6, 4, and 8 degree rotated distributions<jupyter_code>samples = distribution.get_samples(n=n_samples)
final_samples1 = []

for t in range(2000):
    final_samples1.append(np.copy(samples))

    feed_dict = {
        x: samples,
        'sigma_var:0' : covs[0]
    }
    samples = sess.run(output[0], feed_dict)

samples = distribution2.get_samples(n=n_samples)
final_samples2 = []

for t in range(2000):
    final_samples2.append(np.copy(samples))

    feed_dict = {
        x: samples,
        'sigma_var:0' : covs[1]
    }
    samples = sess.run(output[0], feed_dict)

samples = distribution3.get_samples(n=n_samples)
final_samples3 = []

for t in range(2000):
    final_samples3.append(np.copy(samples))

    feed_dict = {
        x: samples,
        'sigma_var:0' : covs[2]
    }
    samples = sess.run(output[0], feed_dict)

samples = distribution4.get_samples(n=n_samples)
final_samples4 = []

for t in range(2000):
    final_samples4.append(np.copy(samples))

    feed_dict = {
        x: samples,
        'sigma_var:0' : covs[3]
    }
    samples = sess.run(output[0], feed_dict)

samples = distribution5.get_samples(n=n_samples)
final_samples5 = []

for t in range(2000):
    final_samples5.append(np.copy(samples))

    feed_dict = {
        x: samples,
        'sigma_var:0' : covs[4]
    }
    samples = sess.run(output[0], feed_dict)<jupyter_output><empty_output><jupyter_text>We compute the HMC chains with **auto-correlation spectrums** as well.<jupyter_code>L2HMC_samples1 = np.array(final_samples1)
L2HMC_samples2 = np.array(final_samples2)
L2HMC_samples3 = np.array(final_samples3)
L2HMC_samples4 = np.array(final_samples4)
L2HMC_samples5 = np.array(final_samples5)
scale = np.sqrt(np.trace(cov))
L2HMC1 = acl_spectrum(L2HMC_samples1, scale=scale)
L2HMC2 = acl_spectrum(L2HMC_samples2, scale=scale)
L2HMC3 = acl_spectrum(L2HMC_samples3, scale=scale)
L2HMC4 = acl_spectrum(L2HMC_samples4, scale=scale)
L2HMC5 = acl_spectrum(L2HMC_samples5, scale=scale)<jupyter_output><empty_output><jupyter_text>Saving data from this run for later comparison:<jupyter_code>l2hmc_data={"problem_type": "SCG rotation invariance test",
            "network_type": "10 unit width",
            "training_time": training_time,
            "loss": loss_hist,
            "acceptance": acceptance_hist,
            "samples1": L2HMC_samples1,
            "samples2": L2HMC_samples2,
            "samples3": L2HMC_samples3,
            "samples4": L2HMC_samples4,
            "samples5": L2HMC_samples5,
            "acl1": L2HMC1,
            "acl2": L2HMC2,
            "acl3": L2HMC3,
            "acl4": L2HMC4,
            "acl5": L2HMC5,
            "ess1": ESS(L2HMC1),
            "ess2": ESS(L2HMC2),
            "ess3": ESS(L2HMC3),
            "ess4": ESS(L2HMC4),
            "ess5": ESS(L2HMC5)
            }

np.save("/content/drive/MyDrive/repos/l2hmc/data/rotation-invariance10.npy", l2hmc_data)<jupyter_output><empty_output><jupyter_text>We can plot auto-correlation.<jupyter_code>xaxis = 10 * np.arange(300)
plt.plot(xaxis, L2HMC1[:300], label='Trained Distribution')
plt.plot(xaxis, L2HMC2[:300], label='Rotated 2 degrees')
plt.plot(xaxis, L2HMC3[:300], label='Rotated 4 degrees')
plt.plot(xaxis, L2HMC4[:300], label='Rotated 6 degrees')
plt.plot(xaxis, L2HMC5[:300], label='Rotated 8 degrees')
plt.ylabel('Auto-correlation')
plt.xlabel('Gradient Computations')
plt.legend()
plt.show()<jupyter_output><empty_output><jupyter_text>We now compute the **Effective Sample Size** (ESS).<jupyter_code>print('ESS L2HMC Base: %.2e -- ESS L2HMC Rotated 2deg: %.2e -- ESS L2HMC Rotated 4deg: %.2e' % (ESS(L2HMC1), ESS(L2HMC2), ESS(L2HMC3) ))<jupyter_output>ESS L2HMC Base: 3.02e-01 -- ESS L2HMC Rotated 2deg: 1.67e-01 -- ESS L2HMC Rotated 4deg: 9.10e-02<jupyter_text>We can visualize a single chain of L2HMC for $50$ time steps.<jupyter_code>plt.scatter(S1[:, 0], S1[:, 1])
plt.plot(L2HMC_samples1[:50, 2, 0], L2HMC_samples1[:50, 2, 1], color='orange', marker='o', alpha=0.8)
plt.show()
plt.scatter(S2[:, 0], S2[:, 1])
plt.plot(L2HMC_samples2[:50, 2, 0], L2HMC_samples2[:50, 2, 1], color='orange', marker='o', alpha=0.8)
plt.show()
plt.scatter(S3[:, 0], S3[:, 1])
plt.plot(L2HMC_samples3[:50, 2, 0], L2HMC_samples3[:50, 2, 1], color='orange', marker='o', alpha=0.8)
plt.show()
plt.scatter(S4[:, 0], S4[:, 1])
plt.plot(L2HMC_samples4[:50, 2, 0], L2HMC_samples4[:50, 2, 1], color='orange', marker='o', alpha=0.8)
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14461054287962236, 'fraction_numerical': 0.06129032258064516, 'mean_word_length': 4.026097271648873, 'pattern_counts': {'":': 20, '<': 46, '<?xml version=': 0, '>': 48, 'https://': 3, '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': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '104010', 'license': 'Apache-2.0', 'n_tokens_mistral': 5201, 'n_tokens_neox': 4670, 'n_words': 1069, 'path': 'rotation-invariance.ipynb', 'repo_name': 'evbecker/l2hmc', 'size': '12710'}