# Technical Details

As with all packages, numerous technical details that are abstracted away from the user. Now to ensure a clean interface, this abstraction is entirely necessary. However, it can sometimes be confusing when navigating a package's source code to pin down what's going on when there's so many _under the hood_ operations taking place. In this notebook, I'll aim to shed some light on all of the tricks that we do in GPJax to help elucidate the code to anyone wishing to extend GPJax for their own uses.


## Parameter Transformations

### Motivations

Many parameters in a Gaussian process are what we call a _constrained parameter_. By this, we mean that the parameter's value is only defined on a subset of $\mathbb{R}$. One example of this is the lengthscale parameter in any of the stationary kernels. It would not make sense to have a negative lengthscale, and as such the parameter's value is constrained to exist only on the positive real line. 

Whilst mathematically correct, constrained parameters can become a pain when optimising as many optimisers are designed to operate on an unconstrained space. Further, it can often be computationally inefficient to restrict the search space of an optimiser. For these reasons, we instead transform the constrained parameter to exist in an unconstrained space. The optimisation is then done on this unconstrained parameter before we transform it back when we need to evaluate its value. 

Only bijective transformations are valid as we cannot afford to lose our original parameter value when transforming. As such, we have to be careful about which transformations we use. Some common choices include the log-exponential bijection and the softplus transform. We, by default, opt for the softplus transformation in GPJax as it is less prone to numerical overflow in comparison to log-exp transformations.


### Implementation

When it comes to implementations, we attach the transformation directly to the `Parameter` class. It is an optional argument that one can specify when instantiating their parameter. To see this, consider the following example


```python
from gpjax.parameters import Parameter
from gpjax.transforms import Softplus
import jax.numpy as jnp

x = Parameter(jnp.array(1.0), transform=Softplus())
```

    WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)


Now we know that the softplus transformation operation on an input $x \in \mathbb{R}_{>0}$ can be written as 
$$\alpha(x) = \log(\exp(x)-1)$$
where $\alpha(x) \in \mathbb{R}$. In this instance, it can be seen that $\alpha(1)=0.54$. Now this unconstrained value is stored within the parameter's `value` property


```python
print(x.value)
```

    0.541324854612918


whilst the original constrained value can be computed by accesing the parameter's `untransform` property


```python
print(x.untransform)
```

    1.0


### Custom transformation

Should you wish to define your own custom transformation, then this can easily be done by extending the `Transform` class within `gpjax.transforms` and defining a forward transformation and a backward transformation.


```python
class Transform:
    def __init__(self, name="Transformation"):
        self.name = name

    @staticmethod
    def forward(x: jnp.ndarray) -> jnp.ndarray:
        raise NotImplementedError

    @staticmethod
    def backward(x: jnp.ndarray) -> jnp.ndarray:
        raise NotImplementedError
```

The `forward` method is the transformation that maps from a constrained space to an unconstrained space, whilst the `backward` method is the transformation that reverses this. A nice example of this can be seen for the earlier used softplus transformation


```python
from jax.nn import softplus

class Softplus(Transform):
    def __init__(self):
        super().__init__(name='Softplus')

    @staticmethod
    def forward(x: jnp.ndarray) -> jnp.ndarray:
        return jnp.log(jnp.exp(x) - 1.)

    @staticmethod
    def backward(x: jnp.ndarray) -> jnp.ndarray:
        return softplus(x)
```

## Prior distributions

### Motivations

Often when we use Gaussian processes, we do so as they facilitate incorporation of prior information into the model. Implicitly, by the very use of a Gaussian process, we are incorporating our prior information around the functional behaviour of the latent function that we are seeking to recover. However, we can take this one step further by placing priors on the hyperparameters of the Gaussian process. Going into the details of which priors are recommended and how to go about selecting them goes beyond the scope of this article, but it suffices to say that doing so can greatly enhance the utility of a Gaussian process. 

At least in my own experience, when priors are placed on the hyperparameters of a Gaussian process they are specified with respect to the constrained parameter value. As an example of this, consider the lengthscale parameter $\ell \in \mathbb{R}_{>0}$. When specifying a prior distribution $p_{0}(\ell)$, I would typically select a distribution that has support on the positive real line, such as the Gamma distribution. An opposing approach would be to transform the parameter so that it is defined on the entire real line (as discussed in §1) and then specify a prior distribution such as a Gaussian that has unconstrained support. Deciding which of these two approaches to adopt in GPJax is somewhat a moot point to me, so I've opted for priors to be defined on the constrained parameter. That being said, I'd be more than open to altering this opinion is people felt strongly that priors should be defined on the unconstrained parameter value.

