models.py

import tensorflow as tf
import numpy as np
from utils import lognormdens, weight_variable, bias_variable


class MLP(object):

    def __init__(self, dims, activations, stddev=1., bias_value=0.0):
        self.dims = dims
        self.activations = activations
        self.layers = []
        previous_dim = dims[0]
        for i, dim, activation in zip(xrange(len(activations)),
                                      dims[1:], activations):
            with tf.variable_scope('layer' + str(i)):
                weights = weight_variable((previous_dim, dim),
                                          stddev / np.sqrt(previous_dim))
                if i < len(activations) - 1:
                    biases = bias_variable((dim,), value=bias_value)
                else:
                    biases = bias_variable((dim,), value=0.0)

            self.layers.append((weights, biases, activation))
            previous_dim = dim

    def __call__(self, x, add_bias=True, return_activations=False):
        h = x
        hidden = []
        for weights, biases, activation in self.layers:
            h = tf.matmul(h, weights)
            if add_bias:
                h += biases
            if activation:
                h = activation(h)
            hidden.append(h)
        self.hidden = hidden
        if return_activations:
            return hidden
        else:
            return h

    def get_forward_derivative(self, x, fprimes):
        h = x
        for layer, fprime in zip(self.layers, fprimes):
            weights, biases, activation = layer
            h = tf.matmul(h, weights)
            h *= fprime
        return h


class MLPBlock(object):
    """Applies a separate MLP to each dimension of the input.

    The output dimensionality is assumed to be identical to the input.
    """

    def __init__(self, input_dim, hidden_dim, n_layers=1,
            stddev=1., bias_value=0.0):
        # bias value will only be applied to the hidden layers
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.hidden_layers = []
        with tf.variable_scope('block_mlp'):
            self.w_in_var = weight_variable((input_dim, input_dim * hidden_dim),
                                   stddev / np.sqrt(hidden_dim),
                                   name='w_in')
            self.w_out_var = weight_variable((input_dim * hidden_dim, input_dim),
                                    stddev / np.sqrt(hidden_dim),
                                    name='w_out')
            mask = np.zeros((input_dim, input_dim * hidden_dim),
                            dtype='float32')
            hid_to_hid_mask = np.zeros((input_dim * hidden_dim,
                                        input_dim * hidden_dim),
                                       dtype='float32')
            self.bias_hid = bias_variable((hidden_dim * input_dim,),
                                          value=bias_value,
                                          name='bias_first_hid')
            self.bias_out = bias_variable((input_dim,),
                                          name='bias_out')
            for i, row in enumerate(mask):
                row[i * hidden_dim:(i + 1) * hidden_dim] = 1.0

            for i in range(0, input_dim * hidden_dim, hidden_dim):
                hid_to_hid_mask[i:i + hidden_dim, i:i + hidden_dim] = 1.0

            self.hid_to_hid_mask = tf.convert_to_tensor(hid_to_hid_mask)
            self.in_out_mask = tf.convert_to_tensor(mask)
            self.w_in = self.w_in_var * self.in_out_mask
            self.w_out = self.w_out_var * tf.transpose(self.in_out_mask)
            for i in range(n_layers - 1):
                with tf.variable_scope('layer_' + str(i)):
                    w_hid = weight_variable((input_dim * hidden_dim,
                                             input_dim * hidden_dim),
                                             stddev / np.sqrt(hidden_dim))
                    b_hid = bias_variable((hidden_dim * input_dim,),
                                          value=bias_value)
                    self.hidden_layers.append((w_hid * self.hid_to_hid_mask,
                                               b_hid))

    def __call__(self, y, **kwargs):
        return self.forward(y, **kwargs)

    def forward(self, y, activation=None):
        h = tf.matmul(y, self.w_in) + self.bias_hid
        if activation is not None:
            h = activation(h)
        for w_hid, b_hid in self.hidden_layers:
            h = tf.matmul(h, w_hid) + b_hid
            if activation is not None:
                h = activation(h)
        x = tf.matmul(h, self.w_out) + self.bias_out
        return x


