# 1. Kernel Ridge Regression

\begin{align*}
\text{Primal form:}~&\Vert Ax - y \Vert_2^2 + \lambda \Vert x \Vert_2^2\\
\textbf{Closed-form solution:}~&x = \left(A^TA + \lambda I_n\right)^{-1}A^Ty\\
\end{align*}

From the identity $\left(A^TA + \lambda I_n\right)^{-1}A^T = A^T\left(AA^T + \lambda I_m\right)^{-1}$, we can obtain a closed-form solution with the alternate form:
\begin{align*}
\textbf{Alternate closed-form solution:}~x = A^T\left(AA^T + \lambda I_m\right)^{-1}y
\end{align*}

Let $A$ have $m$ rows and $n$ columns, such that each row is a data-point (or sample) and each column is a feature. In typical applications of kernel ridge regression, $m >> n$. This means that computing $\alpha$ directly requires computing the $m \times m$ matrix, $AA^T$. For very large $m$ this can be an expensive $m^2n$ operation and, in the extreme, we many not even have space to store $AA^T$.

The Kernelized variant of ridge regression replaces the $m \times n$ matrix $A$ with a high-dimensional feature space for each sample $\Phi(A)$. $\Phi(A)$ is computational-expensive to form, especially when the chosen kernel function has infinite-dimensional feature space (e.g. Radial Basis Function). However, the $m \times m$ kernel matrix $K$ can be computed since the kernel functions are defined by $K_{i,j} = \sum k(a_i, a_j)$, where $a_i$ is the $i$-th row of $A$ and $k(...)$ is a user-defined kernel function.

One example of such a kernel function, is the polynomial kernel: $k(a_i, a_j) = (a_ia_j^T + r)^d$ s.t. $r \geq 0$ and $d \geq 1$. *When $r = 0$ and $d = 1$, we obtain a simple dot-product (or linear kernel).* In replacing the dot-products with a kernel function, we obtain a feature-space that can have very high dimensions (possibly infinite-dimensions). This means that computing the solution, $x$, is often very expensive. So instead, we will focus on obtaining the solution:
\begin{align*}
\bf
\text{Alternate solution:}~&\alpha = \left(\frac{1}{\lambda}K + I_m\right)^{-1}y\\
\text{and}~&w = \frac{1}{\lambda}\Phi(A)^T\alpha
\end{align*}

# 2. Subspace Iteration for Kernel Ridge Regression

Note that the kernel matrix, $K$, is more expensive to compute than $AA^T$ (due to the kernel function), so we will instead compute the solution to $\alpha$ iteratively using ***subspace iteration*** by selecting, and solving for a few samples (row of $A$) every iteration.
\begin{align*}
\text{Dual form:}~&\Vert \alpha - y \Vert_2^2 + \frac{1}{\lambda} \Vert \Phi(A)^T\alpha \Vert_2^2\\

\bf\text{Iterative update:}~&\alpha_h = \alpha_{h-1} + \mathbb{I}_h^T {u_h}\\
\end{align*}
The matrix $\mathbb{I}_h$ has $b$ rows and $m$ columns such that rows of $\mathbb{I}_h$ are sub-sampled rows of the $m$-dimensional identity matrix, $I_m$.

By substituting the iterative update into the dual form, we obtain the following minimization problem:
\begin{align*}
{\arg\min}~\left\Vert \alpha_{h-1} + \mathbb{I}_h^Tu_h - y \right\Vert_2^2 + \frac{1}{\lambda} \left\Vert \Phi(A)^T\alpha_{h-1} + \Phi(A)^T\mathbb{I}_h^T u_h\right\Vert_2^2
\end{align*}

By taking the derivative w.r.t $u_h$ we can obtain the gradient for just the $b$ rows of $A$ that we are sampling at every iteration:
\begin{align*}
0 &= \mathbb{I}_h\alpha_{h-1} + u_h - \mathbb{I}_hy + \frac{1}{\lambda} \mathbb{I}_h \Phi(A)\Phi(A)^T\alpha_{h-1} + \frac{1}{\lambda}\mathbb{I}_h\Phi(A)\Phi(A)^T\mathbb{I}_h^T u_h\\
 \textbf{Gradient w.r.t chosen samples: }u_h &= \left(\frac{1}{\lambda} \mathbb{I}_h\Phi(A)\Phi(A)^T\mathbb{I}_h^T + I_b \right)^{-1}\left(\mathbb{I}_h y - \mathbb{I}_h \alpha_{h-1} - \frac{1}{\lambda}\mathbb{I}_h \Phi(A)\Phi(A)^T\alpha_{h-1}\right)
