def __init__(self,

spatial_width,

n_colors,

trajectory_length=1000,

n_temporal_basis=10,

n_hidden_dense_lower=500,

n_hidden_dense_lower_output=2,

n_hidden_dense_upper=20,

n_hidden_conv=20,

n_layers_conv=4,

n_layers_dense_lower=4,

n_layers_dense_upper=2,

n_t_per_minibatch=1,

n_scales=1,

step1_beta=0.001,

uniform_noise = 0,

):

"""

Implements the objective function and mu and sigma estimators for a Gaussian diffusion

probabilistic model, as described in the paper:

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, Surya Ganguli

International Conference on Machine Learning, 2015

Parameters are as follow:

spatial_width - Spatial_width of training images

n_colors - Number of color channels in training data.

trajectory_length - The number of time steps in the trajectory.

n_temporal_basis - The number of temporal basis functions to capture time-step

dependence of model.

n_hidden_dense_lower - The number of hidden units in each layer of the dense network

in the lower half of the MLP. Set to 0 to make a convolutional-only lower half.

n_hidden_dense_lower_output - The number of outputs *per pixel* from the dense network

in the lower half of the MLP. Total outputs are

n_hidden_dense_lower_output*spatial_width**2.

n_hidden_dense_upper - The number of hidden units per pixel in the upper half of the MLP.

n_hidden_conv - The number of feature layers in the convolutional layers in the lower half

of the MLP.

n_layers_conv - How many convolutional layers to use in the lower half of the MLP.

n_layers_dense_lower - How many dense layers to use in the lower half of the MLP.

n_layers_dense_upper - How many dense layers to use in the upper half of the MLP.

n_t_per_minibatch - When computing objective, how many random time-steps t to evaluate

each minibatch at.

step1_beta - The lower bound on the noise variance of the first diffusion step. This is

the minimum variance of the learned model.

uniform_noise - Add uniform noise between [-uniform_noise/2, uniform_noise/2] to the input.

"""

super(DiffusionModel, self).__init__()

self.n_t_per_minibatch = n_t_per_minibatch

self.spatial_width = np.int16(spatial_width)

self.n_colors = np.int16(n_colors)

self.n_temporal_basis = n_temporal_basis

self.trajectory_length = trajectory_length

self.uniform_noise = uniform_noise

self.mlp = regression.MLP_conv_dense(

n_layers_conv, n_layers_dense_lower, n_layers_dense_upper,

n_hidden_conv, n_hidden_dense_lower, n_hidden_dense_lower_output, n_hidden_dense_upper,

spatial_width, n_colors, n_scales, n_temporal_basis)

self.temporal_basis = self.generate_temporal_basis(trajectory_length, n_temporal_basis)

self.beta_arr = self.generate_beta_arr(step1_beta)

self.children = [self.mlp]

def generate_beta_arr(self, step1_beta):

"""

Generate the noise covariances, beta_t, for the forward trajectory.

"""

# lower bound on beta

min_beta_val = 1e-6

min_beta_values = np.ones((self.trajectory_length,))*min_beta_val

min_beta_values[0] += step1_beta

min_beta = theano.shared(value=min_beta_values.astype(theano.config.floatX),

name='min beta')

# (potentially learned) function for how beta changes with timestep

# TODO add beta_perturb_coefficients to the parameters to be learned

beta_perturb_coefficients_values = np.zeros((self.n_temporal_basis,))

beta_perturb_coefficients = theano.shared(

value=beta_perturb_coefficients_values.astype(theano.config.floatX),

name='beta perturb coefficients')

beta_perturb = T.dot(self.temporal_basis.T, beta_perturb_coefficients)

# baseline behavior of beta with time -- destroy a constant fraction

# of the original data variance each time step

# NOTE 2 below means a fraction ~1/T of the variance will be left at the end of the

# trajectory

beta_baseline = 1./np.linspace(self.trajectory_length, 2., self.trajectory_length)

beta_baseline_offset = util.logit_np(beta_baseline).astype(theano.config.floatX)

# and the actual beta_t, restricted to be between min_beta and 1-[small value]

beta_arr = T.nnet.sigmoid(beta_perturb + beta_baseline_offset)

beta_arr = min_beta + beta_arr * (1 - min_beta - 1e-5)

beta_arr = beta_arr.reshape((self.trajectory_length, 1))

return beta_arr

def get_t_weights(self, t):

"""

Generate vector of weights allowing selection of current timestep.

(if t is not an integer, the weights will linearly interpolate)

"""

n_seg = self.trajectory_length

t_compare = T.arange(n_seg, dtype=theano.config.floatX).reshape((1,n_seg))

diff = abs(T.addbroadcast(t,1) - T.addbroadcast(t_compare,0))