class LinearMISEP(object):
    """Implements the linear version of the MISEP model.

    MISEP, as described in [almeida2003]_,
    is an infomax-based ICA model. A normalized neural network replaces
    the usual sigmoid non-linearity of the infomax model.

    .. [almeida2003]
    Almeida, L. B. (2003).
    *MISEP--Linear and Nonlinear ICA Based on Mutual Information.*
    Journal of Machine Learning Research, 4(Dec), 1297-1318.

    """

    def __init__(self, input_dim, hidden_dim):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        with tf.variable_scope('misep'):
            self.F = MLP([input_dim, input_dim], [None])
            # Psi should be contstrained to be non-decreasing.  In the paper
            # this is done by using sigmoid functions and normalizing the norm
            # of the 'output' weights to be 1 / sqrt(h) and ensuring that the
            # weights are positive.  According to the paper, initializing the
            # weights with positive values ensures that they remain positive due
            # to the nature of the objective.  The input matrix should ensure
            # independent Psi functions. An n-diagonal martrix should work.  In
            # practice, tanh is used instead of the logistic sigmoid.
            w_in = tf.random_uniform(shape=(input_dim, input_dim * hidden_dim),
                                     minval=-.1, maxval=.1)
            w_out = tf.random_uniform(shape=(input_dim * hidden_dim, input_dim),
                                      minval=0, maxval=.2)
            self.Psi_weights_in = tf.Variable(w_in, name='Psi_weights_in')
            self.Psi_weights_out = tf.Variable(w_out, name='Psi_weights_out')
            mask = np.zeros((input_dim, input_dim * hidden_dim),
                            dtype='float32')
            self.bias = bias_variable((hidden_dim * input_dim,))
        for i, row in enumerate(mask):
            row[i * hidden_dim:(i + 1) * hidden_dim] = 1.0

        self.Psi_mask = tf.convert_to_tensor(mask)
        self.w_in = self.Psi_weights_in * self.Psi_mask
        self.w_out = self.Psi_weights_out * tf.transpose(self.Psi_mask)

    def forward(self, o):
        y = self.F(o)
        h = tf.matmul(y, self.w_in) + self.bias
        h = tf.tanh(h)
        w_out = self.w_out / tf.sqrt(tf.reduce_sum(self.w_out**2, 0,
                                                   keep_dims=True))
        w_out = w_out / tf.sqrt(tf.cast(self.hidden_dim, 'float32'))
        x = tf.matmul(h, w_out)
        return x, h, y

    def compute_jacobian(self, o):
        x, h, _ = self.forward(o) # h is (batch, input_dim*hidden_dim)
        y_prime = self.F(tf.eye(self.input_dim), add_bias=False)
        h_prime = tf.matmul(y_prime, self.w_in)
        # h_prime is now (input_dim, input_dim*hidden_dim)
        h_prime = tf.expand_dims(h_prime, 0)
        h_prime = h_prime * tf.expand_dims(1 - h**2, 1)
        # h_prime should now be (batch, input_dim, input_dim*hidden_dim)
        w_out = self.w_out / tf.sqrt(tf.reduce_sum(self.w_out**2, 0,
                                                   keep_dims=True))
        w_out = w_out / tf.sqrt(tf.cast(self.hidden_dim, 'float32'))
        # J should be (batch, input_dim, input_dim)
        # FIXME: einsum only seems to work when the batch dimension is known.
        J = tf.einsum('aij,jk->aik', h_prime, w_out) 
        return J

    def get_log_det_jacobian(self, o):
        J = self.compute_jacobian(o)
        def step(J_i):
            return tf.log(tf.abs(tf.matrix_determinant(J_i)))
        return tf.map_fn(step, J)

    def get_log_det_jacobian2(self, o):
        # requires tensorflow >= 1.0 but is much faster
        J = self.compute_jacobian(o)
        operator = tf.contrib.linalg.LinearOperatorFullMatrix(J)
        return operator.log_abs_determinant()