### Implementation

Regarding the implementational details of enabling prior specification, this is hopefully a more lucid concept upon code inspection. As with the earlier discussed parameter transformations, the notion of a prior distribution is acknowledged in the definition of a parameter. To exactly specify a prior distribution, one should simply call in the relevant distribution from TensorFlow probability's distributions module. For an example of this, consider the parameter `x` that was earlier defined.


```python
from tensorflow_probability.substrates.jax import distributions as tfd

x.prior = tfd.Gamma(concentration = 3., rate = 2.)
```

If we momentarily pause to consider the state of this parameter now, then we have a constrained parameter value with a corresponding prior distribution. When it comes to deriving our posterior distribution, then we know that it is proportional to the product of the likelihood and the prior density function. As addition is less prone to numerical overflow than multiplication, we take the log of this produce. The log of a product is just a sum of logs, meaning that our log-posterior is then proportional to the sum of our log-likelihood and the log-prior density. Therefore, to connect the value of our parameter and its respective prior distribution, the only implementational point left to cover is how to evaluate the parameters log-prior density. This can be done through the following `@property`


```python
print(x.log_density)
```

    -0.613706111907959


Naturally, should one wish to evaluate the prior density of the parameter, then the exponent can be taken


```python
print(jnp.exp(x.log_density))
```

    0.5413408768770793


## Cholesky decomposition

### Motivation


#### Single-step of the Cholesky decomposition 

Before we examine the use of Cholesky decompositions in GPJax, I think it's worthwhile to first see why we even need the Cholesky decomposition. In mathematics and statistics, we are often presented with the task of solving a linear system of equations $A\mathbf{x}=\mathbf{y}$ for $A \in \mathbb{R}^{m \times n}$ and $\mathbf{x} \in \mathbb{R}^{n}$, $\mathbf{y}\in\mathbb{R}^{m}$. In this scenario, our goal is to identify the values of $\mathbf{x}$ and for people who studied numerical methods, this probably evokes flashbacks to manually calculating $\mathbf{x}$ using Gaussian elimination. Whilst manageable for a handful of unknown variables, we certainly would not want to do Gaussian elimination for systems containing thousands of variables. Fortunately, computers are quite good at solving these types of problem.

Let our matrix $A$ now be symmetric in $\mathbb{R}^{n \times n}$ (analogously termed Hermitian in $\mathbb{C}^{n \times n}$) and positive definite (SPD) i.e. 
$$\begin{align*} A & = A^{\top} \quad & \text{(symmetry)} \\
\lambda^{\top}A\lambda & > 0 \text{ , for all } \lambda \in \mathbb{R}^{n}\quad & \text{(positive-definite)}\end{align*}.$$
If we were to now begin to apply Gaussian elimination to $A$ with a 1 element in the upper-left, then the first step would yield
$$A = \begin{bmatrix}1 & \mathbf{v}^{\top}\\ 
\mathbf{v} & K \end{bmatrix} = \begin{bmatrix}1 & \mathbf{0} \\ 
\mathbf{v} & I \end{bmatrix}\begin{bmatrix}1 & \mathbf{v}^{\top} \\ 
\mathbf{0} & K - \mathbf{v} \mathbf{v}^{\top} \end{bmatrix}. \tag{1}$$
Under the process of regular Gaussian elimination, we would now move onto the second column and introduce $(n-2)$-length $\mathbf{0}$ vector here too. However, doing so would invalidate the matrix's symmetry. Therefore, before proceeding to the second column, a Cholesky decomposition introduce an of $(n-1)$-length $\mathbf{0}$ vector into the fist row too. To achieve this upper-right triangulation of our reduced matrix system in (1) we can write
$$\begin{bmatrix}1 & \mathbf{v}^{\top} \\ 
\mathbf{0} & K - \mathbf{vv}^{\top} \end{bmatrix} = \begin{bmatrix}1 & \mathbf{0} \\ 
\mathbf{0} & K - \mathbf{vv}^{\top} \end{bmatrix}\begin{bmatrix}1 & \mathbf{v}^{\top} \\ \mathbf{0} & I \end{bmatrix}. \tag{2}$$
Bringing the expressions in (1) and (2) together now we get the following factorisation of $A$
$$\begin{align}A & = \begin{bmatrix}1 & \mathbf{v}^{\top}\\ 
\mathbf{v} & K \end{bmatrix} \nonumber \\ & = \begin{bmatrix}1 & \mathbf{0} \\ 
\mathbf{v} & I \end{bmatrix} \begin{bmatrix}1 & \mathbf{0} \\ 
\mathbf{0} & K - \mathbf{vv}^{\top} \end{bmatrix}\begin{bmatrix}1 & \mathbf{v}^{\top} \\ \mathbf{0} & I \end{bmatrix}. \tag{3}\end{align}$$