\end{align*}

Notice that the term $\mathbb{I}_h \Phi(A)\Phi(A)^T \alpha_{h-1}$ does not permit the use of an auxiliary vector to represent $\Phi(A)^T \alpha_{h-1}$ due to the kernelization. This means that we will have to explicitly compute $\mathbb{I}_h\Phi(A)\Phi(A)^T$ at every iteration. Additionally, the quantity $\mathbb{I}_h\Phi(A)\Phi(A)^T\mathbb{I}_h^T$ does not need to be computed at every iteration because we can simply select the relevant columns (namely the columns defined by right-multiplying by $\mathbb{I}_h^T)$ from the quantity $\mathbb{I}_h\Phi(A)\Phi(A)^T$.

# 3. Communication-Avoiding Subspace Iteration for Kernel Ridge Regression

In this section, we will derive the communication-avoiding variant of subspace iteration for solving the Kernel Ridge Regression problem. Lets start by re-stating the subspace iteration's update step:
\begin{align*}
    u_{h+1} &= \left(\frac{1}{\lambda} \mathbb{I}_{h+1}\Phi(A)\Phi(A)^T\mathbb{I}_{h+1}^T + I_b \right)^{-1}\left(\mathbb{I}_{h+1} y - \mathbb{I}_{h+1} \alpha_{h} - \frac{1}{\lambda}\mathbb{I}_{h+1} \Phi(A)\Phi(A)^T\alpha_{h}\right)\\
    \alpha_{h+1} &= \alpha_{h} + \mathbb{I}_{h+1}^T u_{h+1}
\end{align*}
Given this two-line recurrence, we can write the next update for iteration $h+2$ as:
\begin{align}
    u_{h+2} &= \left(\frac{1}{\lambda} \mathbb{I}_{h+2}\Phi(A)\Phi(A)^T\mathbb{I}_{h+2}^T + I_b \right)^{-1}\left(\mathbb{I}_{h+2} y - \mathbb{I}_{h+2} \alpha_{h+1} - \frac{1}{\lambda}\mathbb{I}_{h+2} \Phi(A)\Phi(A)^T\alpha_{h+1}\right)\tag{1}\\
    \alpha_{h+2} &= \alpha_{h+1} + \mathbb{I}_{h+2}^T u_{h+2}.\tag{2}
\end{align}
Notice that in the update step above, we can replace all instances of $\alpha_{h}$ with $\alpha_{h-1} + \mathbb{I}_h^T u_h$. By doing this, we can write the update steps for $u_{h+1}$ and $\alpha_{h+1}$ in terms of $\alpha_{h-1}$ and $u_h$, instead of $\alpha_h$:

\begin{align*}
    u_{h+2} &= \left(\frac{1}{\lambda} \mathbb{I}_{h+2}\Phi(A)\Phi(A)^T\mathbb{I}_{h+2}^T + I_b \right)^{-1}\left(\mathbb{I}_{h+2} y \color{blue}{- \mathbb{I}_{h+2} \alpha_{h} - \frac{1}{\lambda}\mathbb{I}_{h+2} \Phi(A)\Phi(A)^T\alpha_{h}} \color{red}{- \mathbb{I}_{h+2}\mathbb{I}_{h+1}^T u_{h+1} - \frac{1}{\lambda}\mathbb{I}_{h+2} \Phi(A)\Phi(A)^T\mathbb{I}_{h+1}^T u_{h+1}}\right)\tag{3}\\
    \alpha_{h+2} &= \alpha_{h} \color{red}{+ \mathbb{I}_{h+1}^Tu_{h+1} + \mathbb{I}_{h+2}^T u_{h+2}}.
\end{align*}
Comparing eq. (3) and (1), we can see that the terms in blue differ in that they depend on $\alpha_{h}$ instead of $\alpha_{h+1}$. The terms in red in eq. (3) are additional corrections which depend on the previous iteration's gradient, $u_{h+1}$, which also depend on $\alpha_{h}$. Finally, we can update $\alpha_{h+2}$ by summing the gradients $u_{h+1}$ and $u_{h+2}$.

