# Chapter 1: Introduction

## Introduction: Explicit layers in deep learning

At the heart of modern deep learning methods is the notion of a _layer_.  Deep learning models are traditionally built by stacking many of these layers together, to create an archictecture designed to solve some particular task.  Convolutional networks, for example, consist of convolutional layers, typically followed elementwise nonlinearitis like the ReLU, with additional operations like normalization or dropout, and possibly connected together in multiple different ways, to form things like residual layers.  Likewise, architectures like Transformer networks consist of combinations of so-called self-attention layers and fully-connected layers, again stacked together in a manner than results in the final form of the model.

A common defining characteristic, which at this point is so standard that it often goes unnoticed by practioners, is that the vast majority of these layers in modern deep learning are defined _explicitly_.  That is, they are specified by an exact sequence of operations that it takes to do from the input to the output a layer.  Let's take a scaled self-attention layer as an example [[Vaswani et al.](https://arxiv.org/abs/1706.03762)].  This layer is a mapping from three matrices $K,Q,V \in \mathbb{R}^{T \times n}$ to an output $Z \in \mathbb{R}^{T \times n}$, and is defined by the operation
\begin{equation}
Z = \mathrm{SelfAttention}(K,Q,V) \equiv \mathrm{softmax}\left (\frac{K Q^T}{\sqrt{n}} \right) V
\end{equation}
(this is a simplified version of self-attention, just for illustration, with no masks or multi-head structure).  We could write this layer as a simple Python function (again, just for illustration, not how you would actually want to write e.g., the softmax operation, let alone the fact that you'd likely want to use an automatic differentiation library rather than plain numpy to write these functions).

```
import numpy as np

def self_attention(K,Q,V):
    A = np.exp(K @ Q.T) / np.sqrt(K.shape[1])
    return (A / np.sum(A,1)) @ V

K, Q, V = np.random.randn(3, 5, 4)
print(self_attention(K, Q, V))
```

    [[ 2.40681810e+00 -4.51330415e+00 -2.56135201e-01  5.83548198e-01]
     [ 1.38511508e-01 -2.45408681e-01  8.25641586e-02  7.51321715e-02]
     [-1.17558712e+00 -3.76026491e-01 -1.62134513e+00  8.56850990e-04]
     [ 1.65305095e+00 -2.59930825e+00 -4.29699547e-01 -8.58770304e-01]
     [ 1.11891553e-01 -1.36663381e-01  5.76039039e-02  3.96598953e-02]]

Naturally, things start to get a bit more complex as we add more functionality to the layers themselves, implement them within automatic differentiation libraries, etc, but this explicit form of the typical layer remains through it all: layers are built largely like typical computer programs, where we directly write the code to generate the output of the layer as a function of its input.  This may be so ingrained, in fact, that it's hard to imagine that there is an entirely different way that layers can be defined, namely via _implicit_ layers.

## Implicit layers

The crux of an implicit layer, as we will use the term throughout this document, is that _instead of specifying how to compute the layer's output from the input, we specify the conditions that we want the layer's output to satisfy_.  That is, if we were to write explicit layers (with input $x \in \mathcal{X}$ and output $z \in \mathcal{Z}$) as an application of some explicit function $f : \mathcal{X} \rightarrow \mathcal{Z}$
\begin{equation}
z = f(x)
\end{equation}
then an implicit layer would instead be defined via a function $g : \mathcal{X} \times \mathcal{Z} \rightarrow \mathbb{R}^n$, which is a joint function of both $x$ and $z$, and where the output of the layer $z$ is requied to satisfy some _constraint_ e.g., finding a root of the equation,
\begin{equation}
\mbox{Find $z$ such that } g(x,z) = 0.
\end{equation}
The notation here may indicate that $g(x,z)$ is a simple alegebraic equation, but in practice this same formalism can capture algebraic equations and fixed points, leading to recurent backprop models \cite{} or deep equilibrium models \cite{}; differential equations,leading to Neural ODEs \cite{}; or the optimality conditions of optimization problems, leadings to differentiable optimization approaches \cite.

Before we move on to a concrete example, we should highlight the fact that initially moving to this implicit formulation may seem like a trivial point. After all, in order to actually implement a layer like this, we would need to specify some way of actually _computing_ a root of the equation $g$.  But as we will see shortly, there are numerous practical advantages to considering the implicit form of a layer.

Most fundamentally, implicit form layers _separate the solution procedure of the layer from the definition of the layer itself_.  This level of modularity has proven extremely useful in a number of domains.  Differential equation solvers, for example, which attempt to find a numerical solution to an ordinary differntial equation, can implement all sorts of adaptive step sizes, corections for so-called ``stiff'' equations, etc, all in service of attempting to find a low-error solution to the differentiable equation.  Or as another example, optimization solvers often involve very complex heuristics for solving certain types of problems, but they are all aimed at finding the minimum-objective solution to an optimization task.  Indeed, because we rarely find _exact_ solutions to e.g., algebraic or differential equations, the different solutions methods can be evaluated objectively against each other based upon how well they satisfy the conditions that the layer is attempting to satisfy.

This separatation of the layer's objective and its solution method is desirable enough in and of itself, but a second advantage to implicit layers emerges specifically in the context of deep learning and automatic differentiation.  The traditional approach to automatic differentiation (AD) in machine learning is to implement all layers within an automatic differentiation framework (such as PyTorch, Tensorflow, or JAX), which immediately lets us include these layers in deep models that require gradients for fitting the models to data.  Yet implementing solution procedures, especially those involving iterative updates, like standard differential equation or optimization solvers, directly within AD library would mean that we need the store the computation graph for the complete solution procedure, along with the value of temporary iterates created during this solution.  This requires storing a great deal of information in memory, which can often be a bottleneck during training of large deep learning models.  Fortunately, as we will illustrate below, and highlight several times within this tutorial, implicit layers have the notably advantage that _we can use the implicit function theorem to directly compute gradients at the solution point of these equations_, without having to store any intermediate variables along the way.  This vastly improves the memory consumption and often the numerical accuracy of these methods, providing another notable benefit for implicit models in the setting of deep learning in particular.

### Applications and illustrations

Since the remainder of this chapter will focus on an extremely simple demonstration intended as a pedogogical illustration of the methods (rather than an illustration of state of the art performance), we want to briefly highlight the wide array of applications that have been addressed using implicit layers.  The following are just a small sampling (of instances that the authors happen to be most familiar with), but they hopefully give an illustration of the breadth of approaches addressed by current research in implicit layers.  We'll dive into more detail about a few of these in the text, but for the most part, you will want to look at the current research in the field to stay abreast of all the application areas being addressed by these approaches.

Implicit layers have been used to:
- Solve arbitrary structured convex problems (using the `cvxpy` library) in a differentiable manner.
- Solve smoothed relaxtions of combinatorial optimization problems, such as graph cuts, satisfiability, and many others.
- Integrate differential equations as layers in deep networks (with numerous applications in and of itself, such as integrating continuous time observations, or approximating continuous version of traditional residual networks).
- Create architectures for efficient representation of smooth densities, for use in generative modeling an beyond.
- Achieve performance on par with state-of-the-art Transformer models (at the same parameter count), for language modeling and on par with state-of-the-art computer vision architectures on tasks such as classification and semantic segmentation.

## Outline of this work

With the above brief introduction as context, we outline the remainder of this work and how the chapters fit together.

- In the remainder of **Chapter 1**, we will give an brief introduction to your first implicit layer, defined via a fixed point iteration.  This is essentially a version of recurrent backpropagation that was one of the first forms of implicit layers, tracing back to the late 80s \cite{}, and is also the approach that underlies deep equilibrium (DEQ) models.
- In **Chapter 2** we will discuss the mathematical background behind implicit models, including the implicit function theorem and its implementation in automatic differentiation tools.
- In **Chapter 3** we will present Neural ODEs, an instatiation of implicit layers that has received substantial attention in recent years.  We will present the basic mathematical framework as well as highlight a number of applications of the model.
- In **Chapter 4** we will present Deep Equilibrium Models in greater detail, focusing on extending the basic idea presented in Chapter 1 to modern deep learning frameworks, and highlighting some ongoing directions and applications of the models.
- In **Chapter 5** we will present differentiable optimization, which embeds the solution to optimization problems as layers.  We will specifically show how this relates to an captures and generalizes many existing layers in deep learing.

## Your first implicit model: fixed point iteration

Before diving into the mathematical details and many different forms of implicit models, let's start with a particularly simple example: a network layer defined by a fixed point iteration.  As mentioned above, this type of layer dates back to some of the original formulations of recurrent backpropagation, and is also the basis for the Deep Equilibrium Models we will discuss soon.

### A fixed point iteration layer

Athough we will shortly adopt a view of this layer as the root of a particular equation, to introduce the layer, suppose we have an inputs and outputs $x,z \in \mathbb{R}^n$ of the same dimension, and consider the following approach to computing the output $z$ as a function of $x$
\begin{equation}
\begin{split}
& z := 0 \\
& \mbox{Repeat until convergence:} \\
& \quad z := \tanh(Wz + x)
\end{split}
\end{equation}
for some network parameters $W \in \mathbb{R}^{n \times n}$.  This is an instance of a _fixed point iteration_: under certain conditions the procedure will converge to some fixed output $z^\star$, which of course has the property that
\begin{equation}
z^\star = \tanh(W z^\star + x).
\end{equation}
We'll delay for now, any discussion of _why_ this might be a particularly nice form for a layer to take, but briefly say that this type of layer can e.g., be interpreted as a simple recurrent network where $z$ is the hidden layer, and where we repeatedly apply the network to the _same_ input $x$.  A layer like this can also, e.g., reap some of the benefits of a "deep" neural network (due to the fact that it involves repeatedly application of a nonlinearity), while only having the parameters $W$ of a "single" layer.  But we will discuss these advantages layer, and for now focus on simply using a layer such as this.

Note that, of course, this is can be written as an implicit layer in the form above, i.e., that $z^\star$ is the solution to the root-finding equation
\begin{equation}
\mbox{Find $z$ such that } g(x,z) = 0, \quad \mbox{where } g(x,z) \equiv z - \tanh(W z + x).
\end{equation}
Note that this iteration need not actually converge: although the $\tanh$ activation will enforce that the values of $z$ never leave the range $[-1,+1]$, depending on the value of $W$ it could be that e.g., the values cycle endlessly and never reach a fixed point.  On the other hand, if e.g., $W=0$, then the iteration reaches a "fixed point" $z^\star = \tanh(x)$ after a single iteration.  All that we will say here is that for "typical" values of $W$ (read as: the default values of a linear layer used by most deep learning library, plus the values that they reach over optimization), this iteration will indeed converge, and we'll cover the issues of existence and uniqueness of fixed points later.

### Implementing a fixed point iteration layer

Implicit layers certainly take a bit more effort to implement, compared to traditional layers, within an automatic differentiation library like PyTorch or JAX.  But the actual core of the implementation is still quite straightforward, and still made much easier via these tools.

To start, let's consider the simplest possible implementation of a layer like this, which simply repeats the fixed point iteration to converge, all via the normal autograd functionality of the library (i.e., we are just "unrolling" the fixed point computation).  Becaues this is happening via the normal autograd mechanisms, each intermediate iterate has to be stored in memory, and the backward pass must proceed similarly over the same iterations but in the reverse order.  For now, we'll make a single layer that implements exactly the composition of the $\tanh$ and linear layer above (plus with other hacks for simplicitly, like storing the most recent iteration count and error), but in later chapters we'll make this much more modular so that we can find similar fixed points of generic layers implemented using the same library.

```
import torch
import torch.nn as nn

class TanhFixedPointLayer(nn.Module):
    def __init__(self, out_features, tol = 1e-4, max_iter=50):
        super().__init__()
        self.linear = nn.Linear(out_features, out_features, bias=False)
        self.tol = tol
        self.max_iter = max_iter

    def forward(self, x):
        # initialize output z to be zero
        z = torch.zeros_like(x)
        self.iterations = 0

        # iterate until convergence
        while self.iterations < self.max_iter:
            z_next = torch.tanh(self.linear(z) + x)
            self.err = torch.norm(z - z_next)
            z = z_next
            self.iterations += 1
            if self.err < self.tol:
                break

        return z
```

We can run this layer on random output, to see that it does in fact reach a fixed point.

```
layer = TanhFixedPointLayer(50)
X = torch.randn(10,50)
Z = layer(X)
print(f"Terminated after {layer.iterations} iterations with error {layer.err}")
```

    Terminated after 14 iterations with error 5.921083356952295e-05

Although it may ultimately not be _that_ much more informative that simply running the layer on random data, it's a bit more interesting if we can see the layer used within a real model.  So with this in mind, we'll also present below a simple model trained on the MNIST dataset, using a single fixed point layer (with an additional linear input layer before the fixed point layer, and linear layer after the fixed point layer).  The model isn't intended to break any records here, but it provides a slightly more useful example upon which to base further experiments than just running the layer in isolation.

```
# import the MNIST dataset and data loaders
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

mnist_train = datasets.MNIST(".", train=True, download=True, transform=transforms.ToTensor())
mnist_test = datasets.MNIST(".", train=False, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(mnist_train, batch_size = 100, shuffle=True)
test_loader = DataLoader(mnist_test, batch_size = 100, shuffle=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
```

```
# construct the simple model with fixed point layer
import torch.optim as optim

torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(),
                      nn.Linear(784, 100),
                      TanhFixedPointLayer(100, max_iter=200),
                      nn.Linear(100, 10)
                      ).to(device)
opt = optim.SGD(model.parameters(), lr=1e-1)
```

```
# a generic function for running a single epoch (training or evaluation)
def epoch(loader, model, opt=None, monitor=None):
    total_loss, total_err, total_monitor = 0.,0.,0.
    model.eval() if opt is None else model.train()
    for X,y in loader:
        X,y = X.to(device), y.to(device)
        yp = model(X)
        loss = nn.CrossEntropyLoss()(yp,y)
        if opt:
            opt.zero_grad()
            loss.backward()
            if sum(torch.sum(torch.isnan(p.grad)) for p in model.parameters()) == 0:
              opt.step()

        total_err += (yp.max(dim=1)[1] != y).sum().item()
        total_loss += loss.item() * X.shape[0]
        if monitor is not None:
            total_monitor += monitor(model)
    return total_err / len(loader.dataset), total_loss / len(loader.dataset), total_monitor / len(loader)
```

Let's finally train the model for 10 epochs.  In  addition to the train/test error/loss, we'll also print the average number of fixed point iteration that out layer required to converge to a fixed point.

```
for i in range(10):
    if i == 5:
        opt.param_groups[0]["lr"] = 1e-2

    train_err, train_loss, train_fpiter = epoch(train_loader, model, opt, lambda x : x[2].iterations)
    test_err, test_loss, test_fpiter = epoch(test_loader, model, monitor = lambda x : x[2].iterations)
    print(f"Train Error: {train_err:.4f}, Loss: {train_loss:.4f}, FP Iters: {train_fpiter:.2f} | " +
          f"Test Error: {test_err:.4f}, Loss: {test_loss:.4f}, FP Iters: {test_fpiter:.2f}")
```

Again, not exactly breaking any records here (a single hidden layer network achieves the same performance, with much faster execution/training), but it's nice that the network at least trains with this layer.  There are a few things to notice, however.  The first one is that we end up running the fixed point iteration for a fairly large number of layers in order to get convergence to within $10^{-4}$ of a fixed point.  And if you were the look at the individual iterations required for each minibatch, you'd see that several of them don't even reach this tolerance at all, but exit after 200 steps at a lower tolerance level (it's possible that the fixed point iteration can even go unstable at some points during training, and this typically degrades the model significantly if it does happen without any error handling in place).  This seems to be a pretty substantial downside.  We are effectively running a 50-80 "layer" network in practice, and not seeing much advantage over a standard MLP (not that, because we re-add the input to our layer at each iteration $z := \tanh(Wz + x)$, this is _not_ the same as a traditional MLP of this depth, which would suffer from vashing/exploding gradients).