A Cholesky factorisation will apply the operation described in (3) iteratively for all remaining row-column pairs.

#### Generalisation

At this point, you may well ask the question of how this factorisation can ever be generalisable as it is by no means guaranteed that $A_{1, 1} = 1$ is true. Simply put, we can introduce a constant $\alpha$ such that $\alpha = \sqrt{A_{1,1}}$. Working through (3) with an upper-left entry of $\alpha^2$ now instead of 1 we get the general form of a Cholesky decomposition
$$
\begin{align}A & = \begin{bmatrix}\alpha^2 & \mathbf{v}^{\top}\\ 
\mathbf{v} & K \end{bmatrix} \nonumber \\ 
& = \begin{bmatrix}\alpha & \mathbf{0} \\ 
\frac{\mathbf{v}}{\alpha} & I \end{bmatrix} \begin{bmatrix}1 & \mathbf{0} \\ 
\mathbf{0} & K - \frac{\mathbf{vv}^{\top}}{\alpha^2} \end{bmatrix}\begin{bmatrix}1 & \frac{\mathbf{v}^{\top}}{\alpha} \\ \mathbf{0} & I \end{bmatrix} \\ 
& = L_{1}^{\top}A_{1}L_{1}\end{align}$$
When this process is iteratively applied from the upper-left element of $A$ down to the lower-right element, we get the complete factorisation which we know as the Cholesky decomposition 
$$\begin{align}A & = L_1^{\top}L_2^{\top}\ldots L_n^{\top} L_n \ldots L_2 L_1 \\ 
& = L^{\top}L\end{align}$$
for an upper-triangular $L$. 

#### Uniqueness

For a Cholesky decomposition to work, the upper-left element of $K-\frac{\mathbf{vv}^{\top}}{\alpha^2}$ must be positive, and it's not trivial as to why this is true. However, when considering the first step of the Cholesky decomposition, we established at that $A$ and $L_1^{-\top}A L_1^{-1}$ are both SPD. Further, $K-\frac{\mathbf{vv}^{\top}}{\alpha^2}$ is clearly the principal submatrix of $L_1^{-\top}A L_1^{-1}$ and is therefore positive definite. Right at the start when we defined a positive-definite matrix we saw that for a PD matrix the diagonals are all positive, and therefore the upper-left element of $K-\frac{\mathbf{vv}^{\top}}{\alpha^2}$ is guaranteed to be positive. Having proved this now for $n=1$, we can apply proof by induction to show that this positiveness is true for all steps of a Cholesky factorisation.

A generalisation of this result is that every SPD matrix has a unique Cholesky factorisation. This can be proved by considering the fact that each $\alpha$ is unique as it is determined by the form of $L^{\top}L$. Once $\alpha$ is determined, the first row of $L^{\top}$ is also available. This is true for each step of the decomposition, and therefore the factorisation is unique.

#### Computational details