The example above performs a recurrence unrolling of two iterations (i.e. computing $\alpha_{h+2}$ w.r.t $\alpha_h$), but we can unroll for an arbitrary number of iterations. We will use the variable $s$ to represent the number of iterations we unroll the recurrences. Notice that we still compute all of the gradients, $\{u_h, u_{h+1}, \ldots, u_{h+s}\}$, but only need $\alpha_h$ in order to compute $\alpha_{h+s}$. This means that we can defer updates to $\alpha$ for $s$ iterations.
\begin{align*}
u_{h + j} &= \left(\frac{1}{\lambda} \mathbb{I}_{h + j}\Phi(A)\Phi(A)^T\mathbb{I}_{h + j}^T + I_b\right)^{-1}\left(\mathbb{I}_{h+j} y \color{blue}{- \mathbb{I}_{h+j}\alpha_{h} - \frac{1}{\lambda}\mathbb{I}_{h+j}\Phi(A)\Phi(A)^T\alpha_{h}} \color{red}{- \sum_{k = 1}^{j-1}\mathbb{I}_{h + j}\mathbb{I}_{h + k}^T u_{h + k} - \sum_{k = 1}^{j-1} \frac{1}{\lambda}\mathbb{I}_{h+j}\Phi(A)\Phi(A)^T\mathbb{I}_{h + k}^T u_{h + k}}\right)\\
\alpha_{h+j} &= \alpha_{h} \color{red}{+ \sum_{k = 1}^{j} \mathbb{I}_{h + k}^T u_{h + k}}.\\
\text{for } j &= \{1, 2, \ldots, s\} \text{ where $s$ is a free parameter.}
\end{align*}

Notice that in the summation $\sum_{k = 1}^{j-1} \mathbb{I}_{h+k}\mathbb{I}_{k + j}^T u_k$, the computation $\mathbb{I}_{h+k}\mathbb{I}_k^T$ is computing a $b \times b$ identity-like matrix. This matrix is essentially computing the intersection between the sampled indices from iteration $h + j$ and iterations $h + k$ (where $k$ goes from $1$ to $j - 1$). This is crucial to ensure that gradient updates computed at iterations $h + k$ are incorporated into the iteration $h + j$.