To really see the potential upside of these layers, then, we need to introduce a few more ideas.

### Alternative root finding techqniues

Recall that one benefit of implicit layers is that they provide a separation between _what_ is computed by the layer, and _how_ the layer computes this.  In the above example, our goal of the fixed point iteration was to find some $z$ such that
\begin{equation}
z = \tanh(W z + x).
\end{equation}
One way to do this is to simply iterate this equation, but it is by no means the only way.  Alternatively, we can employ a much faster root-finding method, such as Newton's method, to try to find this solution more efficiently.

Newton's method is a generic root-solving technique.  For some function $g : \mathbb{R}^n \rightarrow \mathbb{R}^n$, if we wish to find a root $g(z) = 0$, then Newton's method repeats the update
\begin{equation}
z := z - \left ( \frac{\partial g}{\partial z} \right ) ^{-1} g(z)
\end{equation}
where $\frac{\partial g}{\partial z}$ denotes the Jacobian of $f$ with respect to $z$ (a "guarded" update is often required in practice that makes smaller steps to ensure sufficient decrease of the residual $\|g(z)\|$, but we won't consider this here).  Although we could resort to automatic differentiation to compute the Jacobians (and we'll need to do this in later chapters, when we have a more generic layer we're using within fixed point iteration), for the case of our $\tanh$ plus lineaer layer, it's easy to compute the Jacobian in closed form.  Specifically, we are attempting to find the root of the equation $g(x,z) = 0$ (returning to the notation from the previous sections of this chapter, where we make the dependence on layer input $x$ explicit), where
\begin{equation}
g(x,z) = z - \tanh(Wz + x).
\end{equation}
Then our Jacobian is given by
\begin{equation}
\frac{\partial g}{\partial z} = I - \mathrm{diag}(\tanh'(Wz + x)) W
\end{equation}
where $\tanh'$ denotes the derivative of the $\tanh$ function, given by
\begin{equation}
\tanh'(x) = \mathrm{sech}^2(x).
\end{equation}

Let's see what an implementation of Newton's method looks like in code.  The implementation is slightly more involved than the simple fixed point iteration, owing to the need to compute the Newton step, but this ends up being only a few additional lines of code.

```
class TanhNewtonLayer(nn.Module):
    def __init__(self, out_features, tol = 1e-4, max_iter=50):
        super().__init__()
        self.linear = nn.Linear(out_features, out_features, bias=False)
        self.tol = tol
        self.max_iter = max_iter

    def forward(self, x):
        # initialize output z to be zero
        z = torch.tanh(x)
        self.iterations = 0

        # iterate until convergence
        while self.iterations < self.max_iter:
            z_linear = self.linear(z) + x
            g = z - torch.tanh(z_linear)
            self.err = torch.norm(g)
            if self.err < self.tol:
                break

            # newton step
            J = torch.eye(z.shape[1])[None,:,:] - (1 / torch.cosh(z_linear)**2)[:,:,None]*self.linear.weight[None,:,:]
            z = z - torch.solve(g[:,:,None], J)[0][:,:,0]
            self.iterations += 1

        g = z - torch.tanh(self.linear(z) + x)
        z[torch.norm(g,dim=1) > self.tol,:] = 0
        return z
```

```
layer = TanhNewtonLayer(50)
X = torch.randn(10,50)
Z = layer(X)
print(f"Terminated after {layer.iterations} iterations with error {layer.err}")
```

    Terminated after 3 iterations with error 1.2266763178558904e-06

This is able to nicely coverge more quickly than the fixed point iteration, but with the (major) caveat that we now have to solve a linear system at each iteration.  And again, since we're implementing this whole procedure using automatic differentiation, means we also need to backprop through the solves in the backward pass.  But we can simply plug this into the same training process as before.

```
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(),
                      nn.Linear(784, 100),
                      TanhNewtonLayer(100, max_iter=40),
                      nn.Linear(100, 10)
                      ).to(device)
opt = optim.SGD(model.parameters(), lr=1e-1)

for i in range(8):
    if i == 5:
        opt.param_groups[0]["lr"] = 1e-2

    train_err, train_loss, train_fpiter = epoch(train_loader, model, opt, lambda x : x[2].iterations)
    test_err, test_loss, test_fpiter = epoch(test_loader, model, monitor = lambda x : x[2].iterations)
    print(f"Train Error: {train_err:.4f}, Loss: {train_loss:.4f}, Newton Iters: {train_fpiter:.2f} | " +
          f"Test Error: {test_err:.4f}, Loss: {test_loss:.4f}, Newton Iters: {test_fpiter:.2f}")
```

```
Train Error: 0.1156, Loss: 0.4131, FP Iters: 6.46 | Test Error: 0.0633, Loss: 0.2117, FP Iters: 7.03
Train Error: 0.0610, Loss: 0.2065, FP Iters: 6.98 | Test Error: 0.0490, Loss: 0.1681, FP Iters: 6.78
Train Error: 0.0463, Loss: 0.1555, FP Iters: 6.73 | Test Error: 0.0438, Loss: 0.1483, FP Iters: 6.04
Train Error: 0.0380, Loss: 0.1280, FP Iters: 7.01 | Test Error: 0.0373, Loss: 0.1276, FP Iters: 6.28
Train Error: 0.0328, Loss: 0.1087, FP Iters: 7.29 | Test Error: 0.0332, Loss: 0.1167, FP Iters: 7.06
Train Error: 0.0238, Loss: 0.0817, FP Iters: 7.41 | Test Error: 0.0323, Loss: 0.1134, FP Iters: 7.63
Train Error: 0.0217, Loss: 0.0764, FP Iters: 8.07 | Test Error: 0.0311, Loss: 0.1064, FP Iters: 7.74
Train Error: 0.0201, Loss: 0.0719, FP Iters: 8.69 | Test Error: 0.0298, Loss: 0.1057, FP Iters: 8.28
```

Again, the method works reasonably well.  However, there are some notable issues with the approach as implemented.  The first, as you will immediately notice if you run the code, is that the approach is noticably _slower_ than the the simpler fixed point iteration method above.  Even though the number of iterations needed is much smaller that for fixed point iteration, each individual iteration is also much slower, as it involves forming and inverting a separate (in this case, $100 \times 100$) Jacobian matrix for each sample in the minibatch.  And for larger hidden unit sizes (especially, e.g. for convolutional networks), it will quickly become intractable to invert or even store these matrices.  And indeed, in pratice, an exact Newton method is rarely used, and instead we can employ quasi-Newton methods to improve convergence over the standard fixed point iteration, while also improving wall clock time.

The second issue with this approach is a bit more subtle, but actually an even larger problem.  Because we implemented Newton's method directly within an automatic differentiation toolkit, there are a few large downsides to this method as presented.  First, as with the fixed point iteration, the automatic differentiation tool will need to save the intermediate iterates of the hidden units; but here, this means that we _also_ need to store in memory the intermediate iterates of the Jacobian terms as well, which drastically increases memory consumption even for the case where we _could_ store and invert the full Jacobian.  Furthermore, backpropagation through repeated inverses can be a numerically unstable routine: if an inverse is close to singular then even if the forward pass converges properly, the backward pass can still generate numerical errors in the graidents.  Indeed, you will notice that we included a "NaN check" in our `epoch()`.  If we did _not_ do this, then for Newton's method the approach will immediately fail: if you check, you will see that around 5% of the updates actually have NaN values in the gradient, due to conditioning of the Jacobians, which is also what causes the method to actually converge _slower_ than the fixed point iteration version.

This paints a pretty dim picture for "efficient" methods for solving implicit models.  Fortunately, however, there is a much better way to implement these layers, courtesy of the implicit function theorem.

### Differentiation in implicit layers

Thus far, we implemented our solver for our implicit layer in exactly the same manner as we would implement any other layer, and let the automatic differentiation library take care of the backward pass.  However, there is a much nicer way to differentiate with respect to the fixed point of a hidden layer.  To see how to do this, let's consider the generic form of our implicit layer, namely given $x$, finding some $z$ such that
\begin{equation}
g(x,z) = 0.
\end{equation}
Let's denote $z^\star(x)$ as the value that solves this fixed point, written this way to emphasize that the output of the implicit layer is still of course an (implicit) function of the input.

Now let's consider how we can compute the Jacobian of this output with respect to the input
\begin{equation}
\frac{\partial z^\star(x)}{\partial x}.
\end{equation}
Unlike the traditional functions you're used to, where we are given an explicit form for computing the output from the input, it may not be obvious how to determine such a Jacobian.  But in fact it's very straightforward to compute the term using _implicit differentiation_, a technique that goes back several centuries in Calculus.  In particular, to derive an expression for this Jacobian, we start with the fixed point condition, which we know holds for $z^\star(x)$, and differentiate both sides with respect to $x$,
\begin{equation}
\frac{\partial g(x,z^\star(x))}{\partial x} = 0.
\end{equation}
Now we just use the chain rule to expand this partial derivative: since $g$ is a function of two variables, there will be a term involving the derivative with respect to each
\begin{equation}
\frac{\partial g(x,z^\star)}{\partial x} + \frac{\partial g(x,z^\star)}{\partial z^\star}\frac{\partial z^\star(x)}{\partial x} = 0
\end{equation}
where the notation $z^\star$ (not indicated as a funciton of $x$), just indicates means we are treating $z^\star$ as a _fixed_ value here (i.e., the Jacobian $\frac{\partial g(x,z^\star)}{\partial x}$ would just be the Jacobian of $g$ with respect to $x$, evaluated at the point $(x,z^\star)$).  Thus, this term, along with the $\frac{\partial g(x,z^\star)}{\partial z^\star}$ term, can be computed themselves using ordinary automatic differentiation libraries.  Finally, we then just rewrite this equation to give us the expression we are after in terms of the expressions we know
\begin{equation}
\frac{\partial z^\star(x)}{\partial x} = - \left ( \frac{\partial g(x,z^\star)}{\partial z^\star} \right )^{-1} \frac{\partial g(x,z^\star)}{\partial x}.
\end{equation}
Technically, in order to ensure that we can actaully apply this theorem, we require that certain conditions must be satisfied, so that the implicit function $z^\star(x)$ is guaranteed to exist: these conditions are reflected in what is known as the implicit function theorem, which will be discussed in the following chapter.  Additionally, just as with Newton vs. quasi-Newton methods, in practice it is often not possible to compute this inverse directly, but instead an iterative process is needed.  We'll cover the mathematical details and formalisms a lot more in the next chapter, but for the purposes of most of what we actually need to derive, this "informal" derivation is virtually all you'll need.  Finally, although we wrote the above formula for the Jacobian with respect to $x$, when $g$ is also a function of some parameters $\theta$ (e.g., the weights and biases), precisely the same derivation holds to find the Jacobian with respect to these parameters.

Moving back from the detailed derivation of this formula though, the implicit function theorem leads to a very practical consequence.  Namely, the formula gives a form for the necessary Jacobian _without needing to backpropagate through the method used to obtain the fixed point_.  In other words, it doesn't matter at all _how_ we compute the zero of the function (whether via fixed point iteration, Newton's method, or quasi-Newton methods).  All that matters is finding the fixed point (using whatever technique you want), at which point we can directly compute the necessary Jacobians using this analytical form (or more precisely, compute the backward pass, which often will not require explicit computation of the Jacobian in practice).  No intermediate terms of the iterative method used to compute the fixed point need to be stored in memory (making the methods much more memory efficient), and there is no need to unroll the forward computations within an automatic differentiation layer.

### Implementing implicit differentiation

Let's see how an implementation of implicit differentiation will work in practice.  First, let's consider our tanh plus linear layer again, where the $g(x,z)$ function is given by
\begin{equation}
g(x,z) = z - \tanh(Wz + x)
\end{equation}
In this case the Jacobian $\frac{\partial g}{\partial z^\star}$, needed for implicit differentation, is given by
\begin{equation}
\frac{\partial g}{\partial z^\star} = I - \mathrm{diag}(\tanh'(Wz^\star+x)) W.
\end{equation}
You may notice that this is the _exact_ same Jacobian that we formed when using Newton's method to solve for the fixed point.  This is not an accident: indeed, precisely the same Jacobian term is needed in finding a root within Newton's method as for computing the backward pass via implicit differentiation.  This results in a very nice property: forthe situation where we find a solution to our root via a method like Newton's method (or any approach that computes and inverts the Jacobian), then computing the backward pass via Newton's method is effectively "free" (at least, relative to the complexity of solving the fixed point to begin with): we can simply reuse the Jacobian (and its inverse) that we did in the forward pass.  Of course, since in practice we often use quasi-Newton or first order methods for finding fixed points of implicit layers, this is not quite as big an advantage as it may seem.  But, nonetheless, in the case where we _do_ compute even an approximation to the Jacobian during the forward pass, it can be beneficial to leverage this computation in the backward pass as well.

Before we move on to our implmentation, we should emphasize how the actual implicit differentiation process works within backpropation (i.e., reverse mode autodifferentiation).  In backpropagation, we don't actually need to compute full Jacobians of the intermediate layers in a network.  Rather, the goal of backprop is to compute the gradient with respect to some _scalar_ loss.  If we write this out in terms of our gradient above, it would look something like
\begin{equation}
\frac{\partial \ell}{\partial x} = \frac{\partial \ell}{\partial z^\star} \frac{\partial z^\star}{\partial x} = - \frac{\partial \ell}{\partial z^\star} \left (\frac{\partial g}{\partial z^\star} \right )^{-1} \frac{\partial g}{\partial x}
\end{equation}
where we applied the implicit differentiation formula above in the last equality.  In backpropagation, this term is computed left-to-right, meaning instead of actually needing to compute the full Jacobian $\frac{\partial z^\star}{\partial x}$, we just need to the compute the _vector Jacobian product_ shown above.  As a matter of convention, most automatic differentiation frameworks frame this in terms of the operations on the _gradient_ (the transpose of the Jacobian for a scalar-valued function)
\begin{equation}
\nabla_{z^\star} \ell = \left ( \frac{\partial \ell}{\partial z^\star} \right )^T
\end{equation}
so we need to multiply by the _transpose_ of the Jacobian
\begin{equation}
\nabla_x \ell = \left (\frac{\partial g}{\partial x} \right )^T \left (\frac{\partial g}{\partial z^\star} \right )^{-T} \nabla_{z^\star} \ell.
\end{equation}
Again, however, we emphasize that we don't actually need to store and compute the actual inverse $\left (\frac{\partial g}{\partial z^\star} \right )^{-T}$, just be able to solve the (linear) equation that arises in this formula.

Finally, let's discuss how we can implement a formula like this within an automatic differentiation toolkit.  The details will of course vary from framework to framework, but since we're ultimately talking about implementing a _new_ type of function here (i.e., in the forward pass computing a fixed point outside of any automatic differentiation, and then computing a "custom" backward pass) you may be tempted to use a feature like the `autograd.Function` interface (if you were implementing this in PyTorch, for instance), which lets you specify a forward and backward pass entirely outside the normal automatic differentiation pass of the library.  But this would actually be a bit cumbersome in practice: after all, one of the benefits of automatic differentiation is that we could potentially implement the function $g$ (whether it used convolutions, self-attention, or any other feature) inside the same automatic differentiation library, and we would to automatically include all these gradients without writing a new function for each particular function $g$ we want to implement.  Fortunately, there is a fairly straightforward though subtle way to deal with this issue.  We'll return to several examples of
efficient implicit differentiation in later sections, which each have their own implementation quirks, but for a simple example like this, a common paradigm that works is the following three steps:
1. _Outside_ the automatic differentiation tape, solve for the root of the implicit layer $g(x,z^\star) = 0$.
2. "Reengage" the automatic differentiation by running the following assignment _within_ the automatic differentation tape:
\begin{equation}
z := z^\star - g(x,z^\star).
\end{equation}
This has the effect of "reinserting" the partial derivatives $-\frac{\partial g}{\partial x}$ to the autograd tape (and is a no-op in terms of the value of $z$, since $g(x,z^\star) = 0$ ).
3. Add a "backward hook" to the backward pass that multiplies by $(\frac{\partial g}{\partial z^\star})^{-T}$.  This will fix the backward pass so that it correctly implements the gradient according to the implicit function theorem.

For the tanh + linear layer from before, this results in an implementation like that below.  Notice that the layer is essentially identical to the version we implemented previously, except that Newton's method runs within a `torch.no_grad():` block, and we add the backward pass hook via the `register_hook` function.  For the second step above, given the $g$ function as highlighted before, the assignment is simply
\begin{equation}
z := z^\star - g(x,z^\star) = z^\star - z^\star + \tanh(Wz + x) = \tanh(Wz^\star + x)
\end{equation}
i.e., we run a single fixed point iteration, within the automatic differentiation tape, after finding the fixed point with Newton's method.

```
class TanhNewtonImplicitLayer(nn.Module):
    def __init__(self, out_features, tol = 1e-4, max_iter=50):
        super().__init__()
        self.linear = nn.Linear(out_features, out_features, bias=False)
        self.tol = tol
        self.max_iter = max_iter

    def forward(self, x):
        # Run Newton's method outside of the autograd framework
        with torch.no_grad():
            z = torch.tanh(x)
            self.iterations = 0
            while self.iterations < self.max_iter:
                z_linear = self.linear(z) + x
                g = z - torch.tanh(z_linear)
                self.err = torch.norm(g)
                if self.err < self.tol:
                    break

                # newton step
                J = torch.eye(z.shape[1])[None,:,:] - (1 / torch.cosh(z_linear)**2)[:,:,None]*self.linear.weight[None,:,:]
                z = z - torch.solve(g[:,:,None], J)[0][:,:,0]
                self.iterations += 1

        # reengage autograd and add the gradient hook
        z = torch.tanh(self.linear(z) + x)
        z.register_hook(lambda grad : torch.solve(grad[:,:,None], J.transpose(1,2))[0][:,:,0])
        return z
```

Note that this is a fairly non-standard implementation: we are implementing an element of the forward pass _outside_ the normal automatic differentiation tape, and then adding a backwards hood to "fix" the gradient.  We can verify the correctness of this layer using the building `gradcheck` command.  Note that this implmentation won't work for double backprop (i.e., `gradgradcheck` will not work), but this can be addressed by a slightly more involved approach, and often isn't needed in practice, so we ignore it for now.

```
from torch.autograd import gradcheck

layer = TanhNewtonImplicitLayer(5, tol=1e-10).double()
gradcheck(layer, torch.randn(3, 5, requires_grad=True, dtype=torch.double), check_undefined_grad=False)
```

Finally, again for demonstration, we'll train our MNIST network with this new variant of the implicit layer.  As hoped, the method is indeed quite a bit faster _and_ substantially more stable than the previous implementation of Newton's method.  And while we again highlight that using Newton's method exactly isn't typically a reasonable approach for settings like these, somemthing very similar will in fact be quite useful when we discuss differentiable optimimzation in later chapters.

```
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(),
                      nn.Linear(784, 100),
                      TanhNewtonImplicitLayer(100, max_iter=40),
                      nn.Linear(100, 10)
                      ).to(device)
opt = optim.SGD(model.parameters(), lr=1e-1)

for i in range(10):
    if i == 5:
        opt.param_groups[0]["lr"] = 1e-2

    train_err, train_loss, train_fpiter = epoch(train_loader, model, opt, lambda x : x[2].iterations)
    test_err, test_loss, test_fpiter = epoch(test_loader, model, monitor = lambda x : x[2].iterations)
    print(f"Train Error: {train_err:.4f}, Loss: {train_loss:.4f}, Newton Iters: {train_fpiter:.2f} | " +
          f"Test Error: {test_err:.4f}, Loss: {test_loss:.4f}, Newton Iters: {test_fpiter:.2f}")
```

```
Train Error: 0.1130, Loss: 0.4061, Newton Iters: 6.48 | Test Error: 0.0605, Loss: 0.2040, Newton Iters: 6.92
Train Error: 0.0577, Loss: 0.1949, Newton Iters: 6.95 | Test Error: 0.0487, Loss: 0.1660, Newton Iters: 6.67
Train Error: 0.0449, Loss: 0.1500, Newton Iters: 6.76 | Test Error: 0.0412, Loss: 0.1398, Newton Iters: 5.79
Train Error: 0.0369, Loss: 0.1236, Newton Iters: 6.68 | Test Error: 0.0374, Loss: 0.1238, Newton Iters: 6.63
Train Error: 0.0316, Loss: 0.1056, Newton Iters: 7.25 | Test Error: 0.0346, Loss: 0.1145, Newton Iters: 7.12
Train Error: 0.0214, Loss: 0.0736, Newton Iters: 7.43 | Test Error: 0.0308, Loss: 0.1018, Newton Iters: 7.30
Train Error: 0.0191, Loss: 0.0682, Newton Iters: 7.84 | Test Error: 0.0285, Loss: 0.0989, Newton Iters: 7.79
Train Error: 0.0182, Loss: 0.0651, Newton Iters: 8.32 | Test Error: 0.0288, Loss: 0.0992, Newton Iters: 8.12
Train Error: 0.0179, Loss: 0.0626, Newton Iters: 9.07 | Test Error: 0.0288, Loss: 0.1001, Newton Iters: 8.64
Train Error: 0.0171, Loss: 0.0607, Newton Iters: 9.76 | Test Error: 0.0277, Loss: 0.0990, Newton Iters: 9.29
```

## Final chapter remarks

Before delving into the much more realistic and varied world of practical implicit models, we want to highlight our accomplishments so far.  Using very little additional code from a "traditional" deep model (and definitely not that much more than a traditional recurrent model), we're able to code a layer that 1) solves a non-linear root-finding problem via Newton's method, equivalent to finding the fixed-point of an infinite-depth network, and 2) integrates easily into automatic differentiation tools.  The relative ease of these approaches, once your push past a bit of the mathematical notation of implicit differentiation, is indeed one of the more compelling factors of using implicit layers within deep learning as a whole.

Throughout the rest of this tutorial, we will provide you with the tools and background you need to apply implicit layers to a wide variety of problems and settings, with code examples throughout.  Our hope is that this will enable readers to quickly integrate and make progress in this new directions.


# Variational Auto Encoder in pyTorch for MNIST dataset

## Overview

The basic idea of a Variational Auto Encoder (VAE) is shown below:

Hopefully, this cartoon should convey a couple of important ideas:
1. the VAE is composed of a (probabilistic) Encoder and a (probabilistic) Decoder which do NOT need to have the same architecture (represented by boxes of different color and shapes)
2. the output of the VAE is probabilistic. So it is often usefull to visualize the expected (i.e. averaged) output.
3. the input and the expected output of the VAE are NOT identical (the difference between input and output is greatly exagerated just for making the poiunt clear)

The encoder compress the input $x$ into a latent variable $z$ which lives in small dimensions latent space. The decoder tries to recover the original image from the latent space representation.

We now zoom-in in the central part of the figure and clarify what we mean by **probabilistic encoder**.

The important feature is that the latent variable $z=\mu+\sigma \epsilon$ is a random variable since $\epsilon$ is itself a random variable distribuited according to a normal distribution $p(\epsilon)=N(\epsilon|0,I)$. This means that each input vector $x$ can be mapped into many differrent values of $z$, i.e. we have a **probabilistic encoder**. Fortunately, these value of $z$ are all close to each other (i.e. they are centered around $\mu$ and have a spreading of size $\sigma$).

From the figure above we can easily read off the core of the pyTorch implementation of the VAE:
mu_z, sigma_z = encoder(x_input)
ep = np.random.normal(0, 1, size=sigma_z.size())
z = mu_z + sigma_z * ep
x_output = decoder(z)
We will train the VAE by demanding that the reconstruction error is small (i.e. the input $x$ and output $\tilde{x}$ are close to each other according to some proper metric) This requirement, together with the use of a **probabilstic encoder**, imposes a nice structure on the latent space. To understand this imagine that two inputs $x_1$ and $x_2$ representing two different digits (for example a $0$ and a $7$) were mapped to similar values of $\mu$ and $\sigma$ by the encoder. Then we would calculate:

$$z_1=\mu_1+\sigma_1 \epsilon_1,\quad\text{and}\quad z_2=\mu_2+\sigma_2 \epsilon_2$$

Since $\epsilon\sim N(0,I)$ is a random variable it is possible that $z_1=z_2$ (or more precisely that $z_1$ and $z_2$ are arbitrarely close to each other). Then when we feed this into the decoder, it would be impossible for it to decide wheather to generate a $0$ or a $7$ and the recontruction error would be (on average) very large. The only way to avoid this problem is for the encoder to learn how to map different inputs (for example a $0$ and a $7$) into different regions of the latent space and conversely to map similar inputs (for example two different version of a $0$) into close-by regions in the latent space.

We will now describe in detail how this works. As an example we will use the MNIST dataset of handwritten digits. The VAE will take as input 28 x 28 gray scale images and will simultaneously learn:
1. to compress the input vector $x$ in $D_x=728$ dimensions (obtained by flattening the $28 \times 28$ images) into a hidden variables $z$ in a low dimensional space, for example $D_z=2$.
2. to generate a $D_x$ dimensional output vector $\tilde{x}$ for each particular value of the hidden variable $z$.

As we have already seen, the VAE is composed of a **stochastic encoder** which maps $x\to z$ and a **stochasitc decoder** which maps $z \to \tilde{x}$. We will train the VAE by:
1. demanding that the reconstruction error is small, i.e. $\tilde{x}$ and $x$ are similar to each other (according to a proper metric)
2. enforcing a suitable structure in the latent space. This is important since we eventually want to sample the latent space and use the decoder to generate new (never seen before) images. This is only possible if we know how to meaningfully sample the latent space which in turns requires that the latent space has some "structure".

## Mathematical Preliminaries

We start by showing three mathematical identities:

\begin{equation}
\log p(X)=\int dZ q(Z)\log p(X) = \int dZ q(Z) \log \frac{p(X,Z)}{p(Z|X)}  \\
= \int dZ q(Z) \log \frac{p(X,Z)}{q(Z)} + \int dZ q(Z) \log \frac{q(Z)}{p(Z|X)} \\
= ELBO + KL(q(Z)\,||\,p(Z|X))
\end{equation}

where $KL(a(.)\,||\,b(.))$ is the KL divergence between the two distribution $a(.)$ and $b(.)$ which is always positive and becomes zero only when $a(.)=b(.)$. From equation (1) and the positivity of the KL divergence it immediately follows that:

\begin{equation}
\log p(X) \ge ELBO
\end{equation}

We can rewrite ELBO in two ways:

\begin{equation}
ELBO \equiv \int \,dz \,q(Z)\,\log \frac{p(X,Z)}{q(Z)} = E_{q(Z)} [\log p(X,Z) - \log q(Z)] \\
ELBO \equiv \int \,dz \,q(Z)\,\log \frac{p(X|Z)p(Z)}{q(Z)} = E_{q(Z)} [\log p(X|Z)] - KL(q(Z)\,||\,p(Z))
\end{equation}

At this point, the equations 1,2,3 are not very usefull and might seem a little arbitrary. Interestingly the equations above remain true even when the distributions depend on some parameters. When this is the case these equations can be used to:
1. to find the value of the parameters which maximize the probability of observing the data (this is known as Expectation-Maximization algorithm) or
2. to approximate the true posterior distribution in terms of a variational family of distribution (Bayesian Variational Inference).

## Probabilistic Decoder (i.e. $z\to \tilde{x}$)

The task of the probabilistic decoder is to take a hidden variable $z$ in $D_z$ dimension and transform it into an output vector $\tilde{x}$ in $D_x$ dimension. In mathematical terms the probabilistic decoder is defined by:

\begin{equation}
p_\theta(\tilde{x})=p_\theta(\tilde{x}|z)p(z)
\end{equation}

where $p(z)$ is the prior for the hidden variable $z$ which is chosen to be a multidimensional normal distribution with zero average and unit variance:

\begin{equation}
p(z)=N(z|0,I),\quad\text{where $I$ is the identity matrix in $D_z$ dimension}
\end{equation}

and $p_\theta(\tilde{x}|z)$ is the likelyhood which descibes the probability, under the current model, that the hidden variables $z$ generate the output $\tilde{x}$. This probability depends on the model under consideration. For us the model is a Neural Network (NN). The parameters of the NN are collectively indicated by $\theta$. In the equation above, the term on the LHS is known as the EVIDENCE which is the probability of observing the data under the current model. Clearly we should optimize the parameters $\theta$ so that the EVIDENCE is maximized (more of this later).

To be concrete, we will use a NN with a single hidden layer. Therefore $p_\theta(\tilde{x}|z)$ is built explicitely in the following way:

\begin{equation}
a_2 = relu(W_1 z + b_1) \\
a_3 = sigmoid(W_2 a_2 + b_2) \\
p(\tilde{x}|a_3) = a_3^\tilde{x}(1-a_3)^{(1-\tilde{x})}
\end{equation}

In other words, the latent variable $z$ is fed into a NN with two dense layers. The output of the NN is a vector $a_3$ containing $D_x$ probabilities (i.e. the dimensionality of the MNIST images). Each entry in $a_3$ bescribes the probability that a particular pixel in the image $\tilde{x}$ is black $\tilde{x}_i=0$ or white $\tilde{x}_i=1$. Clearly this is a Bernulli distribution therefore the expected value of $\tilde{x}$ is:

$$ E[\tilde{x}] = \sum_{\tilde{x}=0,1} \tilde{x} p(\tilde{x}|a_3) = 0\times(1-a_3)+1\times a_3 = a_3 $$

Therefore to visualize the average output we can simply image the probabilities $a_3$.
The parameters of the decoder NN are $\theta=\{W_1,b_1,W_2,b_2\}$.

As a side, we note that:

\begin{equation}
\log p(\tilde{x}|a_3) = \sum_i \tilde{x}_i \log (a_3)_i + (1-\tilde{x}_i) \log (1-(a_3)_i)
\end{equation}

is just the binary cross entropy between $\tilde{x}$ and $a_3$.

## Probabilistic Encoder (i.e. $x\to z$)

The task of the probabilistic encoder is to take a input vector $x$ in $D_x$ dimension and compress it into the hidden variables $z$ in $D_z$ dimensions. Mathematically the probabilistic encoder is defined by the posterior $p_\theta(z|x)$ which describes the probability that an input vector $x$ is mapped to a particular value $z$ in the latent space. In principle the posterior can be learned using the Baysian framework:

\begin{equation}
p_\theta(z|x)=\frac{p_\theta(x|z)p(z)}{p_\theta(x)}
\end{equation}

where all the term on the RHS have been introduced when discussing the probabilistic decoder. The expression above is the mathematically way of expressing the fact that the encoder and decoder are complementary to each other. For the problem at hand, the equation above is intractable (both analytically and computationally). This is analytically intractable because we do not have an analitical form for $p_\theta(x|z)$ since this is encoded in a NN. Moreover this is computationally intractable since in the EVIDENCE ($p_\theta(x)=\int dz p_\theta(x|z)p(z)$) which appears in the denominator is exponentially hard to compute in the dimension of the hidden space $D_z$.

We are therefore force to use Approximate Bayesian Inference. We approximate the true posterior $p_\theta(z|x)$ by a variational family $q_\phi(z|x)$. Here we restrict ourself to the family of Gaussian posterior $q(z|x)=N(z\,|\,\mu(x),diag(\sigma(x)))$. In principle, for each $x$ we would need to learn $\mu(x)$ and $\sigma(x)$. However, in this case the number of parameters scales with the dataset size which is a BAD idea. Instead we use an inference network descibed by the parameters $\phi$ which will learn how to map $x$ to $\mu(x)$ and $\sigma(x)$. This is called **amortization**.

To be concrete we choose the inference network to be a NN with a single hidden layer:

\begin{equation}
a_x=relu(W_x x + b_x) \\
\mu_z = W_\mu a_x + b_\mu \\
\log\sigma_z = W_\sigma a_x + b_\sigma
\end{equation}

In other words, the input variable $x$  is fed into a NN with two dense layers. The output of the NN are two vectors of dimension $D_z$. One containts the average and the other the variance (actually the log variance) of the Normal distribution. The parameters of this inference network are $\phi=\{W_x,b_x,W_\mu,b_\mu,W_\sigma,b_\sigma\}$. The output of the inference network are $\mu_\phi(x),\sigma_\phi(x)$ which are used to build the approximate posterior:

\begin{equation}
q_\phi(z|x)=N(z\,|\,\mu_\phi(x),diag(\sigma_\phi(x)))
\end{equation}

## Loss Function

We now rewrite the inital three equations making sure to keep track of all the parameters we have introduced thus far:

\begin{equation}
\log p_\theta(x)= ELBO(\theta,\phi) + KL(q_\phi(z|x)\,||\,p_\theta(z|x))
\end{equation}
\begin{equation}
\log p_\theta(x) \ge ELBO(\theta,\phi)
\end{equation}
\begin{equation}
ELBO(\theta,\phi) = E_{q_\phi(z|x)} [\log p_\theta(x|z)] - KL(q_\phi(z|x)\,||\,p(z))
\end{equation}

As a recap:
1. $p(z)=N(z | 0,I)$ is the prior for the hidden variables
2. $p_\theta(x|z)$ and $p_\theta(x)$ are the "likelihood" and "evidence" which depend on the parameters $\theta=\{W_1,b_1,W_2,b_2\}$ of the decoder NN.
3. $p_\theta(z|x)$ is the true posterior which is not accessible to us (i.e. the problem is intractable)
3. $q_\phi(z|x)$ is the approximate posterior which is given by gaussian family $q_\phi(z|x)=N(z\,|\,\mu_\phi(x),\sigma_\phi(x))$ where $\mu_\phi(x)$ and $\sigma_\phi(x)$ are the output of the inference NN which depend on the parameters $\phi=\{W_x,b_x,W_\mu,b_\mu,W_\sigma,b_\sigma\}$.

Our goal is now to maximize the evidence (or equivalently the log-evidence since the log is a monothonic function), i.e. maximize the probability of observing the data under the current model. Unfortunately the log-evidence depends on the true posterior (see Eq. 11) which is unaccessible to us. Fortunately, we have shown that the ELBO is a lower bound for the evidence (this justifies its name) therefore we decide to maximize the ELBO as a proxi for the maximization of the evidence. The ELBO depends on the parameters $\theta$ in the decoder NN and $\phi$ in the inference NN. Note that:
1. when we maximize ELBO w.r.t. $\phi$ we are reducing $KL(q_\phi(z|x)\,||\,p_\theta(z|x))$ (this can be seen by noticing that the LHS in Eq. 11 is independent on $\phi$ therefore increasing ELBO must be accompanied by a reduction in KL) and therefore making the approximate posterior more similar to the true posterior
2. when we maximize ELBO w.r.t. $\theta$ we are most likely pushing the evidence higher and therefore making the model more tuned to the observed data.

Therefore maximizing the ELBO with respect to both $\theta$ and $\phi$ is a very desirable thing to do.

In a Machine Learning framework we are usually trying to minimize the loss function. If we are interested in maximizing the ELBO we can simply set:

\begin{equation}
Loss=-ELBO(\theta,\phi)=Loss_1+Loss_2
\end{equation}

where $Loss_1$ and $Loss_2$ correspond to the two terms (the KL and the Expected Log Likelyhood) in the ELBO. Conceptually, we are done. Practically, we still need to "massage" the ELBO expression a bit in order to write it as a simple analytical function of the outputs of the decoder NN and inference NN. This is a necessary step if we want to use modern ML framework (such as Keras, pyTorch etc) which need to be able to automatically compute the derivatives of the Loss function and backpropagate to update the parameters $\phi,\theta$.

Let us start with the KL divergence term in the ELBO. Since both the prior $p(z)=N(z\,|\,0,I)$ and the approximate posterior $q_\phi(z|x)=N(z\,|\,\mu_\phi(x),\sigma_\phi(x))$ are gaussian then it is easy to compute the KL term analytically:

$$ Loss_1 = KL(q_\phi(Z|X)\,||\,p(Z)) = -\frac{1}{2} \sum_{k=1}^{D_z} \left(1+\log\sigma^2(x)_k- \mu(x)^2_k - \sigma^2(x)_k \right) $$

Note that this expression above is a simple (analytical) function of outputs of the inference network ($\mu_\phi(x),\sigma_\phi(x)$) therefore this expression is ready for use in ML framework based on backpropagation. **This part of the Loss function encorauges the posterior distribution to be close to the prior and can be interpreted as a regularization term which prevent the parameters $\mu_\phi(x)$ and $\sigma_\phi(x)$ from becoming to different from $0$ and $1$ respectively.**

The second term in the ELBO is the Expected Log Likelyhood (ELL):

$$Loss_2 = - ELL(\theta,\phi) = - E_{q_\phi(z|x)} [\log p_\theta(x|z)] = - \int dz\,q_\phi(z|x)\,\log p_\theta(x|z) $$

This expression is **not** simple enough to be used directly for backpropagation. Note that this expression involves an integral over the distribution $q_\phi(z|x)$ which depend on the parameter $\phi$ which we are trying to optimize. Interestingly, the expression above can be rewritten exactly as:

$$ Loss_2= - \int dz\,q_\phi(z|x)\,\log p_\theta(x|z) = - \int d\epsilon\,N(\epsilon|0,I)\,\log p_\theta(x|\mu_\phi(x)+\sigma_\phi(x) \epsilon)$$

The equivalence between the two expressions can be seen by thinking about the "Monte-Carlo way" of computing these integrals. The first expression can be computed by:
1. drawing multiple $z$ from the distribution $q_\phi(z|x)=N(z\,|\,\mu_\phi(x),\sigma_\phi(x))$
2. then using those values to compute $\log p_\theta(x|z)$.

The second expression can be computed by:
1. draw multiple $\epsilon$ from the distribution $N(0,I)$
2. use those value of $\epsilon$ to compute $z=\mu_\phi(x)+\sigma_\phi(x) \epsilon$
3. compute $\log p_\theta(x|z)$

It is easy to see that, if $\epsilon\sim N(0,I)$ then $z=\mu_\phi(x)+\sigma_\phi(x)$ is drawn from $N(z\,|\,\mu_\phi(x),\sigma_\phi(x))=q_\phi(z|x)$ and therefore steps 1 and 2 are equivalent to draw $z$ from the distribution $q_\phi(z|x)$. This means that the two expressions for $Loss_2$ are equivalent.

The final step is to approximate the integral over $d\epsilon$ with a single Monte Carlo sample, i.e.:

$$ Loss_2 \approx - \log p_\theta\left(x|\mu_\phi(x)+\sigma_\phi(x) \epsilon\right),\quad\text{where}\quad \epsilon\sim N(0,I)$$

We have finally reached an expression simple enough to be used in ML framework based on backpropagation. **The ELL term in the Loss function is the reconstruction error. It encourages the encoder and the decoder to be inverse of each other.** To see this imagine that a region in the latent space is mapped to a region in the original space , i.e. $p_\theta(x|z)$ is high among those two regions. The $Loss_2$ will be small if  $q_\phi(z|x)$ is high among those same two regions. So we can say that, at the end of training, the mapping from $z\to x$ described by $p_\theta(x|z)$ is opposite to the mapping from $x\to z$ described by $q_\phi(z|x)$.

### Raparametrization Trick

More formally, the manipulation that we have done in the ELL, is known as reparametrization trick. The reparameterization trick is a straightforward change of variables that expresses the random variable $z\sim q_\phi(z|x)$ as a deterministic transformation $g_\phi(x,\epsilon)$ of the input $x$ and another random variable $\epsilon$ with parameters $\phi$:

$$ z=g_\phi(x,\epsilon),\quad \epsilon\sim p(\epsilon) $$

Note that $p(\epsilon)$ is parameter-free and independent $x$ and $\phi$ then, for any function $f(x,z)$ we then have:

$$ E_{q_\phi(z|x)} \left[f(x,z)\right] = E_{\epsilon} \left[f\left(x,g_\phi(x,\epsilon)\right)\right] \approx f\left(x,g_\phi(x,\epsilon)\right),\quad\text{with}\quad \epsilon\sim p(\epsilon)$$

the advantage of this reparametrization is that now the parameter $\phi$ appears in the integrad but not in the distribution with respect to which the intrand is computed. In the final step we have replaced the expecatation value with a sigle Monte Carlo drawn. **Note that, in general, the change of variable form $z\to\epsilon$ lead to the introduction of a Jacobian.**

## Implementation of VAE using pyTorch

The equations for the encoder are:
$$
a_x=relu(W_x x + b_x) \\
\mu_z = W_\mu a_x + b_\mu \\
\log\sigma_z = W_\sigma a_x + b_\sigma
$$
where $x$ has dimension $D_x$, $a_x$ has dimension $D_\text{hidden_encoder}$ and $z$ has dimension $D_z$

The equation for the decoder are:
$$
a_2 = relu(W_1 z + b_1) \\
a_3 = sigmoid(W_2 a_2 + b_2) \\
p(\tilde{x}|a_3) = Bern(\tilde{x},a_3) = a_3^{\tilde{x}} (1-a_3)^{(1-\tilde{x})}
$$
where $z$ has dimension $D_z$, $a_2$ has dimension $D_\text{hidden_decoder}$ and $a_3$ has the same dimension as $D_x$.

The way that encoder and decoder are coupled together is:
$$
\mu_z, \log\sigma_z = \text{encoder}(x_\text{input}) \\
\sigma_z = \exp\left[\log\sigma_z\right] \\
\epsilon \sim N(0,I) \\
z = \mu_z + \sigma_z * \epsilon \\
x_\text{output} = \text{decoder}(z)
$$

Let us implement this in pyTorch now!

### Let's start by importing some packages and define helper functions

```python
import torch
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision.utils import save_image
import torch.optim as optim
from torch import nn
import matplotlib.pyplot as plt
%matplotlib inline

# Decide if the system has GPUs
if torch.cuda.is_available():
    device = torch.device('cuda')
    dtype = torch.cuda.FloatTensor
else:
    device = torch.device('cpu')
    dtype = torch.FloatTensor

print(device)
print(dtype)

def save_everything(model_object,optimizer_object,filepath):
        """
        Save all the quantities related to model and optimizer
        """
        state = model_object.save_model()
        state['optimizer'] = optimizer_object.state_dict()
        torch.save(state, filepath)

def load_everything(model_object,optimizer_object,filepath):
        """
        Load all the quantities relateds to the object from file
        """
        state = torch.load(filepath)
        model_object.load_model(state)
        optimizer_object.load_state_dict(state['optimizer'])

def show_batch(images):
    """
    Simple function to show the images
    """
    im = torchvision.utils.make_grid(images)
    plt.imshow(np.transpose(im.numpy(), (1, 2, 0)))
```

    cuda
    <class 'torch.cuda.FloatTensor'>

## Define the Encode, Decoder and CAE class

```python
# Define the encoder class by specifying its layers and the connections between layers
class Encoder(torch.nn.Module):
    def __init__(self, DIM_x, DIM_hidden_encoder, DIM_z):
        super().__init__()
        self.linear1 = torch.nn.Linear(DIM_x, DIM_hidden_encoder)
        self.z_mu = torch.nn.Linear(DIM_hidden_encoder, DIM_z)
        self.z_log_sigma2 = torch.nn.Linear(DIM_hidden_encoder, DIM_z)

    def forward(self, x):
        x = F.relu(self.linear1(x))
        z_mu = self.z_mu(x)
        z_log_sigma2 = self.z_log_sigma2(x)
        return z_mu,z_log_sigma2

# Define the encoder class by specifying its layers and the connections between layers
class Decoder(torch.nn.Module):
    def __init__(self, DIM_z, DIM_hidden_decoder, DIM_x):
        super().__init__()
        self.linear1 = torch.nn.Linear(DIM_z, DIM_hidden_decoder)
        self.linear2 = torch.nn.Linear(DIM_hidden_decoder, DIM_x)

    def forward(self, z):
        z = F.relu(self.linear1(z))
        a3= torch.sigmoid(self.linear2(z))
        return a3

# Define the encoder VAE class by combining encoder and decoder
class VAE(torch.nn.Module):

    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.loss_history = []
        self.epoch = 0
        self.training_time = 0
        self.BCE_loss = torch.nn.BCELoss(reduction='sum')

    def save_model(self):
        state={
            'state_dict': self.state_dict(),
            'loss_history': self.loss_history,
            'epoch': self.epoch,
            'training_time' : self.training_time
        }
        return state

    def load_model(self,state):
        self.load_state_dict(state['state_dict'])
        self.loss_history = state['loss_history']
        self.epoch = state['epoch']
        self.training_time = state['training_time']

    def forward(self, x_input):

        # encoder
        z_mu, z_log_sigma2 = self.encoder(x_input)

        # connection between encoder and decoder
        z_sigma = torch.exp(0.5*z_log_sigma2)
        if(device == 'cpu'): ### create a random variable on cpu
            epsilon = dtype.from_numpy(np.random.normal(0, 1, size=z_sigma.size())) #N(0,I)
        else: ### create a random variable on cuda
            epsilon = dtype(size=z_sigma.size()).normal_()
        z = z_mu + z_sigma * Variable(epsilon, requires_grad=False) # Reparameterization trick

        # decoder
        x_output = self.decoder(z)
        return x_output, z_mu, z_log_sigma2

    def compute_loss(self, x_input, x_output, z_mu, z_log_sigma2):
        # binary cross entropy
        BCE = self.BCE_loss(x_output,x_input) #Binary cross entropy between output and target
        # KL divergence
        z_sigma2 = torch.exp(z_log_sigma2)
        KLD = 0.5 * torch.sum(z_sigma2 + z_mu.pow(2) - 1 - z_log_sigma2)
        #print("BCE, KLD",BCE.item(),KLD.item())
        return BCE+KLD

    def train_one_epoch(self,trainloader):

        tmp = []
        for i, data in enumerate(trainloader, 0): #loop over minibatches

            # Get images and flatten them into vectors
            images_in, labels  = data
            images_in = images_in.to(device) #put the data on GPU if necessary
            x_input = Variable(images_in.view(-1,DIM_x),requires_grad=False)

            # Run the VAE forward
            x_output, z_mu, z_log_sigma2 = self.forward(x_input)
            #print("xout",x_output.dtype)
            #print("xout",x_output.device)

            # Compute the loss
            loss = self.compute_loss(x_input, x_output, z_mu, z_log_sigma2)
            #print("loss",loss.dtype)
            #print("loss",loss.device)

            # compute the average loss in this epoch
            tmp.append(loss.item()) # add the loss to the tmp list

            # For each minibatch set the gradient to zero
            optimizer.zero_grad()
            loss.backward()  # do backprop and compute all the gradients
            optimizer.step() # update the parameters

        # Svae the average loss during the epoch and the final value at the end of epoch
        self.loss_history.append(np.mean(tmp))
        self.epoch += 1

    def train(self,trainloader,DELTA_EPOCHS):
        import timeit
        tic=timeit.default_timer()

        for nn in range(DELTA_EPOCHS):
            self.train_one_epoch(trainloader)
            print("EPOCH loss=",self.epoch,self.loss_history[-1])

        toc=timeit.default_timer()
        self.training_time += toc - tic

    def evaluate_one_batch(self,trainloader):
        with torch.no_grad():
            images_in, labels  = next(iter(trainloader))
            x_input = images_in.view(-1,DIM_x).to(device)
            x_output, z_mu, z_log_sigma = self.forward(x_input)
            out_reshaped = x_output.view(-1,1,height, width)
            images_out = out_reshaped.cpu()
        return images_in,images_out
```

### Run the model

```python
if __name__ == '__main__':

    BATCH_SIZE = 64

    # torchvision.datasets.MNIST outputs a set of PIL images
    # We transform them to tensors
    transform = transforms.ToTensor()
    #transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])

    # Load and transform data
    trainset = torchvision.datasets.MNIST('./data', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True)

    testset = torchvision.datasets.MNIST('./data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE, shuffle=False)

    # Visualize the data and extract the image size
    images, labels =next(iter(trainloader))
    _, _, height, width = images.size()

    # The other dimensions are arbitrary and chosen by user
    DIM_x              = height*width
    DIM_z              = 2
    DIM_hidden_encoder = 256
    DIM_hidden_decoder = 128

    # Save the same set of test images for testing
    images_test_in, labels_test  = next(iter(testloader))
    x_test_one_batch = images_test_in.view(-1,DIM_x).to(device)

    # Create instances of Encoder, Decoder and VAE on either GPU or CPU
    encoder = Encoder(DIM_x, DIM_hidden_encoder, DIM_z).to(device)
    decoder = Decoder(DIM_z, DIM_hidden_decoder, DIM_x).to(device)
    vae = VAE(encoder, decoder).to(device) #note that here I am creating the model to either gpu or cpu

    # Select the optimizer
    #optimizer = optim.Adam(vae.parameters(), lr=0.0001)
    optimizer = optim.RMSprop(vae.parameters())
```

```python
    for k in range(10,11):
        if(k>0):
            load_everything(vae,optimizer,'./vae_'+str(k-1)+'.pth')

        vae.train(trainloader,10)  # train the model for 10 EPOCHS
        save_everything(vae,optimizer,'./vae_'+str(k)+'.pth')

        with torch.no_grad():
            x_test_out_one_batch, _ , _ = vae(x_test_one_batch)
            images_in = x_test_one_batch.view(-1,1,height, width).cpu()
            images_out = x_test_out_one_batch.view(-1,1,height, width).cpu()

        if(k==0):
            save_image(images_in,'./input.png')
        save_image(images_out,'./output_' + str(vae.epoch) + '.png')
```

```python
plt.semilogy(vae.loss_history[10:])
```

```python
plt.plot(vae.loss_history[10:])
images_in,images_out = vae.evaluate_one_batch(trainloader)
```

```python
show_batch(images_out)
```

```python
show_batch(images_in)
```

```python
z1 = list()
z2 = list()
mu_z1 = list()
mu_z2 = list()
labs = list()
with torch.no_grad():
    for i, data in enumerate(testloader, 0):
    #for i in range(0,20):
        images_in, labels  = next(iter(testloader))
        x_input = images_in.view(-1,DIM_x).to(device)
        z_mu,z_log_sigma2 = vae.encoder(x_input)

        z_sigma = torch.exp(0.5*z_log_sigma2)
        epsilon = dtype(size=z_sigma.size()).normal_()
        z = z_mu + z_sigma * epsilon

        for a,b in z_mu.cpu().detach().numpy():
            mu_z1.append(a)
            mu_z2.append(b)
        for c,d in z.cpu().detach().numpy():
            z1.append(c)
            z2.append(d)
        for ll in labels.cpu().detach().numpy():
            labs.append(ll)
```

```python
#plt.figure(figsize=golden_size(6))
plt.scatter(z1, z2, c=labs, cmap='nipy_spectral')
plt.colorbar()
plt.xlabel('$z_1$')
plt.ylabel('$z_2$')
plt.savefig('./VAE_MNIST_latent_z.pdf')
plt.show()
```

```python
#plt.figure(figsize=golden_size(6))
plt.scatter(mu_z1, mu_z2, c=labs, cmap='nipy_spectral')
plt.colorbar()
plt.xlabel('$mu_{z1}$')
plt.ylabel('$mu_{z2}$')
plt.savefig('./VAE_MNIST_latent_mu_z.pdf')
plt.show()
```

```python
display('./input.png')
```

    './input.png'

```python
from IPython.display import Image, display

listOfImageNames = ['./input.png',
                    './output_10.png',
                    './output_20.png',
                    './output_30.png',
                    './output_40.png',
                    './output_50.png',
                    './output_60.png',
                    './output_70.png',
                    './output_80.png',
                    './output_90.png',
                    './output_110.png',
                    './output_100.png']

for imageName in listOfImageNames:
    display(Image(filename=imageName))
```

```python

```


```python
import pytex
import numpy as np
from IPython.core.display import Latex

# make sympy equations pretty
from sympy.interactive import printing
printing.init_printing(use_latex=True)
```

```python
(pytex.magnitude(4),
 pytex.magnitude(-.04),
 pytex.magnitude(-499))
```

```python
# calculate magnitude to which error will be rounded
(pytex.prec_by_err(0.29),
 pytex.prec_by_err(0.31),
 pytex.prec_by_err(100))
```

```python
(pytex.round_to_mag(-0.24, -1),
 pytex.round_to_mag(0.05, -2, rdg='UP'),
 pytex.round_to_mag(0.01, -1, rdg='DOWN'))
```

```python
(pytex.repr_float(2.1),
 pytex.repr_float(2.1e-33))
```

    ('2.1', '2.1 \\times 10^{-33}')

```python
Latex('$%s$' % pytex.repr_float(2.1e33))
```

$2.1 \times 10^{33}$

```python
(pytex.repr_decimal(1e-23),
 pytex.repr_decimal(1e-23, -4),
 pytex.repr_decimal(1.2345, -2),
 pytex.repr_decimal(1.2345, -2, rdg='UP'))
```

    ('0.000000000000000000000001', '0.0000', '1.23', '1.24')

```python
pytex.repr_error(1.282e-33)
```

    '1.3 \\times 10^{-33}'

```python
pytex.align_num(321.2, dleft=4, dright=4)
```

    '\\phantom{+0} 321.2 \\phantom{000}'

```python
aligned_list = pytex.align_num_list([123, 123.132, 2e22, -10, 100e22])
print('\n'.join(aligned_list))
Latex('<br>'.join([r'$%s$' % e for e in aligned_list]))
```

    \phantom{+} 123.0 \phantom{00}\phantom{ \times 10^{22}}
    \phantom{+} 123.132\phantom{ \times 10^{22}}
    \phantom{+00} 2 \phantom{.000} \times 10^{22}
    -\phantom{0} 10.0 \phantom{00}\phantom{ \times 10^{22}}
    \phantom{+00} 1 \phantom{.000} \times 10^{24}

$\phantom{+} 123.0 \phantom{00}\phantom{ \times 10^{22}}$
$\phantom{+} 123.132\phantom{ \times 10^{22}}$
$\phantom{+00} 2 \phantom{.000} \times 10^{22}$
$-\phantom{0} 10.0 \phantom{00}\phantom{ \times 10^{22}}$
$\phantom{+00} 1 \phantom{.000} \times 10^{24}$

```python
Latex(pytex.format_valerr(1231, 1.2))
```

$1231.0 \pm 1.2$

```python
Latex(pytex.format_valerr(1231e10, 1.2e10))
```

$(12310 \pm 12) \times 10^{9}$

```python
values = (np.random.rand(10)-.5) * 2015
errors = np.random.rand(10) * 10

valerrs = pytex.format_valerr_list(values, errors)
Latex('<br>'.join(valerrs))
```

$-917 \phantom{.0} \pm 10 \phantom{.0}$
$\phantom{+} 654 \phantom{.0} \pm \phantom{0} 8 \phantom{.0}$
$-220 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$
$-335 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$
$\phantom{+} 621.3 \pm \phantom{0} 1.9$
$\phantom{+} 674.0 \pm \phantom{0} 2.9$
$\phantom{+} 747 \phantom{.0} \pm 10 \phantom{.0}$
$-420.0 \pm \phantom{0} 1.6$
$-841 \phantom{.0} \pm 10 \phantom{.0}$
$\phantom{+} 650.0 \pm \phantom{0} 2.2$

```python
Latex(pytex.table_html([valerrs + valerrs, valerrs], fill='-', just='lc', vsep='| |', hsep={0, -1}))
```

		$-917 \phantom{.0} \pm 10 \phantom{.0}$
		$-917 \phantom{.0} \pm 10 \phantom{.0}$

		$\phantom{+} 654 \phantom{.0} \pm \phantom{0} 8 \phantom{.0}$
		$\phantom{+} 654 \phantom{.0} \pm \phantom{0} 8 \phantom{.0}$

		$-220 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$
		$-220 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$

		$-335 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$
		$-335 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$

		$\phantom{+} 621.3 \pm \phantom{0} 1.9$
		$\phantom{+} 621.3 \pm \phantom{0} 1.9$

		$\phantom{+} 674.0 \pm \phantom{0} 2.9$
		$\phantom{+} 674.0 \pm \phantom{0} 2.9$

		$\phantom{+} 747 \phantom{.0} \pm 10 \phantom{.0}$
		$\phantom{+} 747 \phantom{.0} \pm 10 \phantom{.0}$

		$-420.0 \pm \phantom{0} 1.6$
		$-420.0 \pm \phantom{0} 1.6$

		$-841 \phantom{.0} \pm 10 \phantom{.0}$
		$-841 \phantom{.0} \pm 10 \phantom{.0}$

		$\phantom{+} 650.0 \pm \phantom{0} 2.2$
		$\phantom{+} 650.0 \pm \phantom{0} 2.2$

		$-917 \phantom{.0} \pm 10 \phantom{.0}$
		-

		$\phantom{+} 654 \phantom{.0} \pm \phantom{0} 8 \phantom{.0}$
		-

		$-220 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$
		-

		$-335 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$
		-

		$\phantom{+} 621.3 \pm \phantom{0} 1.9$
		-

		$\phantom{+} 674.0 \pm \phantom{0} 2.9$
		-

		$\phantom{+} 747 \phantom{.0} \pm 10 \phantom{.0}$
		-

		$-420.0 \pm \phantom{0} 1.6$
		-

		$-841 \phantom{.0} \pm 10 \phantom{.0}$
		-

		$\phantom{+} 650.0 \pm \phantom{0} 2.2$
		-

```python
print(pytex.table_latex([valerrs, valerrs]))
```

    \begin{table}[H]
    \centering
    	\begin{tabular}{|c|c|}
    	\hline
    	$-917 \phantom{.0} \pm 10 \phantom{.0}$ & $-917 \phantom{.0} \pm 10 \phantom{.0}$\\ \hline
    	$\phantom{+} 654 \phantom{.0} \pm \phantom{0} 8 \phantom{.0}$ & $\phantom{+} 654 \phantom{.0} \pm \phantom{0} 8 \phantom{.0}$\\
    	$-220 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$ & $-220 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$\\
    	$-335 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$ & $-335 \phantom{.0} \pm \phantom{0} 9 \phantom{.0}$\\
    	$\phantom{+} 621.3 \pm \phantom{0} 1.9$ & $\phantom{+} 621.3 \pm \phantom{0} 1.9$\\
    	$\phantom{+} 674.0 \pm \phantom{0} 2.9$ & $\phantom{+} 674.0 \pm \phantom{0} 2.9$\\
    	$\phantom{+} 747 \phantom{.0} \pm 10 \phantom{.0}$ & $\phantom{+} 747 \phantom{.0} \pm 10 \phantom{.0}$\\
    	$-420.0 \pm \phantom{0} 1.6$ & $-420.0 \pm \phantom{0} 1.6$\\
    	$-841 \phantom{.0} \pm 10 \phantom{.0}$ & $-841 \phantom{.0} \pm 10 \phantom{.0}$\\
    	$\phantom{+} 650.0 \pm \phantom{0} 2.2$ & $\phantom{+} 650.0 \pm \phantom{0} 2.2$\\ \hline
    	\end{tabular}
    \end{table}

```python

```


测试DVI正态近似的准确性

```python
%matplotlib inline

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from IPython.display import clear_output
import os, json

import gaussian_variables as gv
import utils as u
import plot_utils as pu
import bayes_layers as bnn
from bayes_models import MLP, PointMLP, AdaptedMLP

np.random.seed(3)
```

首先构建数据集：
\begin{equation}
    y = -(x+0.5)\sin(3\pi x) + \eta
\end{equation}

\begin{equation}
    \eta = 0.45(x + 0.5)^2
\end{equation}

```python
def base_model(x):
    return -(x+0.5)*np.sin(3 * np.pi *x)

def noise_model(x):
    return 0.45*(x+0.5)**2

def sample_data(x):
    return base_model(x) + np.random.normal(0, noise_model(x))

data_size = {'train': 500, 'valid': 100, 'test': 100}
toy_data = []
for section in ['train', 'valid', 'test']:
    x = (np.random.rand(data_size['train'], 1) - 0.5)
    toy_data.append([x, sample_data(x).reshape(-1)])
x = np.arange(-1,1,1/100)
toy_data.append([[[_] for _ in x], base_model(x)])

pu.toy_results_plot(toy_data, {'mean':base_model, 'std':noise_model})
```

设置模型参数，这里以一个三层、每层5个神经元的网络为例(本项目中的其他网络结果修改“hidden_dims”参数可得到)

```python
hypers = {
    "x_dim": 1,                                "y_dim": 2,
    "hidden_dims": [5,5,5],                  "nonlinearity": "relu",
    "adapter": {
        'in' : {"scale": [[1.0]], "shift": [[0.0]]},
        'out': {"scale": [[1.0, 0.83]], "shift": [[0.0, -3.5]]}
    },

    "method": "bayes",                         "style": "heteroskedastic",
    "homo_logvar_scale": 2*np.log(0.2),        "prior_type": ["empirical", "wider_he", "wider_he"],

    "n_epochs": 20000,                         "batch_size": 500,
    "learning_rate": 0.001,                    "lambda": 1.0,
    "warmup_updates": {'lambda': 14000.0},     "anneal_updates": {'lambda': 1000.0},
    "optimizer": "adam",                       "gradient_clip": 0.1,

    "data_fraction": 1.0,                      "sections_to_run": ["train"]
}
```

构建模型

```python
def make_model(hypers):
    if hypers['method'].lower().strip() == 'bayes':
        MLP_factory = MLP
        prediction = lambda y: tf.reshape(y.mean[:,0], [-1])
        loss = bnn.regression_loss
    else:
        MLP_factory = PointMLP
        prediction = lambda y: tf.reshape(y.mean[:,0], [-1])
        loss = bnn.point_regression_loss

    mlp = MLP_factory(hypers['x_dim'], hypers['y_dim'], hypers)
    mlp = AdaptedMLP(mlp)
    mlp.make_placeholders()
    ipt = mlp.placeholders['ipt_mean']
    y = mlp(ipt)
    target = tf.placeholder(tf.float32, [None])
    mlp.placeholders['target'] = target
    global_step = tf.Variable(0, trainable=False, name='global_step')
    loss, logprob, all_surprise = loss(y, target, mlp, hypers, global_step)

    accuracy = tf.reduce_mean(tf.abs(target - prediction(y)))

    return {
        'model': mlp,
        'metrics': {
            'accuracy': accuracy,   'loss': loss,
            'logprob': logprob,     'all_surprise': all_surprise
            },
        'global_step': global_step}
```

给出一个比较DVI和MCVI算法结果的函数

```python
from scipy import stats
def show_compare(model_and_metrics, sess):
    plt.figure()
    n_samp = 20000
    x = 0.25
    ipt = [[[x]] for _ in range(n_samp)]

    sample_op = model_and_metrics['model'].run_with_MC(
        ipt, n_sample=n_samp)
    approx_op = model_and_metrics['model'](x)
    samples = sess.run(sample_op)
    approx = sess.run([approx_op.mean, approx_op.var])

    # samples_b.shape
    m_min = stats.norm.ppf(
        0.0001, loc=approx[0][0, 0], scale=np.sqrt(approx[1][0, 0, 0]))
    m_max = stats.norm.ppf(
        0.9999, loc=approx[0][0, 0], scale=np.sqrt(approx[1][0, 0, 0]))
    l_min = stats.norm.ppf(
        0.0001, loc=approx[0][0, 1], scale=np.sqrt(approx[1][0, 1, 1]))
    l_max = stats.norm.ppf(
        0.9999, loc=approx[0][0, 1], scale=np.sqrt(approx[1][0, 1, 1]))

    bin_no_m = np.linspace(m_min, m_max, 50)
    bin_no_l = np.linspace(l_min, l_max, 50)

    fig = plt.figure()

    ax1 = fig.add_subplot(221)
    ax1.hist(samples[:, 0, 0], bin_no_m,
             density=True, edgecolor='k', facecolor='#b4c7e7')
    ax1.plot(*gaussian1d(approx[0][0, 0], approx[1][0, 0, 0],
                         m_min, m_max), 'b')
    plt.xlim([m_min, m_max])
    ax1.set_yticks([])
    ax1.set_xlabel('$m$')
    ax1.set_ylabel('$q(m)$')

    ax2 = fig.add_subplot(222)
    ax2.hist(samples[:, 0, 1], bin_no_l,
             density=True, edgecolor='k', facecolor='#b4c7e7', label="MC")

    ax2.plot(*gaussian1d(approx[0][0, 1],
                         approx[1][0, 1, 1],
                         l_min,
                         l_max), 'b', label="ours")
    plt.xlim([l_min, l_max])
    ax2.set_yticks([])
    ax2.set_xlabel('$\ell$')
    ax2.set_ylabel('$q(\ell)$')
    plt.show()
    return None
```

```python
def gaussian1d(mean, var, min, max):
    x_axis = np.linspace(min, max, 1000)
    return x_axis, 1.0 / np.sqrt(2.0 * np.pi * var) * \
        np.exp(-1.0 / (2.0 * var) * (x_axis - mean)**2)
```

定义训练函数

```python
def run(data):
    run_id = u.start_run()

    restricted_training_set = u.restrict_dataset_size(data[0], hypers['data_fraction'])
    hypers['dataset_size'] = len(restricted_training_set[0])

    device_id = 1
    device_string = u.get_device_string(device_id)
    print(hypers)
    with tf.device(device_string):
        if True:
            model_and_metrics = make_model(hypers)
            train_op = u.make_optimizer(model_and_metrics, hypers)
            sess = u.get_session()
            saver = tf.train.Saver()

            all_summaries = []
            best_valid_accuracy = np.inf

        show_compare(model_and_metrics, sess)

        for epoch in range(hypers['n_epochs']):
            verbose = (epoch % 20 == 0)
            if verbose:
                print("Epoch %i:        " % epoch, end='')
            epoch_summary, accuracies = u.train_valid_test(
                {
                    'train': restricted_training_set,
                    'valid': data[1],
                    'test': data[2]
                },
                sess, model_and_metrics, train_op, hypers, verbose)

        show_compare(model_and_metrics, sess)

```

开始训练，并得到两个DVI和MCVI之间的比较

```python
run(toy_data)
```

将之前的数据集整体增大10倍

```python
def base_model_10(x):
    return 10*(-(x+0.5)*np.sin(3 * np.pi *x))

def noise_model_10(x):
    return 10*(0.45*(x+0.5)**2)

def sample_data_10(x):
    return base_model_10(x) + np.random.normal(0, noise_model_10(x))

data_size = {'train': 500, 'valid': 100, 'test': 100}
toy_data = []
for section in ['train', 'valid', 'test']:
    x = (np.random.rand(data_size['train'], 1) - 0.5)
    toy_data.append([x, sample_data_10(x).reshape(-1)])
x = np.arange(-1,1,1/100)
toy_data.append([[[_] for _ in x], base_model_10(x)])

pu.toy_results_plot(toy_data, {'mean':base_model_10, 'std':noise_model_10})
plt.ylim([-11,16])
```

同样的方式进行训练，并得到DVI和MCVI之间的比较

```python
run(toy_data)
```

对于m的结果，DVI和MCVI的一致性变强了。

```python

```