def get_Abs(Ws, vs, Qs):
"""produces the pre activation feature maps of layer ell"""
n = Qs[-1].shape[0]
As = [Ws[0] * T.ones((n, 1, 1))]
bs = [vs[0] * T.ones((n, 1))]
for i in range(len(Qs)):
As.append(T.einsum('db,nb,nbs->nds', Ws[i + 1], Qs[i], As[-1]))
bs.append(T.einsum('db,nb,nb->nd', Ws[i + 1], Qs[i], bs[-1])\
+ vs[i + 1])
return As, bs
def compute_W_derivative(lhs, rhs, Ws, Qs, l):
L = len(Qs)
backward_As = [rhs]
for i in range(0, l):
backward_As.append(T.einsum('na,as,ns->na', Qs[i], Ws[i],
backward_As[-1]))
forward_As = [lhs]
for i in range(l, L):
forward_As.append(T.einsum('ns,sb,nb->nb', forward_As[-1],
Ws[- i], Qs[-1]))
return T.einsum('na,nb->nab', forward_As, backward_As)
def get_forward(Ws, Qs):
"""this function gives the slope matrix that forwards any pre activation
of layer l to the output layer which is ::
W^{L}Q^{L}_{\omega}W^{\ell}Q^{\ell}
for the \ell element of the returned list. For the first one, is returns
the entire A_{\omega} and for the last one it is the identity matrix
"""
N = Qs[-1].shape[0]
L = len(Qs)
forward = [T.identity(Ws[-1].shape[0]) * T.ones((N, 1, 1))]
for i in range(L):
forward.append(T.einsum('ndb,bs,ns->nds', forward[-1], Ws[- 1 - i],
Qs[- 1 - i]))
return forward[::-1]
def create_fns(input,
in_signs,
Ds,
x,
m0,
m1,
m2,
batch_in_signs,
alpha=0.1,
sigma=1,
sigma_x=1,
lr=0.0002):
cumulative_units = np.concatenate([[0], np.cumsum(Ds[:-1])])
BS = batch_in_signs.shape[0]
Ws = [
T.Variable(sj.initializers.glorot((j, i)) * sigma)
for j, i in zip(Ds[1:], Ds[:-1])
]
bs = [T.Variable(sj.initializers.he((j,)) * sigma) for j in Ds[1:-1]]\
+ [T.Variable(T.zeros((Ds[-1],)))]
A_w = [T.eye(Ds[0])]
B_w = [T.zeros(Ds[0])]
A_q = [T.eye(Ds[0])]
B_q = [T.zeros(Ds[0])]
batch_A_q = [T.eye(Ds[0]) * T.ones((BS, 1, 1))]
batch_B_q = [T.zeros((BS, Ds[0]))]
maps = [input]
signs = []
masks = [T.ones(Ds[0])]
in_masks = T.where(T.concatenate([T.ones(Ds[0]), in_signs]) > 0, 1., alpha)
batch_in_masks = T.where(
T.concatenate([T.ones((BS, Ds[0])), batch_in_signs], 1) > 0, 1., alpha)
for w, b in zip(Ws[:-1], bs[:-1]):
pre_activation = T.matmul(w, maps[-1]) + b
signs.append(T.sign(pre_activation))
masks.append(T.where(pre_activation > 0, 1., alpha))
maps.append(pre_activation * masks[-1])
maps.append(T.matmul(Ws[-1], maps[-1]) + bs[-1])
# compute per region A and B
for start, end, w, b, m in zip(cumulative_units[:-1], cumulative_units[1:],
Ws, bs, masks):
A_w.append(T.matmul(w * m, A_w[-1]))
B_w.append(T.matmul(w * m, B_w[-1]) + b)
A_q.append(T.matmul(w * in_masks[start:end], A_q[-1]))
B_q.append(T.matmul(w * in_masks[start:end], B_q[-1]) + b)
batch_A_q.append(
T.matmul(w * batch_in_masks[:, None, start:end], batch_A_q[-1]))
batch_B_q.append((w * batch_in_masks[:, None, start:end]\
* batch_B_q[-1][:, None, :]).sum(2) + b)
batch_B_q = batch_B_q[-1]
batch_A_q = batch_A_q[-1]
signs = T.concatenate(signs)
inequalities = T.hstack(
[T.concatenate(B_w[1:-1])[:, None],
T.vstack(A_w[1:-1])]) * signs[:, None]
inequalities_code = T.hstack(
[T.concatenate(B_q[1:-1])[:, None],
T.vstack(A_q[1:-1])]) * in_signs[:, None]
#### loss
log_sigma2 = T.Variable(sigma_x)
sigma2 = T.exp(log_sigma2)
Am1 = T.einsum('qds,nqs->nqd', batch_A_q, m1)
Bm0 = T.einsum('qd,nq->nd', batch_B_q, m0)
B2m0 = T.einsum('nq,qd->n', m0, batch_B_q**2)
AAm2 = T.einsum('qds,qdu,nqup->nsp', batch_A_q, batch_A_q, m2)
inner = -(x * (Am1.sum(1) + Bm0)).sum(1) + (Am1 * batch_B_q).sum((1, 2))
loss_2 = (x**2).sum(1) + B2m0 + T.trace(AAm2, axis1=1, axis2=2).squeeze()
loss_z = T.trace(m2.sum(1), axis1=1, axis2=2).squeeze()
cst = 0.5 * (Ds[0] + Ds[-1]) * T.log(2 * np.pi)
loss = cst + 0.5 * Ds[-1] * log_sigma2 + inner / sigma2\
+ 0.5 * loss_2 / sigma2 + 0.5 * loss_z
mean_loss = loss.mean()
adam = sj.optimizers.NesterovMomentum(mean_loss, Ws + bs, lr, 0.9)
train_f = sj.function(batch_in_signs,
x,
m0,
m1,
m2,
outputs=mean_loss,
updates=adam.updates)
f = sj.function(input,
outputs=[maps[-1], A_w[-1], B_w[-1], inequalities, signs])
g = sj.function(in_signs, outputs=[A_q[-1], B_q[-1]])
all_g = sj.function(in_signs, outputs=inequalities_code)
h = sj.function(input, outputs=maps[-1])
return f, g, h, all_g, train_f, sigma2
def create_fns(batch_size, R, Ds, seed, leakiness=0.1, lr=0.0002, scaler=1,
var_x=1):
alpha = T.Placeholder((1,), 'float32')
x = T.Placeholder((Ds[0],), 'float32')
X = T.Placeholder((batch_size, Ds[-1]), 'float32')
signs = T.Placeholder((np.sum(Ds[1:-1]),), 'float32')
SIGNS = T.Placeholder((R, np.sum(Ds[1:-1])), 'float32')
m0 = T.Placeholder((batch_size, R), 'float32')
m1 = T.Placeholder((batch_size, R, Ds[0]), 'float32')
m2 = T.Placeholder((batch_size, R, Ds[0], Ds[0]), 'float32')
Ws, vs = init_weights(Ds, seed, scaler)
Ws = [T.Variable(w, name='W' + str(l)) for l, w in enumerate(Ws)]
vs = [T.Variable(v, name='v' + str(l)) for l, v in enumerate(vs)]
var_x = T.Variable(T.ones(Ds[-1]) * var_x)
var_z = T.Variable(T.ones(Ds[0]))
# create the placeholders
Ws_ph = [T.Placeholder(w.shape, w.dtype) for w in Ws]
vs_ph = [T.Placeholder(v.shape, v.dtype) for v in vs]
var_x_ph = T.Placeholder(var_x.shape, var_x.dtype)
############################################################################
# Compute the output of g(x)
############################################################################
maps = [x]
xsigns = []
masks = []
for w, v in zip(Ws[:-1], vs[:-1]):
pre_activation = T.matmul(w, maps[-1]) + v
xsigns.append(T.sign(pre_activation))
masks.append(relu_mask(pre_activation, leakiness))
maps.append(pre_activation * masks[-1])
xsigns = T.concatenate(xsigns)
maps.append(T.matmul(Ws[-1], maps[-1]) + vs[-1])
############################################################################
# compute the masks and then the per layer affine mappings
############################################################################
cumulative_units = np.cumsum([0] + Ds[1:])
xqs = relu_mask([xsigns[None, cumulative_units[i]:cumulative_units[i + 1]]
for i in range(len(Ds) - 2)], leakiness)
qs = relu_mask([signs[None, cumulative_units[i]:cumulative_units[i + 1]]
for i in range(len(Ds) - 2)], leakiness)
Qs = relu_mask([SIGNS[:, cumulative_units[i]:cumulative_units[i + 1]]
for i in range(len(Ds) - 2)], leakiness)
Axs, bxs = get_Abs(Ws, vs, xqs)
Aqs, bqs = get_Abs(Ws, vs, qs)
AQs, bQs = get_Abs(Ws, vs, Qs)
all_bxs = T.hstack(bxs[:-1]).transpose()
all_Axs = T.hstack(Axs[:-1])[0]
all_bqs = T.hstack(bqs[:-1]).transpose()
all_Aqs = T.hstack(Aqs[:-1])[0]
x_inequalities = T.hstack([all_Axs, all_bxs]) * xsigns[:, None]
q_inequalities = T.hstack([all_Aqs, all_bqs]) * signs[:, None]
############################################################################
# loss (E-step NLL)
############################################################################
Bm0 = T.einsum('nd,Nn->Nd', bQs[-1], m0)
B2m0 = T.einsum('nd,Nn->Nd', bQs[-1] ** 2, m0)
Am1 = T.einsum('nds,Nns->Nd', AQs[-1], m1)
ABm1 = T.einsum('nds,nd,Nns->Nd', AQs[-1], bQs[-1], m1)
Am2ATdiag = T.diagonal(T.einsum('nds,Nnsc,npc->Ndp', AQs[-1], m2, AQs[-1]),
axis1=1, axis2=2)
xAm1Bm0 = X * (Am1 + Bm0)
M2diag = T.diagonal(m2.sum(1), axis1=1, axis2=2)
prior = sum([T.mean(w**2) for w in Ws], 0.) / cov_W\
+ sum([T.mean(v**2) for v in vs[:-1]], 0.) / cov_b
loss = - 0.5 * (T.log(var_x).sum() + T.log(var_z).sum()\
+ (M2diag / var_z).sum(1).mean() + ((X ** 2 - 2 * xAm1Bm0 + B2m0\
+ Am2ATdiag + 2 * ABm1) / var_x).sum(1).mean())
mean_loss = - (loss + 0.5 * prior)
adam = sj.optimizers.SGD(mean_loss, 0.001, params=Ws + vs)
############################################################################
# update of var_x
############################################################################
update_varx = (X ** 2 - 2 * xAm1Bm0 + B2m0 + Am2ATdiag + 2 * ABm1).mean()\
* T.ones(Ds[-1])
update_varz = M2diag.mean() * T.ones(Ds[0])
############################################################################
# update for biases IT IS DONE FOR ISOTROPIC COVARIANCE MATRIX
############################################################################
FQ = get_forward(Ws, Qs)
update_vs = {}
for i in range(len(vs)):
if i < len(vs) - 1:
# now we forward each bias to the x-space except the ith
separated_bs = bQs[-1] - T.einsum('nds,s->nd', FQ[i], vs[i])
# compute the residual and apply sigma
residual = (X[:, None, :] - separated_bs) * m0[:, :, None]\
- T.einsum('nds,Nns->Nnd', AQs[-1], m1)
back_error = T.einsum('nds,nd->s', FQ[i], residual.mean(0))
whiten = T.einsum('ndc,nds,n->cs', FQ[i] , FQ[i], m0.mean(0))\
+ T.eye(back_error.shape[0]) / (Ds[i] * cov_b)
update_vs[vs[i]] = T.linalg.solve(whiten, back_error)
else:
back_error = (X - (Am1 + Bm0) + vs[-1])
update_vs[vs[i]] = back_error.mean(0)
############################################################################
# update for slopes IT IS DONE FOR ISOTROPIC COVARIANCE MATRIX
############################################################################
update_Ws = {}
for i in range(len(Ws)):
U = T.einsum('nds,ndc->nsc', FQ[i], FQ[i])
if i == 0:
V = m2.mean(0)
else:
V1 = T.einsum('nd,nq,Nn->ndq', bQs[i-1], bQs[i-1], m0)
V2 = T.einsum('nds,nqc,Nnsc->ndq', AQs[i-1], AQs[i-1], m2)
V3 = T.einsum('nds,nq,Nns->ndq', AQs[i-1], bQs[i-1], m1)
Q = T.einsum('nd,nq->ndq', Qs[i - 1], Qs[i - 1])
V = Q * (V1 + V2 + V3 + V3.transpose((0, 2, 1))) / batch_size
whiten = T.stack([T.kron(U[n], V[n]) for n in range(V.shape[0])]).sum(0)
whiten = whiten + T.eye(whiten.shape[-1]) / (Ds[i]*Ds[i+1]*cov_W)
# compute the residual (bottom up)
if i == len(Ws) - 1:
bottom_up = (X[:, None, :] - vs[-1])
else:
if i == 0:
residual = (X[:, None, :] - bQs[-1])
else:
residual = (X[:, None, :] - bQs[-1]\
+ T.einsum('nds,ns->nd', FQ[i - 1], bQs[i-1]))
bottom_up = T.einsum('ndc,Nnd->Nnc', FQ[i], residual)
# compute the top down vector
if i == 0:
top_down = m1
else:
top_down = Qs[i - 1] * (T.einsum('nds,Nns->Nnd', AQs[i - 1], m1) +\
T.einsum('nd,Nn->Nnd', bQs[i - 1], m0))
vector = T.einsum('Nnc,Nns->cs', bottom_up, top_down) / batch_size
condition = T.diagonal(whiten)
update_W = T.linalg.solve(whiten, vector.reshape(-1)).reshape(Ws[i].shape)
update_Ws[Ws[i]] = update_W
############################################################################
# create the io functions
############################################################################
params = sj.function(outputs = Ws + vs + [var_x])
ll = T.Placeholder((), 'int32')
selector = T.one_hot(ll, len(vs))
for i in range(len(vs)):
update_vs[vs[i]] = ((1 - alpha) * vs[i] + alpha * update_vs[vs[i]])\
* selector[i] + vs[i] * (1 - selector[i])
for i in range(len(Ws)):
update_Ws[Ws[i]] = ((1 - alpha) * Ws[i] + alpha * update_Ws[Ws[i]])\
* selector[i] + Ws[i] * (1 - selector[i])
output = {'train':sj.function(SIGNS, X, m0, m1, m2, outputs=mean_loss,
updates=adam.updates),
'update_var':sj.function(SIGNS, X, m0, m1, m2, outputs=mean_loss,
updates = {var_x: update_varx}),
'update_vs':sj.function(alpha, ll, SIGNS, X, m0, m1, m2, outputs=mean_loss,
updates = update_vs),
'loss':sj.function(SIGNS, X, m0, m1, m2, outputs=mean_loss),
'update_Ws':sj.function(alpha, ll, SIGNS, X, m0, m1, m2, outputs=mean_loss,
updates = update_Ws),
'signs2Ab': sj.function(signs, outputs=[Aqs[-1][0], bqs[-1][0]]),
'signs2ineq': sj.function(signs, outputs=q_inequalities),
'g': sj.function(x, outputs=maps[-1]),
'input2all': sj.function(x, outputs=[maps[-1], Axs[-1][0],
bxs[-1][0], x_inequalities, xsigns]),
'get_nll': sj.function(SIGNS, X, m0, m1, m2, outputs=mean_loss),
'assign': sj.function(*Ws_ph, *vs_ph, var_x_ph,
updates=dict(zip(Ws + vs + [var_x],
Ws_ph + vs_ph + [var_x_ph]))),
'varx': sj.function(outputs=var_x),
'prior': sj.function(outputs=prior),
'varz': sj.function(outputs=var_z),
'params': params,
# 'probed' : sj.function(SIGNS, X, m0, m1, m2, outputs=probed),
'input2signs': sj.function(x, outputs=xsigns),
'S' : Ds[0], 'D': Ds[-1], 'R': R, 'model': 'EM', 'L':len(Ds)-1,
'kwargs': {'batch_size': batch_size, 'Ds':Ds, 'seed':seed,
'leakiness':leakiness, 'lr':lr, 'scaler':scaler}}
def sample(n):
samples = []
for i in range(n):
samples.append(output['g'](np.random.randn(Ds[0])))
return np.array(samples)
output['sample'] = sample
return output