def __call__(self, x):
x_ad = self.adapter['in']['scale'] * x + self.adapter['in']['shift']
self.pre_adapt = self.mlp(x_ad)
mean = self.adapter['out'][
'scale'] * self.pre_adapt.mean + self.adapter['out']['shift']
cov = tf.reshape(self.adapter['out']['scale'], [-1, 1]) * tf.reshape(
self.adapter['out']['scale'], [1, -1]) * self.pre_adapt.var
return gv.GaussianVar(mean, cov)
def __call__(self, x):
A = [aa.value.mean for aa in self.A]
b = [bb.value.mean for bb in self.b]
self.h = [x]
self.h.append(tf.matmul(x, A[0]) + b[0])
for L in range(1, len(A)):
self.h.append(tf.matmul(tf.nn.relu(self.h[-1]), A[L]) + b[L])
return gv.GaussianVar(self.h[-1], tf.constant(0.0, dtype=tf.float32))
def linear(x, A, b):
"""
compute y = x^T A + b
"""
x_mean = x.mean
y_mean = tf.matmul(x_mean, A.mean) + b.mean
x_cov = x.var
y_cov = linear_covariance(x_mean, x_cov, A, b)
return gv.GaussianVar(y_mean, y_cov)
def linear_certain_activations(x_certain, A, b):
"""
compute y = x^T A + b
assuming x has zero variance
"""
x_mean = x_certain
xx = x_mean * x_mean
y_mean = tf.matmul(x_mean, A.mean) + b.mean
y_cov = tf.matrix_diag(tf.matmul(xx, A.var) + b.var)
return gv.GaussianVar(y_mean, y_cov)
def linear_relu_diagonal(x, A, b):
"""
compute y = relu(x)^T A + b
"""
x_var_diag = x.var
sqrt_x_var_diag = tf.sqrt(x_var_diag)
mu = x.mean / (sqrt_x_var_diag + EPSILON)
pdf = bu.standard_gaussian(mu)
cdf = bu.gaussian_cdf(mu)
softrelu = pdf + mu * cdf
z_mean = sqrt_x_var_diag * softrelu
y_mean = tf.matmul(z_mean, A.mean) + b.mean
z_var = x_var_diag * (cdf + mu * softrelu - tf.square(softrelu))
y_cov = linear_covariance_diagonal(z_mean, z_var, A, b)
return gv.GaussianVar(y_mean, y_cov)
def linear_heaviside(x, A, b):
"""
compute y = heaviside(x)^T A + b
"""
x_var_diag = tf.matrix_diag_part(x.var)
mu = x.mean / (tf.sqrt(x_var_diag) + EPSILON)
def heaviside_covariance(x):
mu1 = tf.expand_dims(mu, 2)
mu2 = tf.transpose(mu1, [0, 2, 1])
s11s22 = tf.expand_dims(x_var_diag, axis=2) * tf.expand_dims(
x_var_diag, axis=1)
rho = x.var / (tf.sqrt(s11s22)) # + EPSILON)
rho = tf.clip_by_value(rho, -1 / (1 + EPSILON), 1 / (1 + EPSILON))
return bu.heavy_g(rho, mu1, mu2)
z_mean = bu.gaussian_cdf(mu)
y_mean = tf.matmul(z_mean, A.mean) + b.mean
z_cov = heaviside_covariance(x)
y_cov = linear_covariance(z_mean, z_cov, A, b)
return gv.GaussianVar(y_mean, y_cov)
def linear_relu(x, A, b):
"""
compute y = relu(x)^T A + b
"""
x_var_diag = tf.matrix_diag_part(x.var)
sqrt_x_var_diag = tf.sqrt(x_var_diag)
mu = x.mean / (sqrt_x_var_diag + EPSILON)
def relu_covariance(x):
mu1 = tf.expand_dims(mu, 2)
mu2 = tf.transpose(mu1, [0, 2, 1])
s11s22 = tf.expand_dims(x_var_diag, axis=2) * tf.expand_dims(
x_var_diag, axis=1)
rho = x.var / (tf.sqrt(s11s22)) # + EPSILON)
rho = tf.clip_by_value(rho, -1 / (1 + EPSILON), 1 / (1 + EPSILON))
return x.var * bu.delta(rho, mu1, mu2)
z_mean = sqrt_x_var_diag * bu.softrelu(mu)
y_mean = tf.matmul(z_mean, A.mean) + b.mean
z_cov = relu_covariance(x)
y_cov = linear_covariance(z_mean, z_cov, A, b)
return gv.GaussianVar(y_mean, y_cov)
def simple(x, A, b):
mu = x.mean
y_mean = tf.matmul(mu, A.mean) + b.mean
y_cov = x.var
return gv.GaussianVar(y_mean, y_cov)