t_weights = T.max(T.join(1, (-diff+1).reshape((n_seg,1)), T.zeros((n_seg,1))), axis=1)

return t_weights.reshape((-1,1))

def get_beta_forward(self, t):

"""

Get the covariance of the forward diffusion process at timestep

t.

"""

t_weights = self.get_t_weights(t)

return T.dot(t_weights.T, self.beta_arr)

def get_mu_sigma(self, X_noisy, t):

"""

Generate mu and sigma for one step in the reverse trajectory,

starting from a minibatch of images X_noisy, and at timestep t.

"""

Z = self.mlp.apply(X_noisy)

mu_coeff, beta_coeff = self.temporal_readout(Z, t)

# reverse variance is perturbation around forward variance

beta_forward = self.get_beta_forward(t)

# make impact of beta_coeff scaled appropriately with mu_coeff

beta_coeff_scaled = beta_coeff / np.sqrt(self.trajectory_length).astype(theano.config.floatX)

beta_reverse = T.nnet.sigmoid(beta_coeff_scaled + util.logit(beta_forward))

# # reverse mean is decay towards mu_coeff

# mu = (X_noisy - mu_coeff)*T.sqrt(1. - beta_reverse) + mu_coeff

# reverse mean is a perturbation around the mean under forward

# process

# # DEBUG -- use these lines to test objective is 0 for isotropic Gaussian model

# beta_reverse = beta_forward

# mu_coeff = mu_coeff*0

mu = X_noisy*T.sqrt(1. - beta_forward) + mu_coeff*T.sqrt(beta_forward)

sigma = T.sqrt(beta_reverse)

mu.name = 'mu p'

sigma.name = 'sigma p'

return mu, sigma

def generate_forward_diffusion_sample(self, X_noiseless):

"""

Corrupt a training image with t steps worth of Gaussian noise, and

return the corrupted image, as well as the mean and covariance of the

posterior q(x^{t-1}|x^t, x^0).

"""

X_noiseless = X_noiseless.reshape(

(-1, self.n_colors, self.spatial_width, self.spatial_width))

n_images = X_noiseless.shape[0].astype('int16')

rng = Random().theano_rng

# choose a timestep in [1, self.trajectory_length-1].

# note the reverse process is fixed for the very

# first timestep, so we skip it.

# TODO for some reason random_integer is missing from the Blocks

# theano random number generator.

t = T.floor(rng.uniform(size=(1,1), low=1, high=self.trajectory_length,

dtype=theano.config.floatX))

t_weights = self.get_t_weights(t)

N = rng.normal(size=(n_images, self.n_colors, self.spatial_width, self.spatial_width),

dtype=theano.config.floatX)

# noise added this time step

beta_forward = self.get_beta_forward(t)

# decay in noise variance due to original signal this step

alpha_forward = 1. - beta_forward

# compute total decay in the fraction of the variance due to X_noiseless

alpha_arr = 1. - self.beta_arr

alpha_cum_forward_arr = T.extra_ops.cumprod(alpha_arr).reshape((self.trajectory_length,1))

alpha_cum_forward = T.dot(t_weights.T, alpha_cum_forward_arr)

# total fraction of the variance due to noise being mixed in

beta_cumulative = 1. - alpha_cum_forward

# total fraction of the variance due to noise being mixed in one step ago

beta_cumulative_prior_step = 1. - alpha_cum_forward/alpha_forward

# generate the corrupted training data

X_uniformnoise = X_noiseless + (rng.uniform(size=(n_images, self.n_colors, self.spatial_width, self.spatial_width),

dtype=theano.config.floatX)-T.constant(0.5,dtype=theano.config.floatX))*T.constant(self.uniform_noise,dtype=theano.config.floatX)

X_noisy = X_uniformnoise*T.sqrt(alpha_cum_forward) + N*T.sqrt(1. - alpha_cum_forward)

# compute the mean and covariance of the posterior distribution

mu1_scl = T.sqrt(alpha_cum_forward / alpha_forward)

mu2_scl = 1. / T.sqrt(alpha_forward)

cov1 = 1. - alpha_cum_forward/alpha_forward

cov2 = beta_forward / alpha_forward

lam = 1./cov1 + 1./cov2

mu = (

X_uniformnoise * mu1_scl / cov1 +

X_noisy * mu2_scl / cov2

) / lam

sigma = T.sqrt(1./lam)

sigma = sigma.reshape((1,1,1,1))

mu.name = 'mu q posterior'

sigma.name = 'sigma q posterior'

X_noisy.name = 'X_noisy'

t.name = 't'

return X_noisy, t, mu, sigma

def get_beta_full_trajectory(self):

"""

Return the cumulative covariance from the entire forward trajectory.

"""

alpha_arr = 1. - self.beta_arr

