<jupyter_start><jupyter_text>Implementation of Stein variational gradient descent (SVGD) in TensorFlow. And try to improve its performance in high-dimensional space.<jupyter_code>%matplotlib inline

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow.contrib.distributions as tfd
tf.reset_default_graph()
def expectation(f, xs, axis=0, name=None):
  """Expectation value of mapping function `f` to values `xs`.
  Args:
    f: Callable from tensor to tensor.
    xs: Tensor with the shape `[None]+fis`, where the `fis` is the
      shape of the input tensor of `f`.
  Returns:
    Tensor with the shape of the output tensor of `f`.
  """
  with tf.name_scope(name, 'Expectation', [xs]):
    fs = tf.map_fn(f, xs, parallel_iterations=100,
                   back_prop=False)
    return tf.reduce_mean(fs, axis)

def square_distance(x, y):
  return tf.reduce_sum(tf.square(x - y))

def rbf_kernel(x, y, h, name=None):
  """Implements the RBF kernel.
  Args:
    x: Tensor.
    y: Tensor with the same shape as the `x`.
    h: Non-vanishing scalar of float type.
  Returns:
    Scalar.
  """
  with tf.name_scope(name, 'RBFKernel', [x, y]):
    return tf.exp(-0.5/h * square_distance(x, y))

def grad_rbf_kernel(x, y, h, name=None):
  """Implements the gradients of the RBF kernel.
  Args:
    x: Tensor.
    y: Tensor with the same shape as the `x`.
    h: Non-vanishing scalar of float type.
  Returns:
    Tensor with the same shape as the `x`.
  """
  with tf.name_scope(name, 'RBFKernelGradient', [x, y]):
    return rbf_kernel(x, y, h) * (y - x) / h

def get_phi_star(samples, log_p, kernel, grad_kernel):
  """Get the function :math:`\phi_{\star}`(\cdot)` that maps
  from tensor of an arbitrary shape (say `x_shape`) to a tensor
  of the same shape.
  
  Args:
    samples: Tensor of the shape `[n_samples] + x_shape`.
    log_p: Callable that maps from tensor of `x_shape` to scalar.
    kernel: Callable that maps two tensors of `x_shape` to scaler.
    grad_kernel: The gradient of the `kernel`.
  Returns:
    Callable.
  """
  def phi_star(x):
    def adaptive_h():
      median_square = tf.reduce_mean(
          square_distance(tf.expand_dims(x, axis=0), samples))
      n_samples = samples.get_shape().as_list()[0]
      return median_square / tf.log(float(n_samples))
    def integrand(y):
      return (kernel(x, y, adaptive_h()) * tf.gradients(log_p(y), [y])[0]
              + grad_kernel(x, y, adaptive_h()))
    return expectation(integrand, samples)
  return phi_star

def median(value):
  value = tf.reshape(value, [-1])
  median_idx = value.get_shape().as_list()[0]//2
  return tf.nn.top_k(value, median_idx).values[median_idx-1]
def initialize_samples(n_samples, n_dims):
  samples_shape = [n_samples, n_dims]
  initializer = tf.initializers.random_uniform(minval=-1, maxval=1)
  samples = tf.get_variable(
      'samples', shape=samples_shape,
      dtype='float32', initializer=initializer)
  return samples

n_samples = 200
n_dims = 1
xs = initialize_samples(n_samples, n_dims)
dist = tfd.MultivariateNormalDiag(tf.zeros([n_dims]),
                                  tf.ones([n_dims]))
log_p = dist.log_prob
phi_star = get_phi_star(xs, log_p, rbf_kernel, grad_rbf_kernel)
class AdaGrad(object):
  def __init__(self, xs, phi_star, learning_rate, alpha=0.9):
    self.xs = xs
    self.phi_star = phi_star
    self.learning_rate = learning_rate
    self.alpha = alpha
    self.cache = tf.Variable(tf.zeros(shape=xs.get_shape(),
                                      dtype='float32'),
                             name='AdaGradCache')
  
  def update(self):
    raw_grads = tf.map_fn(self.phi_star, self.xs, name='Gradients')
    new_cache = (self.alpha * self.cache
                 + (1 - self.alpha) * tf.square(raw_grads))
    update_cache_op = tf.assign(self.cache, new_cache, name='UpdateCache')
    with tf.control_dependencies([update_cache_op]):
      adjusted_grads = raw_grads / (tf.sqrt(self.cache) + 1e-8)
      delta_xs = self.learning_rate * adjusted_grads
      update_op = tf.assign(self.xs, self.xs + delta_xs, name='UpdateVars')
      return update_op
    
update_op = AdaGrad(xs, phi_star, learning_rate=0.1).update()
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
%%time

n_iters = 100
for step in range(n_iters):
    sess.run(update_op)
    if step % 1 == 0:
        print('After step {}'.format(step))
        axis=0
        xs_val = sess.run(xs)
        targets = np.random.normal(size=n_samples)
        plt.hist([xs_val[:,axis], targets], bins=20)
        plt.show()
xs_val = sess.run(xs)
axis=0
targets = np.random.normal(size=n_samples)
plt.hist([xs_val[:,axis], targets], bins=10)
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12679324894514768, 'fraction_numerical': 0.008649789029535865, 'mean_word_length': 3.3455545371219064, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '460364', 'license': 'MIT', 'n_tokens_mistral': 1685, 'n_tokens_neox': 1648, 'n_words': 423, 'path': 'nn4infer/svgd.ipynb', 'repo_name': 'physicsliupeng/nn4fun', 'size': '4740'}