In the words of [Linus Torvalds](https://en.wikipedia.org/wiki/Linus_Torvalds), _“Talk is cheap. Show me the code.”_, it helps at this point to see the Cholesky decomposition through an algorithmic lens.
```python
L = A
for i in range(n):
    for j in range(i+1, n):
        L[j, j:n] -= L[i, j:n] * R[i, j]/R[i, i]
    R[i, i:n] = R[i, i:n]/R[i, i]**0.5
```

The triangular nature of L is unveiled here the inner for loop's decreasing span. It is also the multiplication and subtraction operations of this inner loop that dominate the computational cost of a Cholesky decomposition. At the $i^{\text{th}}$ step of the loop, there are $i-j+1$ subtractions and $i-j+1$ multiplications required. By then summing over $i \in 1:n$ and $j \in (i+1):n$ we get a scaling of $\frac{1}{3}n^3$ floating-point operations (flops) with $n$. If we wanted to achieve the same factorisation using Gaussian elimination, then our computational complexity would scale $\frac{2}{3}n^3$ flops; twice as slow in comparison to a Cholesky factorisation.

#### More stable computations

Now that we've established what a Cholesky factorisation is and how one might go about computing the factors, a natural question one might ask is why we would go to this effort? The answer lies in the increased stability of a Cholesky factor in comparison to Gaussian elimination. The stabilisation comes from the fact that the values of $L$ are upper-bounded by the values of $A$. When working with the $2$-norm operator, we can actually go one further and state that $\lVert L^{\top} \rVert := \lVert L \rVert := \lVert A \rVert^{\frac{1}{2}}$. More generally, for any finite norm, the values $\lVert L\rVert$ and $\lVert A \rVert^{\frac{1}{2}}$ cannot differ by more than a factor of $\sqrt{n}$.

#### Matrix inversion

Drawing upon the concept of Cholesky factorisation, we can recast our original problem of solving a system of linear equations using the lower-triangular matrices that are given in a Cholesky factorisation
$$\begin{align} A\mathbf{x} & = \mathbf{b} \nonumber \\
\implies  LL^{\top}\mathbf{x} & = \mathbf{b} \nonumber \\
\implies  LL^{\top}\mathbf{x} & = L\mathbf{y} \ \text{ for some } \mathbf{y} \nonumber \\
\implies  L^{\top}\mathbf{x} & = \mathbf{y}
\end{align}$$
For conciseness, we'll now adopt the backslash notation $A\backslash \mathbf{b}$ to indicate that the vector $\mathbf{x}$ solves the linear system of equations $A \mathbf{x} = \mathbf{b}$. By this, we can equate $\mathbf{x}$ such that
$$\mathbf{x} = L^{\top} \backslash (L \backslash \mathbf{b}).$$
As $L$ is a lower-triangular matrix, we can use forward-substitution to determine the values of $\mathbf{b} = L\mathbf{y}$. Similarily, $\mathbf{x} = L^{\top}\mathbf{y}$ can be solved using backward-substitution due to the upper-triangular structure in $L^{\top}$.

#### Determinants

A final note on the Cholesky factorisation that is relavent within the context of Gaussian processes is that for a $n \times n$ SPD matrix $A$, the matrix's determinant and log-determinant can be computed by
$$\lvert A \rvert = \prod_{i=1}^n L^2_{ii} \quad \text{and} \quad \log\lvert A \rvert = 2\log\sum_{i=1}^n L_{i,i}$$
respectively.

### Implementation

Due to the enhanced stability that is experience when working with Cholesky factors in comparison to regular SPD matrices, all matrix inversions in GPJax are done using the Cholesky factors. To see this, consider the evaluation of the conjugate Gaussian process' marginal log-likelihood term. Mathematically, this quantity is written as
$$\log p(\mathbf{y} | X, \theta) = -0.5 \left(\operatorname{logdet}(K_{xx}+\sigma_n^2 I) + \mathbf{y}^{\top}(K_{xx} + \sigma_n^2 I)^{-1}\mathbf{y} + N \log (2 \pi) \right).$$
However, by using the Cholesky factorisation $L = \operatorname{cholesky}(K_{xx}+\sigma_n^2 I)$ to compute the matrix inverses and determinants, we instead write
$$\log p(\mathbf{y}| X, \theta) = -0.5 \left(\mathbf{y}^{\top}\boldsymbol{\alpha} + \sum\limits_{i=1}^{n}\log L_{ii} + N \log (2 \pi) \right)$$
where 
$$\boldsymbol{\alpha} = L^{\top}\backslash (L \backslash \mathbf{y}).$$

It is this second form of the marginal log-likelihood that is used within GPJax due to its enhanced numerical stability; a crucial attribute when optimising the model's parameters with respect to it.

## Floating point precision

By default, Jax uses 32-bit floats. From the excellent Jax documentation, the rationale behind this is _" to mitigate the Numpy API’s tendency to aggressively promote operands to double"_. For many machine learning algorithms, 32-bit precision is probably sufficient. However, for Gaussian processes, we have to deal with challenging matrix inversion and we want to mitigate the effects of numerical rounding as much as possible.

To see this, consider a SPD matrix $A \in \mathbb{R}^{n \times n}$ and the vectors $\mathbf{x}, \mathbf{y} \in \mathbb{R}^{n}$ which, together, form a linear system of equations
$$ \mathbf{y} = A\mathbf{x}.$$
As the matrix is PD, we know that it is then invertible so we can then write
$$\begin{align*}\mathbf{y}^{\top}A^{-1}\mathbf{y} & = \mathbf{x}K^{\top}K^{-1}K\mathbf{x}\\
& = \mathbf{x}^{\top}K\mathbf{x}>0\end{align*}$$
i.e., for a PD matrix $A$ that is invertible, its inverse $A^{-1}$ is also invertible. 

Whilst we might be happy with this fact theoretically, the result breaks down computationally due to the floating-point rounding that occurs in our machines. To convince ourself of this, we can define a function that checks if a matrix is PD.


```python
from typing import Callable
import jax


def is_positive_definte(A: jnp.ndarray):
    return (jnp.linalg.eigvals(A)>0).all()


def sqexp_kernel(x: jnp.ndarray, y: jnp.ndarray, variance: float = 1.0, lengthscale: float = 1.0):
    tau = jnp.square(x-y)
    return jnp.square(variance)*jnp.exp(-0.5*tau/jnp.square(lengthscale))


def kernel(func: Callable, x: jnp.ndarray, y: jnp.ndarray):
    return jax.vmap(lambda x: jax.vmap(lambda y: func(x, y))(y))(x)
```

In the above chunk of code, we have defined the squared exponential kernel for which we know the resultant Gram matrix is PD. However, as we'll see below, even for a small number of points, the rounding introduced by 32-bit precision yields a non-SPD matrix.


```python
x = jnp.linspace(-3., 3., 20, dtype=jnp.float32)
K = kernel(sqexp_kernel, x, x)
Kinv = jnp.linalg.inv(K)

print(is_positive_definte(K))
print(is_positive_definte(Kinv))
```

    False
    False


If we now redefine our original data vector `x` as an array with 64bit precision and force Jax to use 64-bit precision, then this same matrix is now PD.


```python
from jax.config import config; config.update("jax_enable_x64", True)

x = jnp.linspace(-3., 3., 20, dtype=jnp.float64)
K = kernel(sqexp_kernel, x, x)
Kinv = jnp.linalg.inv(K)

print(is_positive_definte(K))
print(is_positive_definte(Kinv))
```

    True
    True


For this reason, we enforce 64bit precision in GPJax by loading the configuration import above when one imports GPJax. To further enforce numerical stability in the kernel matrix, we add some _jitter_ to the matrix's diagonal. In doing this, we do introduce some numerical error, but the amount is often tolerable. 

Again, we can see the effect of this jitter in the above example by reducing the distance between our points in `x`. This will reduce the size of the kernel matrix's eigenvalues and increase the chance of numerical rounding being a problem, even with 64-bit precision.


```python
x = jnp.linspace(-3., 3., 200, dtype=jnp.float64)
K = kernel(sqexp_kernel, x, x)
Kinv = jnp.linalg.inv(K)

print(is_positive_definte(K))
print(is_positive_definte(Kinv))
```

    False
    False


If we now add a tiny amount ($10^{-12}$) of jitter to the diagonal of K, then our inversion should restabilise.


```python
jitter = jnp.eye(K.shape[0])*1e-12
K += jitter
Kinv = jnp.linalg.inv(K)

print(is_positive_definte(K))
print(is_positive_definte(Kinv))
```

    True
    True


As can be seen, our matrix inversions are now much more stable.

In GPJax, we set the default amount of jitter as $10^{6}$. Often you'll be able to use a far smaller amount of jitter, particularly if you have a small amount of well-spaced input data. Conversely, you may need more jitter for some of your problems. Fortunately, this value can be altered upon instantiation of the GP prior through the option `jitter` argument. For example, this is how you could manually set the jitter to $10^{-10}$


```python
from gpjax.gps import Prior
from gpjax import RBF

f = Prior(RBF, jitter=1e-10)
```

7b468b28f3fe86e1f0e3cbe2bc7d8eed3eac754e
ipynb
Jupyter Notebook
docs/nbs/technical_details.ipynb
jejjohnson/GPJax-1
1ac2c1d5cfd52ca4d58875a83422e6230026b582
Apache-2.0
2020-12-03T14:07:39.000Z
2022-03-14T17:45:34.000Z
docs/nbs/technical_details.ipynb
jejjohnson/GPJax-1
1ac2c1d5cfd52ca4d58875a83422e6230026b582
Apache-2.0
2020-12-05T08:54:45.000Z
2022-03-01T09:56:50.000Z
docs/nbs/technical_details.ipynb
jejjohnson/GPJax-1
1ac2c1d5cfd52ca4d58875a83422e6230026b582
Apache-2.0
2021-02-05T12:37:57.000Z
2022-03-13T13:00:20.000Z
Qwen/Qwen-72B
1. YES
2. YES