beta_full_trajectory = 1. - T.exp(T.sum(T.log(alpha_arr)))

return beta_full_trajectory

def get_negL_bound(self, mu, sigma, mu_posterior, sigma_posterior):

"""

Compute the lower bound on the log likelihood, as a function of mu and

sigma from the reverse diffusion process, and the posterior mu and

sigma from the forward diffusion process.

Returns the difference between this bound and the log likelihood

under a unit norm isotropic Gaussian. So this function returns how

much better the diffusion model is than an isotropic Gaussian.

"""

# the KL divergence between model transition and posterior from data

KL = ( T.log(sigma) - T.log(sigma_posterior)

+ (sigma_posterior**2 + (mu_posterior-mu)**2)/(2*sigma**2)

- 0.5)

# conditional entropies H_q(x^T|x^0) and H_q(x^1|x^0)

H_startpoint = (0.5*(1 + np.log(2.*np.pi))).astype(theano.config.floatX) + 0.5*T.log(self.beta_arr[0])

H_endpoint = (0.5*(1 + np.log(2.*np.pi))).astype(theano.config.floatX) + 0.5*T.log(self.get_beta_full_trajectory())

H_prior = (0.5*(1 + np.log(2.*np.pi))).astype(theano.config.floatX) + 0.5*T.log(1.)

negL_bound = KL*self.trajectory_length + H_startpoint - H_endpoint + H_prior

# the negL_bound if this was an isotropic Gaussian model of the data

negL_gauss = (0.5*(1 + np.log(2.*np.pi))).astype(theano.config.floatX) + 0.5*T.log(1.)

negL_diff = negL_bound - negL_gauss

L_diff_bits = negL_diff / T.log(2.)

L_diff_bits_avg = L_diff_bits.mean()*self.n_colors

return L_diff_bits_avg

def cost_single_t(self, X_noiseless):

"""

Compute the lower bound on the log likelihood, given a training minibatch, for a single

randomly chosen timestep.

"""

X_noisy, t, mu_posterior, sigma_posterior = \

self.generate_forward_diffusion_sample(X_noiseless)

mu, sigma = self.get_mu_sigma(X_noisy, t)

negL_bound = self.get_negL_bound(mu, sigma, mu_posterior, sigma_posterior)

return negL_bound

def internal_state(self, X_noiseless):

"""

Return a bunch of the internal state, for monitoring purposes during optimization.

"""

X_noisy, t, mu_posterior, sigma_posterior = \

self.generate_forward_diffusion_sample(X_noiseless)

mu, sigma = self.get_mu_sigma(X_noisy, t)

mu_diff = mu-mu_posterior

mu_diff.name = 'mu diff'

logratio = T.log(sigma/sigma_posterior)

logratio.name = 'log sigma ratio'

return [mu_diff, logratio, mu, sigma, mu_posterior, sigma_posterior, X_noiseless, X_noisy]

@application

def cost(self, X_noiseless):

"""

Compute the lower bound on the log likelihood, given a training minibatch.

This will draw a single timestep and compute the cost for that timestep only.

"""

cost = 0.

for ii in xrange(self.n_t_per_minibatch):

cost += self.cost_single_t(X_noiseless)

return cost/self.n_t_per_minibatch

def temporal_readout(self, Z, t):

"""

Go from the top layer of the multilayer perceptron to coefficients for

mu and sigma for each pixel.

Z contains coefficients for spatial basis functions for each pixel for

both mu and sigma.

"""

n_images = Z.shape[0].astype('int16')

t_weights = self.get_t_weights(t)

Z = Z.reshape((n_images, self.spatial_width, self.spatial_width,

self.n_colors, 2, self.n_temporal_basis))

coeff_weights = T.dot(self.temporal_basis, t_weights)

concat_coeffs = T.dot(Z, coeff_weights)

mu_coeff = concat_coeffs[:,:,:,:,0].dimshuffle(0,3,1,2)

beta_coeff = concat_coeffs[:,:,:,:,1].dimshuffle(0,3,1,2)

return mu_coeff, beta_coeff

def generate_temporal_basis(self, trajectory_length, n_basis):

"""

Generate the bump basis functions for temporal readout of mu and sigma.

"""

temporal_basis = np.zeros((trajectory_length, n_basis))

xx = np.linspace(-1, 1, trajectory_length)

x_centers = np.linspace(-1, 1, n_basis)

width = (x_centers[1] - x_centers[0])/2.

for ii in xrange(n_basis):

temporal_basis[:,ii] = np.exp(-(xx-x_centers[ii])**2 / (2*width**2))

temporal_basis /= np.sum(temporal_basis, axis=1).reshape((-1,1))

temporal_basis = temporal_basis.T

temporal_basis_theano = theano.shared(value=temporal_basis.astype(theano.config.floatX),

name="temporal basis")