NICE - Nonlinear Independent Components Estimation

Nonlinear Independent Components Estimation (NICE) is a normalizing flow model proposed by Dinh et al., 2014. It is a specially designed neural network architecture that allows computation of exact log-likelihood, exact and efficient sampling and exact and efficient inference of latent variables. The neural architecture is invertible and therefore provides a way of representing only non-self-intersecting geometry. We do this by using a Pre-Aux net with the NICE architecture. In this tutorial, we walk through the theory behind the NICE representation and how to actually use their geodiff implementation in practice.

Theoretical Background¶

The NICE architecture as proposed in Dinh et al., 2014 is based on the idea of stacking simple nonlinear invertible transformations. The architecture is built such that it is easy to compute the inverse and the determinant of the Jacobian of the transformation. Computing the inverse and the determinant of the Jacobian is useful for log-likelihood training of generative models (of say images) but we do not require these properties and are only interested in the architecture being guaranteed invertible. That is, as long as we are concerned we only need a guarantee on the transformation represented by the network being invertible we will never actually want to invert it or compute the Jacobian of the determinant. Therefore, we only focus on the forward pass of the network and do not discuss the computations of the inverse or the determinant of the Jacobian.

The NICE architecture is composed of bijective transformations with triangular Jacobians that serve as a building block for the entire transformation f. In particular, it uses an additive coupling layer as the building block defined as follows:

Additive coupling layer: Let $x \in \mathbb{R}^{D}$ and $I_1, I_2$ be a partition of $\llbracket 1, D \rrbracket$ with $d = |I_1|$ and $m$ be a function defined on $\mathbb{R}^{D}$. We can define $y = (y_{I_1}, y_{I_2})$ where:

where, $m$ is called the additive coupling law. If we choose $I_1 = \llbracket 1, \, d \rrbracket, \ I_2 = \llbracket d+1, \, D \rrbracket$ then the Jacobian is given by:

The inverse mapping is obtained by:

Each additive coupling layer has unit Jacobian determinant as is clear from Eq. (2), i.e. is volume preserving. Therefore, their composition will necessarily have unit Jacobian determinant too. In order to address this issue, we include a diagonal scaling matrix $S$ as the top layer, which multiplies the $i$-th output value by $S_{ii}$ giving $(x_i)_{i \leq D} \rightarrow (S_{ii}x_i)_{i \leq D}$.

Implementation using geodiff¶

We now look at how geodiff allows us to easily use the NICE parameterization to represent shapes.

The NICE class initializer and its expected arguments are shown below:

    def __init__(
        self,
        geometry_dim: int,
        layer_count: int,
        preaux_net: nn.Module,
        coupling_net: nn.Module,
        volume_preserving: bool = False,
        use_batchnormalization: bool = False,
        use_residual_connection: bool = True,
    ) -> None:

To construct a NICE object, we need to supply:

• Geometry Dimension 2 for 2D and 3 for 3D.

• Layer Count the number of layers in the NICE architecture.

• Coupling Net a torch network module to be used as the additive coupling law.

• Pre-Aux Net parameters such as layer count, hidden dimension, activation, normalization and output functions.

• Volume Preserving, to decide whether the NICE architecture is used without the scaling layer at the top. This is best used with a Pre-Aux net with no parameters that preserves the area or volume of the closed circle or sphere.

• Boolean Flags for using batch normalization and residual connections with the NICE layers.

Fitting Shapes¶

We will use a square as our target geometry. We then use the ChamferLoss to compute the geometric difference loss between the target shape and the shape represented by our NICE parameterization. Since the implementation is written in PyTorch we can use the autograd capabilities to compute the gradients of the loss w.r.t. the NICE network parameters and use an optimizer to modify our geometry.

We start by obtaining points on our shapes and normalizing them appropriately such that $x \in [0, 1]$. This is not needed for NICE but is used to offer a comparison with the classical shape representation methods like Hicks-Henne and CST.

We now create a NICE object by specifying the geometry dimension, layer count, coupling network and Pre-Aux net parameters.

We use the ChamferLoss provided by geodiff to compute a geometric loss between the target shape and the shape represented by the NICE object. PyTorch’s autograd capabilities then allow us to compute gradients of the loss w.r.t. the NICE network parameters and modify them to fit the target shape.

We can now visualize the shape represented by our NICE object using its visualize method.