```python
import numpy as np
import scipy as sp
from sklearn.metrics import pairwise_distances
import matplotlib.pyplot as plt


class Kernel_Ridge_Regression(object):
    def __init__(self, data, labels, kernel_params={'kernel_type': 'gaussian', 'sigma': 1e-1}):
        self.kernel_params = kernel_params
        self.nsamples, self.nfeatures = data.shape
        if self.nsamples != len(labels):
            raise ValueError('len(labels) != data.shape[1].')
        self.data = data
        self.labels = labels
        self._set_kernel_function(kernel_params['kernel_type'])
        self.alpha = np.zeros(self.nsamples)

    def _set_kernel_function(self, kernel_type):
        kernel_type = kernel_type.lower()
        if kernel_type == 'gaussian':
             self.kernel_function = self._gaussian_kernel
        elif kernel_type == 'polynomial':
            self.kernel_function = self._polynomial_kernel
        else:
            raise ValueError("Invalid kernel choice: {}".format(kernel_type))
    
    def _gaussian_kernel(self, subsamples, samples):
        kernel_matrix = pairwise_distances(subsamples, samples, metric='euclidean')**2
        kernel_matrix = np.exp(-kernel_matrix/self.kernel_params['sigma']**2)
        return kernel_matrix
    
    def _polynomial_kernel(self, subsamples, samples):
        kernel_matrix = (subsamples.dot(samples.T) + self.kernel_params['shift']) ** self.kernel_params['dim']
        return kernel_matrix
    
    def compute_objective(self):
        obj_val = np.linalg.norm(self.alpha - self.labels)**2
        obj_val += 1/self.C * (self.alpha.T.dot(self.kernel_function(self.data, self.data).dot(self.alpha)))
        return obj_val
    
    def block_coordinate_descent(self, idxs):
        blkrow_kernel_matrix = self.kernel_function(self.data[idxs, :], self.data)
        residual = self.labels[idxs] - self.alpha[idxs] - (1/self.C)*blkrow_kernel_matrix.dot(self.alpha)
        blkdiag_kernel_matrix = ((1/self.C) * blkrow_kernel_matrix[:, idxs]) + 1
        gradient = np.linalg.solve(blkdiag_kernel_matrix, residual)
        self.alpha[idxs] += gradient

    def commavoid_block_coordinate_descent(self, idxs):
        gradient = np.zeros(self.s*self.block_size)
        blkrow_kernel_matrix = self.kernel_function(self.data[idxs, :], self.data)
        residual = self.labels[idxs] - self.alpha[idxs] - (1/self.C)*blkrow_kernel_matrix.dot(self.alpha)
        # compute 's' gradients.
        blk_residual = residual[0:self.block_size]
        blkdiag_kernel_matrix = ((1/self.C) * blkrow_kernel_matrix[0:self.block_size, idxs[0:self.block_size]]) + 1
        gradient[0:self.block_size] = np.linalg.solve(blkdiag_kernel_matrix, blk_residual)
        for i in range(1, self.s):
            residual_correction = np.zeros(self.block_size)
            blk_idxs = range(i*self.block_size, (i+1)*self.block_size)
            blk_residual = residual[blk_idxs]
            blkdiag_kernel_matrix = ((1/self.C) * blkrow_kernel_matrix[blk_idxs, idxs[blk_idxs]]) + 1
            for j in range(i):
                blk_gradient = gradient[j*self.block_size:(j+1)*self.block_size]
                innerblk_idxs = idxs[j*self.block_size, (j+1)*self.block_size]
                commonvals, idxof_outerblk, idxof_innerblk = np.intersect1d(blk_idxs, innerblk_idxs, assume_unique=True, return_indices=True)
                residual_correction[idxof_outerblk] = blk_gradient[idxof_innerblk]
                residual_correction += (1/self.C) * (blkrow_kernel_matrix[blk_idxs, j*self.block_size:(j+1)*self.block_size].dot(blk_gradient))
            blk_residual -= residual_correction
            gradient[blk_idxs] = np.linalg.solve(blkdiag_kernel_matrix, blk_residual)
        self.alpha[idxs] += gradient

    def train(self, epochs=100, optimizer='bcd', optimizer_params={'C': 1e-5, 'block_size': 1, 'shuffle': True}):
        if optimizer.lower() == 'bcd':
            optimizer = self.block_coordinate_descent
        elif optimizer.lower() == 'cabcd':
            optimizer = self.commavoid_block_coordinate_descent
            self.s = optimizer_params['s']
        else:
            raise ValueError('Invalid optimizer choice: {}'.format(optimizer))
        self.C = optimizer_params['C']
        self.block_size = optimizer_params['block_size']
        self.shuffle = optimizer_params['shuffle']
        loss_history = []
        for epoch in range(1,epochs+1):
            if self.shuffle: 
                idxs = np.random.choice(self.nsamples, size=self.nsamples, replace=False)
            else:
                idxs = np.arange(self.nsamples)
            nblocks = len(idxs)/self.block_size
            for block in np.array_split(idxs, nblocks):
                if len(block) < 1:
                    break
                optimizer(block)
            loss_history.append(self.compute_objective())
            epoch_strlen = len(str(epochs))
            print('[{}] loss {:.16f}.'.format(str(epoch).zfill(epoch_strlen), loss_history[-1]))
        plt.plot(loss_history)

    def test(self):
        raise NotImplementedError()

```


```python
from sklearn.datasets import load_svmlight_file

(data, labels) = load_svmlight_file('../data/abalone')
krr = Kernel_Ridge_Regression(data, labels)
krr.train(epochs=100)
```

c9247d8d13226c939234417a949ac6a8eec643d5
ipynb
Jupyter Notebook
notebooks/kernel_ridge.ipynb
aditya08/commavoid-kernel-ml
a3459290e54eedabc4ae254d06c3846e0a943f00
BSD-3-Clause
notebooks/kernel_ridge.ipynb
aditya08/commavoid-kernel-ml
a3459290e54eedabc4ae254d06c3846e0a943f00
BSD-3-Clause
2021-01-11T16:03:26.000Z
2021-08-30T19:26:32.000Z
notebooks/kernel_ridge.ipynb
aditya08/commavoid-kernel-ml
a3459290e54eedabc4ae254d06c3846e0a943f00
BSD-3-Clause
Qwen/Qwen-72B
1. YES
2. YES