variational-inference-with-normalizing-flows
variational-inference-with-normalizing-flows copied to clipboard
ex4sperans
Ensuring the planar flow transformation is invertible
Hey! I believe you should reparameterize u and w in the planar flow like they do in the pymc3 (link below) to ensure the planar flow transformation is invertible, see the appendix in the paper https://arxiv.org/abs/1505.05770.
https://github.com/pymc-devs/pymc3/blob/1cdd1631bea48fef8d140e37c3588a8208498ba0/pymc3/variational/flows.py#L374
Hey! Right, I agree that naive implementation of the planar flow is not always invertible, but I believe the invertibility is not required is the case of KL minimization (what I do). Please correct me if my understanding is wrong.
Anyway, it definitely should be implemented to be always invertible, thank you.
I think it would not be a proper probability distribution if the flow is not invertible, as you need the determinant of inverse jacobian to transform a random variable.
Hi, guys, so have the code been revised for invertible planar flow transformation? @ex4sperans @kingofspace0wzz @talesa
No, it hasn't.
https://github.com/ex4sperans/variational-inference-with-normalizing-flows/blob/master/flow.py#L53
I haven't tested my code but this should do the trick
class PlanarFlow(nn.Module):
def __init__(self, dim):
super().__init__()
self.weight = nn.Parameter(torch.Tensor(1, dim))
self.scale = nn.Parameter(torch.Tensor(1, dim))
self.bias = nn.Parameter(torch.Tensor(1))
self.tanh = nn.Tanh()
self.reset_parameters()
def reset_parameters(self):
self.weight.data.uniform_(-0.01, 0.01)
self.scale.data.uniform_(-0.01, 0.01)
self.bias.data.uniform_(-0.01, 0.01)
def forward(self, z):
activation = F.linear(z, self.weight, self.bias)
scale = self.get_scale(self.scale, self.weight)
return z + scale * self.tanh(activation)
def get_scale(self, scale, weight):
wu = torch.sum(weight * scale, dim=-1, keepdim=True)
mwu = -1 + torch.log(1 + torch.exp(wu))
u_h = scale + (mwu - wu) * (weight + 1e-5)/(torch.norm(weight) + 1e-5)
return u_h
class PlanarFlowLogDetJacobian(nn.Module):
def __init__(self, affine):
super().__init__()
self.weight = affine.weight
self.bias = affine.bias
self.scale = affine.scale
self.tanh = affine.tanh
def forward(self, z):
activation = F.linear(z, self.weight, self.bias)
psi = (1 - self.tanh(activation) ** 2) * self.weight
scale = self.get_scale(self.scale, self.weight)
det_grad = 1 + torch.mm(psi, scale.t())
return safe_log(det_grad.abs())
def get_scale(self, scale, weight):
wu = torch.sum(weight * scale, dim=-1, keepdim=True)
mwu = -1 + torch.log(1 + torch.exp(wu))
u_h = scale + (mwu - wu) * (weight + 1e-5)/(torch.norm(weight) + 1e-5)
return u_h
Quick question, by ensuring planar-flow invertibility is there a closed form formula to find such inverse?
@thipokKub Unfortunately not, one can ensure invertibility without having a closed-form expression of the inverse, see "Invertible Residual Networks" for a more extreme example.