class PNLMISEP(object):
    """Implements a PNL version of the MISEP model.

    MISEP, as described in [almeida2003]_,
    is an infomax-based ICA model. A normalized neural network replaces
    the usual sigmoid non-linearity of the infomax model.

    This version uses an architecture which correponds to a learnable
    non-linearity followed by a linear transfmation. Subsequently, there is
    another constrained learnable non-linearity like in the linear MISEP model.
    This architecture is based on [zheng2007]_.

    .. [almeida2003]
    Almeida, L. B. (2003).
    *MISEP--Linear and Nonlinear ICA Based on Mutual Information.*
    Journal of Machine Learning Research, 4(Dec), 1297-1318.

    .. [zheng2007]
    Zheng, C. H., Huang, D. S., Li, K., Irwin, G., & Sun, Z. L. (2007).
    *MISEP method for postnonlinear blind source separation.*
    Neural computation, 19(9), 2557-2578.


    """

    def __init__(self, input_dim, hidden_dim, scaling=1.0, stddev=1.0):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        with tf.variable_scope('misep'):
            # TODO: should just make matrix for this
            self.F = MLP([input_dim, input_dim], [None], stddev=stddev)
            # Psi should be contstrained to be non-decreasing.  In the paper
            # this is done by using sigmoid functions and normalizing the norm
            # of the 'output' weights to be 1 / sqrt(h) and ensuring that the
            # weights are positive.  According to the paper, initializing the
            # weights with positive values ensures that they remain positive due
            # to the nature of the objective.  The input matrix should ensure
            # independent Psi functions. An n-diagonal martrix should work.  In
            # practice, tanh is used instead of the logistic sigmoid.
            entrance_w_in = tf.random_uniform(shape=(input_dim,
                                                     input_dim * hidden_dim),
                                              minval=-.1, maxval=.1) * scaling
            entrance_w_out = tf.random_uniform(shape=(input_dim * hidden_dim,
                                                      input_dim),
                                               minval=0, maxval=.2) * scaling
            exit_w_in = tf.random_uniform(shape=(input_dim,
                input_dim * hidden_dim), minval=-.1, maxval=.1) * scaling
            exit_w_out = tf.random_uniform(shape=(input_dim * hidden_dim,
                                                  input_dim),
                                           minval=0, maxval=.2) * scaling
            self.entrance_weights_in = tf.Variable(entrance_w_in,
                                                   name='entrance_weights_in')
            self.entrance_weights_out = tf.Variable(entrance_w_out,
                                                    name='entrance_weights_out')
            self.exit_weights_in = tf.Variable(exit_w_in,
                                               name='exit_weights_in')
            self.exit_weights_out = tf.Variable(exit_w_out,
                                                name='exit_weights_out')
            mask = np.zeros((input_dim, input_dim * hidden_dim),
                    dtype='float32')
            self.entrance_bias_h = bias_variable((hidden_dim * input_dim,))
            self.entrance_bias_o = bias_variable((input_dim,))
            self.exit_bias_h = bias_variable((hidden_dim * input_dim,))

        for i, row in enumerate(mask):
            row[i * hidden_dim:(i + 1) * hidden_dim] = 1.0

        self.mask = tf.convert_to_tensor(mask)
        self.entrance_w_in = self.entrance_weights_in * self.mask
        self.entrance_w_out = self.entrance_weights_out * tf.transpose(self.mask)
        self.exit_w_in = self.exit_weights_in * self.mask
        self.exit_w_out = self.exit_weights_out * tf.transpose(self.mask)

    def forward(self, o):
        entrance_h = tf.matmul(o, self.entrance_w_in) + self.entrance_bias_h
        entrance_h = tf.tanh(entrance_h)
        entrance_o = (tf.matmul(entrance_h, self.entrance_w_out) +
                      self.entrance_bias_o)
        y = self.F(entrance_o)  # prediction of the model
        exit_h = tf.matmul(y, self.exit_w_in) + self.exit_bias_h
        exit_h = tf.tanh(exit_h)
        w_out = self.exit_w_out / tf.sqrt(tf.reduce_sum(self.exit_w_out**2, 0,
                                                        keep_dims=True))
        w_out = w_out / tf.sqrt(tf.cast(self.hidden_dim, 'float32'))
        x = tf.matmul(exit_h, w_out)
        return x, entrance_h, exit_h, y

    def compute_jacobian(self, o):
        # 90s style manual gradient computations
        x, entrance_h, exit_h, y = self.forward(o)
        entrance_h_prime = tf.matmul(tf.eye(self.input_dim),
                                     self.entrance_w_in)
        # entrance_h_prime is now (input_dim, input_dim*hidden_dim)
        entrance_h_prime = tf.expand_dims(entrance_h_prime, 0)
        entrance_h_prime *= tf.expand_dims(1 - entrance_h**2, 1)
        # entrance_h_prime should be (batch, input_dim, input_dim*hidden_dim)
        entrance_o_prime = tf.einsum('aij,jk->aik', entrance_h_prime,
                self.entrance_w_out) 
        # we're at (batch, input_dim, input_dim)
        y_prime = tf.einsum('aij,jk->aik',
                            entrance_o_prime,
                            self.F.layers[0][0])
        # still at (batch, input_dim, input_dim)
        exit_h_prime = tf.einsum('aij,jk->aik',
                                 y_prime,
                                 self.exit_w_in) 
        # (batch, input_dim, input_dim*hidden_dim)
        # exit_h should be (batch, input_dim*hidden_dim)
        exit_h_prime = exit_h_prime * tf.expand_dims(1 - exit_h**2, 1)
        w_out = self.exit_w_out / tf.sqrt(tf.reduce_sum(self.exit_w_out**2, 0,
                                                   keep_dims=True))
        w_out = w_out / tf.sqrt(tf.cast(self.hidden_dim, 'float32'))
        # J should be (batch, input_dim, input_dim)
        J = tf.einsum('aij,jk->aik', exit_h_prime, w_out) 
        return J

    def get_log_det_jacobian(self, o):
        J = self.compute_jacobian(o)

        def step(J_i):
            return tf.log(tf.abs(tf.matrix_determinant(J_i)))

        return tf.map_fn(step, J)

    def get_log_det_jacobian2(self, o):
        # requires tensorflow >= 1.0 but is much faster
        J = self.compute_jacobian(o)
        operator = tf.contrib.linalg.LinearOperatorFullMatrix(J)
        return operator.log_abs_determinant()