Intro

In the last post we discussed the meaning of our model family

where each \(\boldsymbol A_i\) is a symmetric matrix. In the last post we discussed what these models predict, and how we can explain them to ourselves and other stakeholders. Before that, we also discussed GPU acceleration to make training and inference faster. Speed is important, but so is cost, and fast GPUs may be expensive. Here, our aim is not only to make it faster, but also cheaper, by making the eigenvalue problem easier to solve even on weaker hardware. We certainly should not be paying for a GPU and waiting more than 5 minutes to train one neuron on a tabular dataset with about 20k rows, even if this one neuron is a fairly complex one! We begin our exploration from theory, which immediately yields practical applications. And as always, we have a notebook to reproduce all experiments in this post.

Simultaneous simplification

Recall that for any orthogonal matrix \({\boldsymbol Q} \in \mathbb{R}^{d \times d}\), we have

So our model family is invariant under such orthogonal similarity transformations, meaning a model with matrices \(\boldsymbol A_i\) is identical to a model with matrices \(\boldsymbol Q^T \boldsymbol A_i \boldsymbol Q\) for any orthogonal \(\boldsymbol Q\).

One of the interesting phenomena in linear algebra is simultaneous diagonalization. A set of matrices \({\boldsymbol A}_i\) is simultaneously diagonalizable if there exists an orthogonal matrix \({\boldsymbol Q}\) such that \({\boldsymbol Q}^T {\boldsymbol A}_i {\boldsymbol Q}\) is diagonal for all \(i\). In other words, the same matrix \(\boldsymbol Q\) diagonalizes all matrices simultaneously.

If we restrict ourselves to models where all of our learned matrices are simultaneously diagonalizable, we can just assume all matrices are diagonal:

So what is the \(k\)-th eigenvalue of this matrix? It’s just the \(k\)-th smallest entry of the vector

On the one hand, it’s an extremely easy eigenvalue problem. But we actually lost almost all of the expressive power, since it’s just a convoluted way to describe a piecewise linear function of \({\boldsymbol x}\). We have ReLU networks for that.

But there is another family of matrices for which the eigenvalue problem is easy - symmetric tridiagonal matrices, meaning, matrices of the form:

Such a matrix is defined by two vectors, the main diagonal \(\boldsymbol a \in \mathbb{R}^d\), and the off-diagonal \(\boldsymbol b \in \mathbb{R}^{d-1}\). Turns out this family strikes a nice balance - eigenvalues of such matrices are efficient to compute, while remaining fairly expressive. Efficiency comes from standing on the shoulders of giants, and using decades of numerical analysis research, given to us in the form of scipy.linalg.eigh_tridiagonal on a silver platter.

To appreciate the speed difference, let’s time eigenvalue and eigenvector computation using SciPy for regular dense matrices, and compare it to tridiagonal matrices. Let’s create a batch of dense matrices:

import scipy.linalg as sla
import numpy as np

# batch of 50 matrices of size 100x100
M = np.random.randn(50, 100, 100)

If you recall from previous posts - we need eigenvectors, in addition to eigenvalues, to compute gradients to train. Now let’s measure eigenvalue and eigenvector computation time. Here it is for eigenvalues:

%%timeit -n 100 -r 30
sla.eigvalsh(M, subset_by_index=(50, 50)).sum()  # 50-th eigenvalue

32.7 ms ± 4.42 ms per loop (mean ± std. dev. of 30 runs, 100 loops each)

Here it is for eigenvectors:

%%timeit -n 100 -r 30
vals, vecs = sla.eigh(M, subset_by_index=(50, 50))
vecs.sum()

34.5 ms ± 4.5 ms per loop (mean ± std. dev. of 30 runs, 100 loops each)

Alright. Now let’s do it for tridiagonal matrices. First, we generate diagonal and off-diagonal vectors:

# A batch of 50 diagonal and 50 off-diagonal vectors for 100x100 matrices.
d = np.random.randn(50, 100)
e = np.random.randn(50, 99)

Now let’s measure. Here is eigenvalue measurement:

%%timeit -n 100 -r 30
sla.eigvalsh_tridiagonal(d, e, select='i', select_range=(50, 50)).sum()

5.11 ms ± 60.1 µs per loop (mean ± std. dev. of 30 runs, 100 loops each)

Here is eigenvector measurement:

%%timeit -n 100 -r 30
vals, vecs = sla.eigh_tridiagonal(d, e, select='i', select_range=(50, 50))
vecs.sum()

6.68 ms ± 105 µs per loop (mean ± std. dev. of 30 runs, 100 loops each)

Between 5x and 6x speedup! Speed is not all we need - we also need representation power, which we shall explore in the next section.

Tridiagonal eigenvalue functions

In the last post, we saw that we can re-write our eigenvalue models as optimization problems over quadratic functions:

where

So we have a latent variable \(\boldsymbol u\) that appears in the quadratic function \({\boldsymbol u}^T \mathcal{A}(\boldsymbol x) \boldsymbol u\), which expresses interactions between all entry pairs \(\boldsymbol u\), since:

If \(\mathcal{A}(\boldsymbol x)\) were diagonal, we would lose all interactions - each entry \(u_i\) interacts only with itself:

This is another manifestation of the loss of expressiveness we discussed before. But if it were tri-diagonal, we do have pairwise interactions:

Even though it’s only between adjacent pairs \(u_i\) and \(u_{i+1}\), it turns out to be enough to produce a fairly rich set of models. Note, these are pairwise interactions between entries of the latent variable \(\boldsymbol u\), not of the raw features \(\boldsymbol x\). In fact, all features of \(\boldsymbol x\) potentially with each other, since each entry of \(\mathcal{A}(\boldsymbol x)\) contains a linear combination of all features.

To visually see that we have nontrivial expressive power, let’s try plotting a univariate function:

where the two matrices are tridiagonal, meaning specified by their diagonal and off-diagonal vectors. Here is a simple implementation of \(f_k(x)\) above:

def tridiagonal_eig_1d(k, diag, off_diag, xs):
    r"""Univariate matrix pencil eigenvalue.
        f(x) = \lambda_k(A + x B)
    where A and B are both tridiagonal.

    Args:
        k (int): the eigenvalue index
        diag (array): a 2 x n array of the diagonals of  A and B
        off_diag (array): a 2 x (n - 1) array of the off-diagonals of A and B
        xs (array): a vector of values x to evaluate f(x) at.

    Returns:
        An array y with y[i] = f(x[i])
    """

    padded_xs = np.c_[np.ones_like(xs), xs]
    mat_diag = padded_xs @ diag         # m x n
    mat_off_diag = padded_xs @ off_diag # m x (n - 1)
    eigval = sla.eigvalsh_tridiagonal(
        mat_diag, mat_off_diag, select='i', select_range=(k, k)
    )
    return eigval

Let’s try plotting a function obtained from random \(5 \times 5\) matrices. Below is a function that plots a grid of eigenvalue functions \(f_k(x)\) for all \(k\), followed by its use to plot our functions:

import matplotlib.pyplot as plt
import math


def plot_tridiag_eig_1d(diag, off_diag, xmin=-3, xmax=3, resolution=1000, fn=tridiagonal_eig_1d):
    dim = diag.shape[1]
    n_rows = int(math.sqrt(dim))
    n_cols = int(math.ceil(dim / n_rows))
    xs = np.linspace(xmin, xmax, resolution)
    fig, axs = plt.subplots(n_rows, n_cols, layout='constrained')
    for i, ax in zip(range(dim), axs.ravel()):
        ax.plot(xs, fn(i, diag, off_diag, xs))
        ax.set_title(f'$\\lambda{1+i}$')
    fig.show()


np.random.seed(42)
plot_tridiag_eig_1d(np.random.randn(2, 5), np.random.randn(2, 4))

Alright! We see functions having non-trivial shapes. As expected from what we saw in previous posts, the smallest eigenvalue \(\lambda_1\) is concave, the largest \(\lambda_5\) is convex, and all other eigenvalue functions have piecewise-smooth shapes that are neigher convex nor concave. What about \(11\times 11\) matrices?

plot_tridiag_eig_1d(np.random.randn(2, 11), np.random.randn(2, 10))

As expected, larger random matrices produce “wilder” shapes - the set of the functions is richer.

Now that we’ve convinced ourselves that tridiagonal matrices have some potential, as a family providing a reasonable balance between speed and expressiveness, let’s move on to a more convincing demonstration of that potential.

Training tridiagonal matrix eigenvalue models

If we want to be able to train with PyTorch, we first need to make sure we can enjoy fast tridiagonal eigenvalue computation there as well. Unfortunately, as of now (PyTorch 2.10), we do not have fast tridiagonal eigenvalue routines in PyTorch, even though tridiagonal and banded matrices do appear in many scientific computing domains. So similarly to a previous post, we will have to implement a custom autograd function that will forward PyTorch tensors to SciPy routines.

As a reminder - we need to subclass torch.autograd.Function and implement two static methods - forward for the computation and backward for the back-propagation of derivatives. This is exactly where we need eigenvectors, and not only the eigenvalues, as we explained in this previous post in the series. As a reminder, for the function \(\lambda_k(\boldsymbol X)\), the “right kind” of generalized derivative is the matrix \(\boldsymbol q_k \boldsymbol q_k^T\), where \(\boldsymbol q_k\) is the corresponding eigenvector. When \(\boldsymbol X\) is tridiagonal, we just need the diagonal and off-diagonal vectors of the \(\boldsymbol q_k \boldsymbol q_k^T\).

So below an autograd function implementing exactly this idea. It appears a bit lengthy, but that’s primarily because it aims to be efficient, and distinguish between two cases: (a) when we need derivatives, e.g., training, and require an eigenvector, and (b) when we do not need derivatives, e.g., inference, and do not require an eigenvector:

import torch

class TridiagEigvalsh(torch.autograd.Function):
    @staticmethod
    def forward(ctx, diag: torch.Tensor, off_diag: torch.Tensor, k: int):
        """Eigenvalue of batch of tridiagonal matrices.

        Args:
            diag (tensor): A M1 x ... x Mn x N tensor representing a batch
                of size M1 x ... x Mn of diagonals of NxN tridiagonal symmetric
                matrices.
            off_diag (tensor): A M1 x ... x Mn x (N - 1) tensor representing
                a batch of size M1 x ... x Mn of off-diagonals of NxN
                tridiagonal symmetric matrices.
            k (int): The eigenvalue index
        """
        need_grad = ctx.needs_input_grad[0] or ctx.needs_input_grad[1]

        diag_np = diag.detach().numpy()
        off_diag_np = off_diag.detach().numpy()
        if need_grad:
            # k-th eigenvalue and eigenvector
            ws_np, Qs_np = sla.eigh_tridiagonal(
                diag_np, off_diag_np, select='i', select_range=(k, k),
                lapack_driver="stemr"
            )
            ws = torch.as_tensor(ws_np, dtype=diag.dtype)
            Qs = torch.as_tensor(Qs_np, dtype=diag.dtype)
            ctx.save_for_backward(Qs.squeeze(-1))
        else:
            # only k-th eigenvalue
            ws_cp = sla.eigvalsh_tridiagonal(
                diag_np, off_diag_np, select='i', select_range=(k, k)
            )
            ws = torch.as_tensor(ws_cp, dtype=diag.dtype)

        return ws.squeeze(-1) # k-th eigenvalue

    @staticmethod
    def backward(ctx, grad_w: torch.Tensor):
        (Qs,) = ctx.saved_tensors  # (..., N) from SciPy

        grad_w = grad_w.to(dtype=Qs.dtype)                 # (...)
        gw = grad_w.unsqueeze(-1)                          # (..., 1)

        grad_diag = gw * Qs.square()                       # (..., N)
        grad_off  = 2 * gw * (Qs[..., :-1] * Qs[..., 1:])  # (..., N-1)

        return grad_diag, grad_off, None

Now let’s try it out. Here is a batch of diagonals of 50 matrices of size 100x100:

diags = torch.randn(50, 100)
off_diags = torch.randn(50, 99)

Now let’s try applying our PyTorch function to the raw tensors. Note - they do not require a gradient, since they aren’t trainable parameters, so we’re going through the no_grad path:

%%timeit -r 30 -n 100
w = TridiagEigvalsh.apply(diags, off_diags, 2).sum()

3.71 ms ± 91.7 µs per loop (mean ± std. dev. of 30 runs, 100 loops each)

Whoa! That’s fast! Now let’s try doing it with trainable parameters:

diags_param = torch.nn.Parameter(diags)
off_diags_param = torch.nn.Parameter(off_diags)

%%timeit -r 30 -n 100
w = TridiagEigvalsh.apply(diags_param, off_diags_param, 2).sum()
w.backward()

4.88 ms ± 205 µs per loop (mean ± std. dev. of 30 runs, 100 loops each)

Pretty fast - a mini-batch of 50 tridiagonal matrices of size 100x100 can compute gradients in approximately 5 milliseconds. Comparing it with approximately 35 milliseconds for full dense matrices - quite a speedup. For convenience, let’s wrap our autograd function class with a simple Python function:

def tridiag_eigvalsh(
        diag: torch.Tensor, off_diag: torch.Tensor, k: int
    ) -> torch.Tensor:
    return TridiagEigvalsh.apply(diag, off_diag, k)

So now, to train a model we need a torch module representing our \(f(\boldsymbol x, \boldsymbol A_{0..n})\) for the tri-diagonal case. This means our trainable parameters are the diagonals and the off-diagonals of the matrices \(\boldsymbol A_0, \dots, \boldsymbol A_n\). Note that both the diagonal vector and the off-diagonal vector of \(\mathcal{A}(\boldsymbol x)\) are just linear functions of \(\boldsymbol x\), so we can express them as simple torch.nn.Linear layers. This yields an almost magically simple class:

from torch import nn

class TridiagSpectral(nn.Module):
    def __init__(self, *, num_features: int, dim: int, eig_idx: int):
        super().__init__()
        self.eig_idx = eig_idx
        self.diag = nn.Linear(num_features, dim)
        self.off_diag = nn.Linear(num_features, dim - 1)

    def forward(self, x):
        return tridiag_eigvalsh(self.diag(x), self.off_diag(x), self.eig_idx)

Now we can use it for training, like any PyTorch model. So let’s try learning a classifier that detects whether we have either two or five ones in a vector:

def toy_function(x: torch.Tensor):
    return torch.maximum(
        x.sum(axis=-1) == 2,
        x.sum(axis=-1) == 5
    ).to(dtype=torch.float32)

Now we shall apply it to learning this function over 12-dimensional vectors. So let’s generate all binary vectors and compute their true label:

n_features = 12
X = torch.cartesian_prod(*([torch.tensor([0., 1.])] * n_features)) 
y = toy_function(X)

This set should contain \(2^12 = 4096\) vectors. And before training, let’s divide the features and labels into a train and evaluation set:

torch.manual_seed(42)
train_mask = torch.rand(len(X)) < 0.5
X_train = X[train_mask, :]
y_train = y[train_mask]
X_test = X[~train_mask, :]
y_test = y[~train_mask]

Alright! We’re ready to train a classifier on (X_train, y_train) and evaluate it on (X_test, y_test). This would be a good time to introduce the fitstream library, which is very convenient for training PyTorch models on small in-memory datasets. Recall that we found it very convenient to hide the training loop behind a Python generator that yields an event on every epoch. So this is what this library does - it performs a pretty standard PyTorch training loop, and yields a dict with some data at the end of each epoch. Let’s first install it in our notebook:

%pip install -q fitstream

Now let’s use it. Below is a short snippet demonstrating how we iterate over the first 3 events, which are simple Python dicts, and use Python’s pprint library to nicely print each dict:

import fitstream as fts
from pprint import pprint

# define model and optimizer
dim = 3
model = TridiagSpectral(num_features=n_features, dim=dim, eig_idx=dim // 2)
optim = torch.optim.Adam(model.parameters(), lr=1e-1)

# use FitStream to obtain the event generator
events = fts.epoch_stream(
    (X_train, y_train), model, optim, nn.BCEWithLogitsLoss(), batch_size=32
)

# iterate over the first three events
for _, event in zip(range(3), events):
    pprint(event)
    print('---')

{'model': TridiagSpectral(
  (diag): Linear(in_features=12, out_features=3, bias=True)
  (off_diag): Linear(in_features=12, out_features=2, bias=True)
),
 'step': 1,
 'train_loss': 0.47607117891311646,
 'train_time_sec': 0.19717628799844533}
---
{'model': TridiagSpectral(
  (diag): Linear(in_features=12, out_features=3, bias=True)
  (off_diag): Linear(in_features=12, out_features=2, bias=True)
),
 'step': 2,
 'train_loss': 0.4695621430873871,
 'train_time_sec': 0.17527212999993935}
---
{'model': TridiagSpectral(
  (diag): Linear(in_features=12, out_features=3, bias=True)
  (off_diag): Linear(in_features=12, out_features=2, bias=True)
),
 'step': 3,
 'train_loss': 0.46067821979522705,
 'train_time_sec': 0.1757020229997579}
---

Now we see what we get - each dict contains our model, the epoch index in the step key, the training loss, and the training time in seconds. Pretty minimal, so the library comes with some helper functions to take this minimal event stream, and enrich it.

It has the pipe function, which lets us pipe the event stream through a sequence of transformations. So let’s introduce the first transformation - take, which simply takes the head of the event stream of the specified size. For example, this will produce 5 events:

# define model and optimizer
dim = 3
model = TridiagSpectral(num_features=n_features, dim=dim, eig_idx=dim // 2)
optim = torch.optim.Adam(model.parameters(), lr=1e-1)

# pipe the epoch stream through the "take" transformation
events = fts.pipe(
    fts.epoch_stream(
        (X_train, y_train), model, optim, nn.BCEWithLogitsLoss(), batch_size=32
    ),
    fts.take(5)
)

for event in events:
    print(event['step'], event['train_loss'])

1 0.4920884072780609
2 0.46535375714302063
3 0.4354994297027588
4 0.4041599929332733
5 0.32225537300109863

The second important transformation is augment, which adds additional keys to each event. Here is an example of adding a key with the training loss squared:

# define model and optimizer
dim = 3
model = TridiagSpectral(num_features=n_features, dim=dim, eig_idx=dim // 2)
optim = torch.optim.Adam(model.parameters(), lr=1e-1)

# pipe the epoch stream through the "take" transformation
events = fts.pipe(
    fts.epoch_stream(
        (X_train, y_train), model, optim, nn.BCEWithLogitsLoss(), batch_size=32
    ),
    fts.augment(lambda event: {"loss_squared": event["train_loss"] ** 2}),
    fts.take(5)
)

for event in events:
    print(event['step'], event["train_loss"], event['loss_squared'])

1 0.4896900951862335 0.23979638932350245
2 0.45949652791023254 0.21113705916155912
3 0.46307340264320374 0.2144369762355547
4 0.44300273060798645 0.19625141932613221
5 0.4168809950351715 0.1737897640215147

Now, this is all nice, but the library is richer. It comes with augmentations for adding a validation loss, early stopping, or simply executing a function on every event. So here is a generator of events for a full-fledged training procedure with a validation loss and early stopping that also prints the losses every 10 epochs:

dim = 3
model = TridiagSpectral(num_features=n_features, dim=dim, eig_idx=dim // 2)
optim = torch.optim.Adam(model.parameters(), lr=1e-1)
events = fts.pipe(
    fts.epoch_stream(
        (X_train, y_train), model, optim, nn.BCEWithLogitsLoss(), batch_size=64
    ),
    fts.augment(fts.validation_loss((X_test, y_test), nn.BCEWithLogitsLoss())),
    fts.early_stop(key="val_loss", patience=5),
    fts.tap(fts.print_keys("train_loss", "val_loss"), every=10),
    fts.take(100)
)

Finally, the library comes with a set of collector functions that iterate over the events and collect them into various data structures. Here it will be convenient to use collect_pd, which collects the event dicts into a Pandas DataFrame. So here is an example of collecting the above event stream into a data frame, and then plotting the training and validation losses:

training_log = fts.collect_pd(events)
training_log.plot(x="step", y=["train_loss", "val_loss"], title='Dim = 3')

step=0001 train_loss=0.4908 val_loss=0.4681
step=0011 train_loss=0.3369 val_loss=0.3147
step=0021 train_loss=0.1540 val_loss=0.1516
step=0031 train_loss=0.0768 val_loss=0.0926
step=0041 train_loss=0.0500 val_loss=0.0725
step=0051 train_loss=0.0363 val_loss=0.0621
step=0061 train_loss=0.0291 val_loss=0.0580
step=0071 train_loss=0.0241 val_loss=0.0557
step=0081 train_loss=0.0210 val_loss=0.0538
step=0091 train_loss=0.0186 val_loss=0.0536
step=0101 train_loss=0.0172 val_loss=0.0539

Nice! We see that the model is learning, which is encouraging. Before we do more experiments, let’s write a function that will construct such a training procedure for us and collect the events to a dataframe:

def run_experiment_bce(model, lr=1e-1, batch_size=64, max_epochs=100):
    optim = torch.optim.Adam(model.parameters(), lr=lr)
    events = fts.pipe(
        fts.epoch_stream(
            (X_train, y_train), model, optim, nn.BCEWithLogitsLoss(), batch_size=batch_size
        ),
        fts.augment(fts.validation_loss((X_test, y_test), nn.BCEWithLogitsLoss())),
        fts.early_stop(key="val_loss", patience=5),
        fts.take(max_epochs)
    )
    return fts.collect_pd(events)

Now we can easily plot similar losses for \(5 \times 5\) tridiagonal matrices:

training_log_5 = run_experiment_bce(
    TridiagSpectral(num_features=n_features, dim=5, eig_idx=2)
)
training_log_5.plot(x="step", y=["train_loss", "val_loss"], title='Dim = 5')

Much better! Now the validation loss is very close to zero as well. Now let’s move to \(9 \times 9\) matrices:

training_log_9 = run_experiment_bce(
    TridiagSpectral(num_features=n_features, dim=9, eig_idx=4)
)
training_log_9.plot(x="step", y=["train_loss", "val_loss"], title='Dim = 9')

Beautiful! Apparently, a model with \(9 \times 9\) tridiagonal symmetric matrices, which has \(13 \times (9 + 8) = 221\) parameters, can learn this function from data almost perfectly. And conceptually, this is just a linear function of the features followed by a non-linear function - the matrix eigenvalue. Just one neuron! You can try it, but a “classical” neuron cannot learn this function.

So now that we’re convinced that the machinery is working, let’s try it on the dataset that accompanies this series - the California Housing dataset we have built into our Colab notebooks.

California housing training

Recall that the dataset is about predicting housing prices in California based on some features. I will skip the part where we read the data, normalize features and targets, and split the data into training and test sets. We’ve already done it in previous posts in this series, and the notebook contains the full code. So here we’ll assume our training data is in X_train, y_train, our evaluation set is X_test, y_test, and the number of features is in num_features. Moreover, since our labels are scaled, we also have label_scale, which is the factor that transforms the training / eval RMSE back to the original units in the dataset - dollars.

First, let’s define a simple function that computes the RMSE in dollars:

def scaled_rmse(y_true, y_pred):
    mse = nn.functional.mse_loss(y_pred, y_true)
    return torch.sqrt(mse) * label_scale

Now, we can define a full-fledged training procedure with the FitStream library we just introduced. When experimenting, I noticed that learning rate scheduling improves convergence substantially and I can work with less epochs, so I also used a learning rate scheduler with warmup - just like we do with LLMs. It first increases the learning rate for a few epochs (warmup), and then decreases it slowly towards zero (cooldown). It is implemented in the OneCycleLR class from PyTorch. So here is our full training procedure:

from torch.optim.lr_scheduler import OneCycleLR

def complete_training_stream(
        dim, n_epochs, warmup_fraction=0.1, lr=5e-3, batch_size=64,
    ):
    model = TridiagSpectral(num_features=num_features, dim=dim, eig_idx=dim // 2)
    optim = torch.optim.Adam(model.parameters(), lr=lr)
    sched = OneCycleLR(
        optim, max_lr=lr, total_steps=n_epochs, pct_start=warmup_fraction, anneal_strategy='linear'
    )

    epoch_events = fts.epoch_stream((X_train, y_train), model, optim, nn.MSELoss(), batch_size=batch_size)
    return fts.pipe(
        epoch_events,
        fts.take(n_epochs),
        fts.augment(fts.validation_loss((X_test, y_test), scaled_rmse)), # <-- here we use scaled_rmse
        fts.augment(lambda event: {"lr": optim.param_groups[0]['lr']}),
        fts.early_stop(key="val_loss", patience=n_epochs // 10),
        fts.tick(sched.step),
    )

Note where we use our scaled_rmse - it is inserted as the validation loss to the stream. Now, let’s try it out with 11-dimensional matrices for 20 epochs:

training_log = fts.collect_pd(complete_training_stream(11, 20))
print(training_log)

This is what I got:

step  train_loss  train_time_sec       val_loss        lr
0      1    0.893857        1.095248  101073.101562  0.000200
1      2    0.404809        1.114020   65408.109375  0.005000
2      3    0.298268        1.138942   62557.234375  0.004722
3      4    0.277326        1.107545   61244.425781  0.004444
4      5    0.268057        1.099500   60413.164062  0.004167
5      6    0.262378        1.105777   59918.421875  0.003889
6      7    0.256854        1.100250   59919.507812  0.003611
7      8    0.253308        1.095398   59363.906250  0.003333
8      9    0.250837        1.091576   59104.089844  0.003056
9     10    0.248961        1.083466   59065.980469  0.002778
10    11    0.246233        1.092109   58896.339844  0.002500
11    12    0.244495        1.114625   58766.113281  0.002222
12    13    0.241918        1.118336   58610.593750  0.001944
13    14    0.240912        1.101395   58488.941406  0.001667
14    15    0.239700        1.100444   58310.894531  0.001389
15    16    0.238621        1.096240   58618.683594  0.001111
16    17    0.237742        1.085707   58579.175781  0.000833

We can see the model is training, the learning rate increased in the first two epochs, as expected, since 10% of the epochs are warmup. It stopped after 17 epochs due to the early stopping mechanism whose patience is two epochs (again, 10% of the maximum).

We can also write a nice function for plotting the learning rate and the validation loss. It’s a bit of boilerplate, for using the primary y-axis for the validation loss, and the secondary y-axis for the learning rate.

def plot_log(log, title=None):
    fig, ax = plt.subplots()

    ax.plot(log.step, log.val_loss, color='blue',
            label=f'RMSE (best=${log.val_loss.min():.2f})')
    ax.set_ylabel("Error ($)")
    ax.grid()

    lr_ax = ax.twinx()
    lr_ax.plot(log.step, log.lr, label='Learning rate',
               color='gray', linestyle='dotted', linewidth=1)
    lr_ax.set_ylabel("Learning rate")

    ax.legend()
    if title is not None:
        ax.set_title(title)
    fig.show()

plot_log(training_log)

In blue we see the validation loss, whereas in dotted black we see the learning rate. We can nicely see the warmup and cooldown stages.

Alright, so now that we have all the machinery in place, let’s try training some model with more epochs. I used 500 epochs in all the experiments, which was enough to train both smaller and larger models. So let’s try 7-dimensional tridiagonal matrices:

training_log_7 = fts.collect_pd(complete_training_stream(7, 500))
plot_log(training_log_7, title='Dim=7')

How about 11-dimensional tridiagonals?

training_log_11 = fts.collect_pd(complete_training_stream(11, 500))
plot_log(training_log_11, title='Dim=11')

Nice! Increasing matrix size reduces the error, meaning that performance scales with model size. But remember - it is just one neuron! How about \(15 \times 15\) matrices?

training_log_15 = fts.collect_pd(complete_training_stream(15, 500))
plot_log(training_log_15, title='Dim=15')

Another slight improvement. What about \(45 \times 45\) matrices?

training_log_45 = fts.collect_pd(tqdm(complete_training_stream(45, 500)))
plot_log(training_log_45, title='Dim=45')

I can share, and you can see it by running the notebook yourself, that each such experiment takes 3-4 minutes. Just to get a feeling - compared to dense matrix experiments we conducted in previous posts, this is much faster, and without any GPU. I’m pretty sure that if PyTorch had tridiagonal support, we could have run each experiment in seconds. But unfortunately - it does not.

Comparing it to dense experiments we conducted with the same dataset and similar matrix sizes in this post, which took us 31 minutes on an NVidia L4 GPU for a \(45 \times 45\) matrix, while achieving a similar test error - we clearly see the difference. No GPU, an order of magnitude faster, and a similar performance at least on this dataset.

Of course - the above are not proper experiments I’d include in a paper. I haven’t conducted any hyperparameter search, perhaps a different optimizer could be better, etc…, but we see the point.

Summary

To summarize, we can see that restricting ourselves to eigenvalue model families where all matrices are simultaneously tri-diagonalizable can be useful to strike a good balance between speed and expressiveness. Let us recall why this model family is interesting - it’s just one neuron, a linear (matrix) function composed with a non-linearity, that is quite expressive, while being fairly interpretable. These nice properties haven’t gone anywhere - spectral norms of our tridiagonal matrices are still a reasonable way to think of importance, and provide a certificate for sensitivity of the model to changes in that feature.

We do, however, see slow convergence. 500 epochs is quite a lot, and even though our training procedure stops beforehand due to the early stopping mechanism, it’s still a few hundred epochs. Even if I throw the best practices at it, such as learning rate scheduling, early stopping, and others - it’s still quite slow. At this stage, this is a price we pay for having a model that is, on the one hand, just one fairly interpretable neuron, but on the other hand can be improved by scaling.

We have many more questions to explore in this series. For example - can we prune any dense eigenvalue model to tridiagonal form? Can we make it converge faster? How do we stack such models as layers of a larger neural network? Stay tuned!

We continue our discussion of machine-learned models of the form

where \({\boldsymbol A}_i\) are learned symmetric matrices, and \(\lambda_k\) is the \(k\)-th smallest eigenvalue. We touched on one aspect of interpretability in a previous post - the importance of each of the features \(x_1, \dots, x_n\). But there are other aspects: what does this model actually compute? How can we reason about it, or explain it to our colleagues or to a regulator? We began this series by saying that this is a “neuron” that is solving an optimization problem. So in this post we shall focus on the different kinds of optimization problems that \(f({\boldsymbol x}; {\boldsymbol A}_{0:n})\) solves, what interpretations we can give to them, and why we should care. All the results in this post are based on Chapter 4 of the Matrix Analysis book by Horn & Johnson, 2nd edition.

A game between two players

You’ve probably seen eigenvalues presented as “stretch factors”. But for our understanding of the model, the optimization view is often more useful. So here is a Courant min-max characterization of the \(k\)-th smallest eigenvalue, named after Richard Courant. It’s a bit “hairy”, so let’s first present it, and then interpret it:

This is a bi-level optimization problem, which we can think of as a two-turn game. The first player chooses \({\boldsymbol C}\) with \(k-1\) rows (i.e., \(k-1\) linear constraints). In response, the second player chooses a unit vector \({\boldsymbol u}\) that is in the null-space of \({\boldsymbol C}\), or equivalently, orthogonal to the rows of \({\boldsymbol C}\).

The objective of the second player is to pay as little as possible, where \({\boldsymbol u}^T {\boldsymbol A} {\boldsymbol u}\) is the cost. But the objective of the first player, of course, is to make their opponent pay as much as possible, so they choose a “worst case” matrix \({\boldsymbol C}\). In case you were wondering, when \(k = 1\) we have no adversarial player, and the vector \(\boldsymbol u\) can be an arbitrary unit vector. And you have probably guessed: at any equilibrium, \({\boldsymbol u}\) is an eigenvector corresponding to \(\lambda_k({\boldsymbol A})\).

One way to read this is: \({\boldsymbol u}\) is a bounded allocation over \(d\) latent resources, and \(A_{i,j}\) is the cost associated with every pairwise interaction of resources, since:

Depending on the context, you can give different interpretations. For example, \({\boldsymbol A}\) represents a set of latent skills a student possesses, and \({\boldsymbol u}\) is a test vector for pairs of skills.

There is a mirror-image of the above game, if we want to sort eigenvalues from largest to smallest. So the \(k\)-th largest eigenvalue, which is also the \(d-k+1\)-th smallest one, can be written as

We can think of it in terms of utility rather than cost - each entry in the matrix is a utility associated with a pair of resources. The first player is choosing the matrix so that the second player will get as little utility as possible, whereas the second player, in response, aims to choose a vector \({\boldsymbol u}\) that will maximize their utility.

Adopting the \(\max-\min\) convention, we can write our “neuron” as:

Consequently, each feature is associated with a matrix of some latent “costs” and our features are just the weights of these costs. Suppose you’re in the insurance business, and someone asks you what your model is doing. You can explain something like “oh, we’re representing each feature of the insured using a table of latent skills to avoid claims, we sum them up, and simulate a game where we aim to elicit their worst-case ability to either avoid damage or absorb it without claiming”. You can give an example of what those “latent skills” could be, just like in matrix factorization people explain what would be the latent features in a movie recommendation system.

As a (kind of) recurrent neural network

Another way to characterize symmetric matrix eigenvalues is as a sequence of optimization problems. Here is a Courant-Fischer formulation (Fischer here is Ernst Sigismund Fischer, not Ronald Fisher):

We just saw that the smallest eigenvalue is just the minimum of a quadratic function over the unit sphere. The second eigenvalue is similar, but the vector has to be orthogonal to one row, since the matrix \(\boldsymbol C\) the first player chooses has one row. The next eigenvalue should be orthogonal to two rows, and so on. But we can actually be more precise about what these constraints are. One of the possible formulations of the Courant-Fischer theorem is

Let \(\boldsymbol A\) be a symmetric matrix with eigenvalues \(\lambda_1 \leq \lambda_2 \leq \cdots \leq \lambda_d\) with corresponding eigenvectors \({\boldsymbol u}_1, \dots, {\boldsymbol u}_d\). Then,

\[\lambda_k = \min_{\boldsymbol u} \left\{ {\boldsymbol u}^T {\boldsymbol A} {\boldsymbol u} : \| {\boldsymbol u} \|_2 = 1,\,\langle {\boldsymbol u}, {\boldsymbol u}_1 \rangle = 0, \dots, \langle {\boldsymbol u}, {\boldsymbol u}_{k-1} \rangle = 0 \right\}\]

In other words, the \(k\)-th smallest eigenvalue is the minimum of \({\boldsymbol u}^T {\boldsymbol A} {\boldsymbol u}\) among all unit vectors orthogonal to eigenvectors corresponding to the previous eigenvalues. Thus, we can think of it as a recurrent process: computing each eigenvalue yields an eigenvector, and all eigenvectors up to \(k-1\) are used to compute the \(k\)-th eigenvalue. Visually, it looks like this:

All steps share the same matrix \({\boldsymbol A}\) as their weights, just like recurrent neural networks share weights in each recurrent step.

Intuitively, as we move from \(\lambda_1\) upward, the function becomes more and more expressive. But it is tempting (and wrong) to conclude that the largest eigenvalue is the most expressive function of \({\boldsymbol A}\). Indeed, we can construct a “mirror-image” of the above process if we order the eigenvalues in decreasing order. In practice, the richest behavior tends to come from the middle of the spectrum, as we saw in the plots.

Again, we can think of this process as a kind of a repeated “game”. This time there is only one player. In each turn the player aims to minimize their cost, but their “strategy” \({\boldsymbol u}\) in each turn becomes more and more restricted - they must try something “different”, or orthogonal to, the strategies they chose in the previous turns.

As a difference of convex functions

Minimization of nonconvex functions is a long-standing challenge in optimization theory and practice. But sometimes knowing some additional information about a nonconvex function can substantially improve both the speed and the reliability of our ability to minimize it. One of these pieces of information is having an explicit representation of the function we aim to minimize (or maximize) as a difference of convex (DC) functions, namely,

such that both \(g\) and \(h\) are convex. In fact, there is an entire stream of literature and algorithms on DC optimization, and many famous optimization software packages have dedicated code paths for this task. For example, the DCCP extension for CVXPY is a famous example.

Turns out our \(f({\boldsymbol x};{\boldsymbol A}_{0:n})\) has such an explicit representation and can be directly used in the DCCP extension, and other similar software packages. Why is it useful? Well, suppose \(f\) models the expected cost of some decision that we would like to minimize.

The idea is based on the Ky Fan Variational Principle, stating that the sum of the eigenvalues \(\lambda_k, \lambda_{k+1}, ..., \lambda_d\) can be written as

This looks a bit hairy, but the term we are maximizing is a linear function of \(\boldsymbol A\), even if it’s a nonlinear function of \(\boldsymbol U\). And the maximum of linear functions is always convex, even if we have an infinite number of linear functions. Consequently, the \(k\)-th smallest eigenvalue can be written in an explicit DC form as:

We can see it visually by plotting \(\lambda_k({\boldsymbol P} + x {\boldsymbol Q})\) and its two convex components as a function of \(x\):

import numpy as np
import matplotlib.pyplot as plt

# choose random matrices
np.random.seed(42)
P = np.random.randn(5, 5)
Q = np.random.randn(5, 5)

# compute eigenvalues of A + x B for x in [-3, 3]
xs = np.linspace(-3, 3, 1000)
eigvals = np.linalg.eigvalsh(P + xs[:, None, None] * Q)

# plot mid eigenvalue and its constituent convex functions
sum_top_3 = np.sum(eigvals[:, 2:], axis=-1)
sum_top_2 = np.sum(eigvals[:, 3:], axis=-1)
plt.plot(xs, sum_top_3, label=r'$\Lambda_3$')
plt.plot(xs, sum_top_2, label=r'$\Lambda_4$')
plt.plot(xs, eigvals[:, 2], label=r'$\lambda_3$')
plt.fill_between(
    xs, sum_top_3, sum_top_2,
    color='skyblue', alpha=0.2, where=(sum_top_3 > sum_top_2)
)
plt.fill_between(
    xs, sum_top_3, sum_top_2,
    color='red', alpha=0.2, where=(sum_top_3 < sum_top_2)
)
plt.grid(True, alpha=0.3)
plt.legend()

Indeed, the orange and blue plots, which are top eigenvalue sums, are convex functions. The gap between them is red when \(\Lambda_3 - \Lambda_4\) is negative, and blue when it is positive. The function in green is the difference, and it exactly reflects the size and the sign of the gap.

Recap

This was a theoretical detour, to understand what kind of functions are we fitting and how we can reason about them. We saw that we can interpret our “neuron” as a game between two players, as a recurrent process where each step solves a simple quadratic optimization problem, and as the difference between convex functions. In the next post we’re back to more practical stuff.

Intro

Efficiency is a quite superpower in research. When training is fast, you can iterate: try an idea, get surprised, debug, tune, and move on. When training is slow, you start avoiding experiments you should be running, simply because they cost too much time.

This became very tangible in this eigenvalue-model series. In the previous posts we looked at models that predict the \(k\)-th eigenvalue of a learned symmetric matrix built from the input features, and we explored what they can represent, plus some robustness/interpretability properties. Naturally, the next step was to scale up the matrix size and see what happens in practice.

I did not show it in the last post, but if you take the California Housing experiment and run \(30\times 30\) matrices for 500 epochs, it takes more than an hour on Colab. And this is with an L4 GPU. I even tried an A100, and it didn’t meaningfully improve anything. So then I asked myself - what’s wrong?

Turns out PyTorch is wrong. torch.linalg.eigvalsh has a note in the official documentation: when the input is on CUDA, it synchronizes the device with the CPU. If eigenvalues sit inside your inner training loop, that synchronization becomes the bottleneck, and the rest of your model almost doesn’t matter.

So this post is a practical detour: we’ll make eigenvalue computation fast enough that it stops getting in the way. We’ll replace the slow call with a faster GPU implementation, and we’ll wrap it in a way that still supports backprop through the \(k\)-th eigenvalue. Once that’s done, the scaling experiments from the previous post become feasible again, and we can go back to asking the interesting questions.

(All execution speeds I measure in this post are on Colab, with an NVIDIA L4 GPU, with the 2025.10 runtime.)

Warm-up

Let’s start from a simple eigenvalue computation test on the CPU. We’ll use Jupyter’s %%time magic keyword to measure time. First, let’s create a mini-batch of 500 matrices of size 100x100:

mats = torch.randn(500, 100, 100)

Now let’s see how fast we can compute the sum of eigenvalues of all these matrices:

%%time
torch.linalg.eigvalsh(mats).sum()

CPU times: user 223 ms, sys: 14.4 ms, total: 237 ms
Wall time: 188 ms
tensor(56.6666)

And now with NumPy:

import numpy as np

# in another cell
%%time
np.linalg.eigvalsh(mats.numpy()).sum()

CPU times: user 942 ms, sys: 7.07 ms, total: 949 ms
Wall time: 337 ms

np.float32(56.665924)

Apparently, on the CPU, NumPy is almost twice slower than PyTorch. So when our tensors are on the CPU, we can continue using PyTorch - it’s pretty fast.

Now let’s move to the GPU. Here is a similar piece of code - create a mini-batch of random matrices and compute the sum of their eigenvalues:

mats = torch.randn(500, 100, 100, device='cuda')
torch.linalg.eigvalsh(mats).sum()

The reason I ran the eigenvalue computation once is as a “warmup” - I want PyTorch to do whatever setup it needs to run CUDA kernels, so next time we invoke eigvalsh it is going to be a “clean” run not contaminated by setup:

%%time
torch.linalg.eigvalsh(mats).sum()

CPU times: user 788 ms, sys: 811 µs, total: 789 ms
Wall time: 788 ms
tensor(-557.2836, device='cuda:0')

Whoa! It’s twice slower than NumPy on CPU, and four times slower than PyTorch on the CPU! Turns out PyTorch developers haven’t invested that much in general-purpose scientific computing on the GPU. It’s quite reasonable - it is not their main focus. So if we want to propose a new computational tool - it’s up to us to make it efficient!

So maybe PyTorch hasn’t invested in eigenvalues on GPU that much, but it doesn’t mean other scientific computing libraries haven’t. CuPy, a library aiming to be “NumPy on CUDA”, is one of those libraries that has a very fast eigenvalue solver we can use. But how can we use it on PyTorch tensors?

Turns out there is a standard called DLPack for representing multi-dimensional tensors in memory, and it is supported both by PyTorch and by CuPy. In PyTorch we have the torch.utils.dlpack package for converting a tensor to a DLPack “capsule” - a wrapper around its memory with appropriate metadata:

from torch.utils import dlpack as torch_dlpack

x = torch.tensor([1, -2, 3], device='cuda')
torch_dlpack.to_dlpack(x)

<capsule object "dltensor" at 0x7b9a1ae26760>

We can use CuPy to consume this “capsule” and access the same tensor, but this time as a CuPy array:

import cupy as cp

cp.from_dlpack(torch_dlpack.to_dlpack(x))

array([ 1, -2,  3])

Now let’s try computing the sum of eigenvalues of our PyTorch mats tensor containing the mini-batch of matrices using CuPy:

cupy_mats = cp.from_dlpack(torch_dlpack.to_dlpack(mats))
cp.linalg.eigvalsh(cupy_mats).sum()

array(-557.284, dtype=float32)

We got the same value, so apparently it’s working. Let’s time it:

%%time
cp.linalg.eigvalsh(cupy_mats).sum()

CPU times: user 1.58 ms, sys: 0 ns, total: 1.58 ms
Wall time: 1.37 ms
array(-557.284, dtype=float32)

Now that’s FAST! 1.37 milliseconds, instead of more than 700 - almost 500 times faster! It’s impressive, but a part of the enormous speedup is because we aren’t doing anything to be prepared to backpropagate.

Of course we got a CuPy array of eigenvalues. But we can easily use DLPack to convert it back to a PyTorch tensor. Since there is no memory copy, it practically incurs no cost, as you can see below:

%%time
eigvals_cupy = cp.linalg.eigvalsh(cupy_mats)
torch_dlpack.from_dlpack(eigvals_cupy).sum()

CPU times: user 1.66 ms, sys: 0 ns, total: 1.66 ms
Wall time: 1.31 ms
tensor(-557.2838, device='cuda:0')

Here, the eigenvalues were computed with CuPy, but their sum was computed with PyTorch. You can see that we got the same result at the same speed, since DLPack conversions just wrap the same GPU memory block, without any copies.

So the process is simple:

1. Wrap our PyTorch tensor’s memory as a CuPy array via DLPack
2. Compute eigenvalues using CuPy
3. Convert eigenvalues back to PyTorch via DLPack

There are more nuances here about memory management, and which object is responsible for actually freeing the memory when it’s no longer needed, and you should learn these DLPack nuances if you wish to use it. But that’s out of the scope of this post.

What we have is still not enough to build a full-fledged function we can use for model training in PyTorch, since for training we also need gradients.

Eigenvalue gradients

Consider the function

of a symmetric matrix \(\boldsymbol X\). Recall from linear algebra that eigenvalues are roots of polynomials, and polynomial roots can have multiplicities - the same root can “repeat” multiple times.

For simplicity, assume for now that at our point of interest we have a simple eigenvalue, namely, with multiplicity 1. In this case, a well-known result from linear algebra is that it has a unique (up to sign) normalized eigenvector \({\boldsymbol q}_k({\boldsymbol X})\). Turns out that the function \(f\) is differentiable at such points, and the gradient is simple:

Thus, the only thing we need for back-propagation is the eigenvector corresponding to our desired eigenvalue - its outer product with itself is the gradient.

Let’s convince ourselves that this works with code. Here is the outer product of the eigenvector corresponding to the middle eigenvalue of a \(5 \times 5\) matrix with itself:

mat = torch.linspace(-5, 5, 25).reshape(5, 5)
w, Q = torch.linalg.eigh(mat)
i = 2
grad_mat = Q[:, i].view(-1, 1) @ Q[:, i].view(1, -1)
grad_mat

tensor([[ 0.1389, -0.2269, -0.0168,  0.2359, -0.1102],
        [-0.2269,  0.3708,  0.0274, -0.3855,  0.1801],
        [-0.0168,  0.0274,  0.0020, -0.0285,  0.0133],
        [ 0.2359, -0.3855, -0.0285,  0.4008, -0.1873],
        [-0.1102,  0.1801,  0.0133, -0.1873,  0.0875]])

And here is the gradient computed by taking the mid eigenvalue and applying tensor.backward() to compute the gradient:

mat_param = torch.nn.Parameter(torch.linspace(-5, 5, 25).reshape(5, 5))
w = torch.linalg.eigvalsh(mat_param)
w[i].backward()
mat_param.grad

tensor([[ 0.1389, -0.2269, -0.0168,  0.2359, -0.1102],
        [-0.2269,  0.3708,  0.0274, -0.3855,  0.1801],
        [-0.0168,  0.0274,  0.0020, -0.0285,  0.0133],
        [ 0.2359, -0.3855, -0.0285,  0.4008, -0.1873],
        [-0.1102,  0.1801,  0.0133, -0.1873,  0.0875]])

Things become more complicated when the eigenvalue is not simple, and has a multiplicity of at least two. In this case the function is not differentiable, and this is exactly the cause of the “kinks” we saw in the first post in the series, where we aimed to understand what kind of functions are representable using our “eigenvalue neuron”.

There are many notions of “generalized derivatives”, and we will have to choose one that is appropriate. Now here is a spoiler alert - we can still take one of the eigenvectors, call it \({\boldsymbol q}_k\), and use the vector \({\boldsymbol q}_k {\boldsymbol q}_k^T\) for back-propagation. So now that we know what code to write, let’s try to understand why.

Consider the well-known ReLU function with a kink at zero. To the left of zero, the derivative is zero. To the right of zero, it is one. At zero there is no derivative, but we can use any number between zero and one. Intuitively, we understand it’s because any line with a slope between zero and one can behave like a tangent - it touches the function at one point. Now note one important point - I said any number between zero and one. So we don’t have one slope we can use - we have an infinity of them.

A generalization of this idea of using the set of vectors “in between neighboring gradients” is known as the Clarke sub-differential¹. In higher dimensions, “in-between” generalizes to the closure of the convex hull. I am not going deep into theory, so we’ll not discuss exactly the convex hull of what we are taking, but intuitively these are gradients in a small neighborhood. If you’re interested, I have a great book² by Frank Clarke himself to recommend :)

Clarke sub-differential is one of these notions of generalized derivatives that are typically accepted as the “right” one for back-prop ³⁴. We are not always guaranteed to get an element in the Clarke sub-differential⁴ when backpropagating through a large graph, but we should do our best at least for our atomic building blocks. And just like we can take any slope between 0 and 1 for ReLU, we can take any vector in sub-differential set. Turns out our outer product of an eigenvector with itself is an element of the Clarke sub-differential set for the \(k\)-th eigenvalue function.

Now we have our two ingredients - a way to quickly compute eigenvalues and eigenvectors on the GPU, and a way to compute the gradient for backpropagation - so let’s finally create our PyTorch function!

A custom \(k\)-th eigenvalue function

First, we’ll need two utilities to convert tensors from PyTorch to CuPy and back via DLPack:

def _torch_to_cupy(x: torch.Tensor):
    """ Zero-copy via DLPack for CUDA """
    return cp.from_dlpack(torch_dlpack.to_dlpack(x))

def _cupy_to_torch(x_cupy):
    """ Zero-copy via DLPack for CUDA """
    return torch_dlpack.from_dlpack(x_cupy)

Implementing a custom PyTorch autograd function is quite simple - we just need to follow a template. We inherit from torch.autograd.Function and implement two static methods - forward and backward. The former computes our function, and optionally caches anything required for computing the derivative. The latter just back-propagates the derivative. Moreover, to make things efficient, typically forward is split into two code paths - one efficient path when no derivatives are required (inference mode), and another one for the case when derivatives are required. So here it is:

class CuPyKthEigval(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A: torch.Tensor, k: int, lower: bool = True):
        # A: PyTorch --> CuPy
        A_ = A if A.is_contiguous() else A.contiguous()
        A_cp = _torch_to_cupy(A_.detach())

        # Which part of A to use, in CuPy language
        uplo = "L" if lower else "U"

        if ctx.needs_input_grad[0]: # for training
            # CuPy eigenvalues and eigenvectors
            ws_cp, Qs_cp = cp.linalg.eigh(A_cp, UPLO=uplo)
            
            # CuPy --> PyTorch
            ws = _cupy_to_torch(ws_cp).to(dtype=A.dtype, device=A.device)
            Qs = _cupy_to_torch(Qs_cp).to(dtype=A.dtype, device=A.device)
            
            # Store k-th eigenvector for the derivative
            ctx.save_for_backward(Qs[..., k].unsqueeze(-1))
        else: # for inference
            ws_cp = cp.linalg.eigvalsh(A_cp, UPLO=uplo)
            ws = _cupy_to_torch(ws_cp).to(dtype=A.dtype, device=A.device)

        return ws[..., k] # k-th eigenvalue

    @staticmethod
    def backward(ctx, grad_w: torch.Tensor):
        (Q,) = ctx.saved_tensors  # (..., n, 1)
        grad_w = grad_w.to(dtype=Q.dtype)
        grad_A = (Q * grad_w[..., None, None]) @ Q.transpose(-1, -2)
        return grad_A, None, None  # no grad for `k` and `lower`

It’s a bit lengthy, but straightforward. We just follow the sketch we laid about above. To use our new function, we just need to call the apply function of our new class:

mat_param = torch.nn.Parameter(torch.linspace(-5, 5, 25, device='cuda').reshape(5, 5))
w = CuPyKthEigval.apply(mat_param, 2)
w.backward()
mat_param.grad

tensor([[ 0.1389, -0.2269, -0.0168,  0.2359, -0.1102],
        [-0.2269,  0.3708,  0.0274, -0.3855,  0.1801],
        [-0.0168,  0.0274,  0.0020, -0.0285,  0.0133],
        [ 0.2359, -0.3855, -0.0285,  0.4008, -0.1873],
        [-0.1102,  0.1801,  0.0133, -0.1873,  0.0875]], device='cuda:0')

Identical to the gradient we previously obtained - so as a sanity check, this appears to be working. Now let’s measure the speed versus PyTorch with a mini-batch of 500 matrices of size \(50 \times 50\):

mats = torch.randn(500, 100, 100, device='cuda')
mat_param = torch.nn.Parameter(torch.randn(500, 100, 100, device='cuda'))

%%time 
w = CuPyKthEigval.apply(mats, 2).sum()

CPU times: user 1.44 ms, sys: 993 µs, total: 2.43 ms
Wall time: 2.1 ms

Alright, 2 milliseconds. Pretty fast. What happens if we need backpropagation?

%%time
w = CuPyKthEigval.apply(mat_param, 2).sum()
w.backward()

CPU times: user 81.3 ms, sys: 1 µs, total: 81.3 ms
Wall time: 81 ms

Much slower! But we also understand why - we need the eigenvector, not just the eigenvalue. What about PyTorch?

CPU times: user 795 ms, sys: 0 ns, total: 795 ms
Wall time: 794 ms

Apparently, our custom function, even with backprop, is almost 10 times faster! Things appear much better, and we can move forward. As a final step, we now wrap it with a convenience function to separate the CUDA tensors from non-CUDA tensors:

def faster_kth_eigvalh(
        A: torch.Tensor, k: int,  *, lower: bool = True
    ) -> torch.Tensor:
    if A.is_cuda:
        return CuPyKthEigval.apply(A, k, lower)
    else:
        return torch.linalg.eigvalsh(A, lower)[..., k]

Nice! So now we have a function that works quickly on a GPU and we can finally do an experiment that I was not able to do in the previous post within a reasonable amount of time - try even larger matrices!

Trying it out in practice

Recall that in the last post we implemented a class, called MultivariateSpectral, for the model family we study in this series:

where the non-decreasing vector \({\boldsymbol \mu}\) and the symmetric matrices \(\mathbf{A}_1, \dots, \mathbf{A}_n\) are the learned parameters. Here is a version of it that uses our new faster_kth_eigvalh function:

class MultivariateSpectral(nn.Module):
    def __init__(self, *, num_features: int, dim: int, eigval_idx: int):
        super().__init__()
        self.eigval_idx = eigval_idx
        self.mu = Nondecreasing(dim) # <-- we wrote it in the last post
        self.A = nn.Parameter(
            torch.randn(num_features, dim, dim) / (math.sqrt(dim) * num_features)
        )

    def forward(self, x):
        # batches of sum of x[i] * A[i]
        nf, dim = self.A.shape[:2]
        feature_mat = (x @ self.A.view(nf, dim * dim)).view(-1, dim, dim)

        # diag(mu) replicated per batch
        bias_mat = torch.diagflat(self.mu()).expand_as(feature_mat)

        # batched eigenvalue computation
        return self._compute_eigval(bias_mat + feature_mat)

    def _compute_eigval(self, mat):
        return faster_kth_eigvalh(mat, self.eigval_idx)

To be able to test it against PyTorch, here is a variant that uses the regular PyTorch eigenvalue function - we just inherit the above class and override the _compute_eigval function:

class MultivariateSpectralTorch(MultivariateSpectral):
    def _compute_eigval(self, mat):
        return torch.linalg.eigvalsh(mat)[..., self.eigval_idx]

So now we will use functions we implemented in the last post to again test ourselves on supervised regression with the California Housing dataset.

In the last post we implemented the function train_model_stream that trains the given model and yields a sequence of dictionaries containing the model and the training loss, and the add_spectral_norms which augments this dictionary with spectral norms of the learned matrices that we used for obtaining a global bound on the model’s sensitivity with respect to features. Here we shall just use these helpers assuming they are defined, and that the dataset is already loaded and pre-processed. The linked notebook at the beginning of this post contains the full code.

So let’s measure how long 5 training epochs take with PyTorch eigenvalues. Again, we shall use the %%time Jupyter magic keyword to measure time:

def training_stream(model, n_epochs, **train_kwargs):
    criterion = nn.MSELoss()
    return add_spectral_norms(train_model_stream(
        model, criterion, n_epochs=n_epochs, **train_kwargs
    ))

%%time
model = MultivariateSpectralTorch(num_features=num_features, dim=45, eigval_idx=22)
for event in training_stream(model, n_epochs=5):
    print('tick')

tick
tick
tick
tick
tick
CPU times: user 1min 16s, sys: 466 ms, total: 1min 16s
Wall time: 1min 16s

Now let’s try it with our faster eigenvalue function:

%%time
model = MultivariateSpectral(num_features=num_features, dim=45, eigval_idx=22)
for event in training_stream(model, n_epochs=5):
    print('tick')

tick
tick
tick
tick
tick
CPU times: user 15.7 s, sys: 205 ms, total: 15.9 s
Wall time: 16.1 s

Nice! Almost five times faster! So now I can actually conduct an experiment I could not in the previous post - see how the model scales if I increase matrix size even further, to \(45 \times 45\). To that end, we shall re-use the plot_progress function that consumes such an iterable stream produced by training and produces a live-updating plot of the progress. Again - I assume the function is given, but you have the full code in the notebook.

def live_plot_training(dim, n_epochs, **train_kwargs):
    model = MultivariateSpectral(
        num_features=num_features, dim=dim, eigval_idx=dim // 2
    )
    events = training_stream(model, n_epochs, **train_kwargs)
    plot_progress(
        events, max_step=n_epochs
    )

%%time
live_plot_training(dim=45, n_epochs=500, lr=5e-5)

CPU times: user 30min 45s, sys: 23.7 s, total: 31min 8s
Wall time: 31min 23s

Well, it took me half an hour. Quite long. But I was able to produce this plot:

Recall that for a \(30 \times 30\) matrix, we got a test error of \(\approx \$54200\), so scaling up indeed improves performance somewhat, but not dramatically. Apparently, with our current training procedure we begin to notice the diminishing returns of this type of scaling.

Now, this does not mean that our training procedure is the best, and this is definitely not an exhaustive scaling experiment, where we choose the best training procedure we can, and perhaps devise some rule of hyperparameter transfer from smaller to larger models. But having the ability to compute eigenvalues quickly lets us actually conduct this research, since PyTorch eigenvalue solver was simply too slow.

Recap

Now that we have the ability to conduct fast experiments we can move forward and do other interesting stuff. Obviously, there might be even better ways to achieve our goal - perhaps writing a custom CUDA kernel for the entire function \(f(\mathbf{x}; {\boldsymbol \mu}, \mathbf{A}_{1:n})\) would even be better. But I just wanted something that doesn’t get in my way when I’m experimenting - that’s all.

The next post in the series will be very different - it will be theoretical. We did a lot of practical things here, but we need to understand some things before we move forward. So stay tuned!

References

Intro

We all want our models to perform well. But some of us would also like our models to be efficient, robust, or interpretable. So in this post we will discuss some mathematical properties of these models that are related to these three pillars. Robustness and interpretability may mean different things to different people, so let’s explain what I mean in this post. As a general note - many things I am going to talk about are true for complex Hermitian matrices, but we focus on real symmetric matrices in the post. So this is the first and the last time I mention complex numbers in this series.

The robustness that we shall explore means robustness to corruption or noise, meaning that bounded changes to the input yield bounded changes to the output, and this bound is known. This is important when we want to know that a small perturbation will not make our model “go wild” and predict something totally unreasonable.

Interpretability can also mean many things. It can be interpretability for us, scientists, so that we can explain what the model does to ourselves. Alternatively, it can mean that we can explain what the model does to a business stakeholder or a regulator. Or in the extreme case, it means we can actually explain to a user why our system made the decision it made based on their input, i.e., why am I not getting a better insurance premium? In this post we shall mostly talk about the first two aspects.

But let’s get started with a small debt I believe I owe you from the previous post - eliminating some of the redundancy.

Eliminating redundancy

We defined our models as

Some of us may remember from linear algebra that eigenvalues of symmetric matrices are invariant under orthogonal transformations. So the representation of our model is not unique - we can just replace all matrices \(\mathbf{A}_i\) by \(\mathbf{Q}\mathbf{A}_i\mathbf{Q}^\intercal\) for some orthogonal matrix \(\mathbf{Q}\) and obtain exactly the same model. Redundancy, of course, is not unique to this family. Matrix factorization models¹ have a similar redundancy. But we can eliminate some of this redundancy.

Since \(\mathbf{A}_0\) is symmetric, it has a spectral decomposition:

where \(\boldsymbol \mu\) is the vector of eigenvalues in some predefined order, such as non-increasing or non-decreasing. Consequently, the model can be written as

Thus, we can assume that the matrix \(\mathbf{A}_0\) is, for example, diagonal and non-decreasing, without losing any representation power, and assume our model is always of the form:

where \(\boldsymbol \mu\) is a non-decreasing vector, and \(\mathbf{A}_i\) are symmetric matrices. So let’s implement such a model in PyTorch. To that end, we will need a way to represent a non-decreasing vector, which is quite easy - use torch.nn.softplus to generate non-negative gaps, and sum them up. Also, I don’t know what is the right initialization for our \(\boldsymbol \mu\), so I chose uniformly spaced points between -1 and 1:

import torch
from torch import nn

class Nondecreasing(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        init = torch.linspace(-1, 1, dim)
        self.start = nn.Parameter(init[:1])
        self.increments = nn.Parameter(init.diff().expm1().log())

    def forward(self):
        return torch.cat([
            self.start,
            self.start + nn.functional.softplus(self.increments).cumsum(dim=0)
        ])

Let’s try it out:

Nondecreasing(10)()

tensor([-1.0000, -0.7778, -0.5556, -0.3333, -0.1111,  0.1111,  0.3333,  0.5556,
         0.7778,  1.0000], grad_fn=<CatBackward0>)

Appears to be working. Now, this may not be the best way to parameterize a non-decreasing vector, and you probably can think of other ways, but it appears to works reasonably well when we train models later in this post.

So now we can use it to implement a PyTorch module for the kind of functions we seek. The code is mostly straightforward, and the only thing requires explaining is the initialization of the matrices \(\mathbf{A}_i\), that we shall talk about right after the code snippet:

import torch.linalg as tla

class MultivariateSpectral(nn.Module):
    def __init__(self, *, num_features: int, dim: int, eigval_idx: int):
        super().__init__()
        self.eigval_idx = eigval_idx
        self.mu = Nondecreasing(dim)
        self.A = nn.Parameter(
            torch.randn(num_features, dim, dim) / (math.sqrt(dim) * num_features)
        )

    def forward(self, x):
        # batches of sum of x[i] * A[i]
        nf, dim = self.A.shape[:2]
        feature_mat = (x @ self.A.view(nf, dim * dim)).view(-1, dim, dim)

        # diag(mu) replicated per batch
        bias_mat = torch.diagflat(self.mu()).expand_as(feature_mat)

        # batched eigenvalue computation
        eigvals = tla.eigvalsh(bias_mat + feature_mat)
        return eigvals[..., self.eigval_idx]

Regarding initialization, I am making an educated guess here. It is known² that the spectrum of \(n \times n\) matrices with random Gaussian entries converges to the semicircle distribution in \([-2\sqrt{n}, 2 \sqrt{n}]\) as \(n\) grows. Moreover, since we will be summing up num_features matrices, it makes sense to initialise our matrices to a normal distribution with a standard deviation of \((\sqrt{n} \cdot \mathtt{num\_features})^{-1}\). Here, too, I don’t know if this is the best initialization, but it works reasonably well.

As a sanity test, let’s try learning the concave function from the previous post:

def f(x, y):
    return -torch.log(torch.exp(x-1) + torch.exp(y+0.5) + torch.exp(-x-y+0.5))
 
# sample 10000 points on the graph of the function
x = torch.empty(10000).uniform_(-3, 3)
y = torch.empty(10000).uniform_(-3, 3)
xy = torch.stack([x, y], dim=-1)
z = f(x, y) + 0.2 * torch.randn(10000)

Here is a simple training loop to see if the loss decreases - let’s fit a concave model (smallest eigenvalue) with \(5 \times 5\) matrices

import math
from itertools import count

model = MultivariateSpectral(num_features=2, dim=5, eigval_idx=0)
optim = torch.optim.Adam(model.parameters(), lr=1e-3)
batch_size = 10
print_every = 100

cum_loss = 0.
for i, xyb, zb in zip(count(), xy.split(batch_size), z.split(batch_size)):
    loss = (model(xyb) - zb).square().mean()
    cum_loss += loss.detach().item()

    optim.zero_grad()
    loss.backward()
    optim.step()
        
    if (i + 1) % print_every == 0:
        print(f'Loss = {cum_loss / print_every:.4f}')
        cum_loss = 0.

Loss = 0.4582
Loss = 0.1355
Loss = 0.0774
Loss = 0.0649
Loss = 0.0626
Loss = 0.0520
Loss = 0.0534
Loss = 0.0516
Loss = 0.0464
Loss = 0.0468

OK. The model appears to be learning - the loss is decreasing. So now that we have eliminated most of the redundancy, let’s move on to more interesting stuff.

Spectral stability and its consequences

First, let us recall that any matrix has an associated operator norm - the maximum amount by which it can stretch a unit vector:

We have np.linalg.norm and torch.linalg.norm to reliably compute it. Why are we recalling it? Turns out there is a useful consequence of the Weyl’s inequality for symmetric matrices - spectral stability:

So if we take a symmetric matrix \(\mathbf{A}\) and “corrupt” or “perturb” it by another symmetric matrix \(\mathbf{B}\), the resulting eigenvalues do not change by more than \(\|\mathbf{B}\|_{\mathrm{op}}\).

Now, consider our model family, and suppose that the first feature \(x_1\) was perturbed by some noise \(\varepsilon\). By the spectral stability property, our model’s output will not change by more than \(\lvert\varepsilon\rvert \| \mathbf{A}_1 \|_{\mathrm{op}}\). And in general, if our feature vector was perturbed by some noise \(\boldsymbol \varepsilon\), we have:

Now, we have two ways to interpret this bound. First, from the standpoint of robustness - we have a direct bound on the possible change of the prediction as a function of the noise \(\boldsymbol \varepsilon\). For example, if we care about the \(\ell_2\) norm of the noise and want to know what happens when \(\|\boldsymbol \varepsilon\|_2 \leq \alpha\), the Cauchy-Schwarz inequality implies that the model’s prediction changes by at most \(\alpha \sqrt{\sum_{i=1}^n \| \mathbf{A}_i \|^2_{\mathrm{op}}}\).

The second way to think of the bound is from the standpoint of interpretability: one notion of feature importance is a worst-case sensitivity bound. The quantity \(\| \mathbf{A}_i \|_{\mathrm{op}}\) upper-bounds how much the prediction can change when only feature \(x_i\) is perturbed, because a small change of \(\varepsilon\) to feature \(x_i\) will make the model’s prediction change by at most \(\varepsilon \| \mathbf{A}_i \|_{\mathrm{op}}\). So this operator norm is a bound on the effect of feature \(x_i\) on the model’s prediction, just like the magnitude of the coefficients in a linear model.

We can use this knowledge in two ways. First, having trained a model, we can interrogate it for its robustness / feature-importance properties by computing the spectral norms of all feature matrices. Second, we can try to impose a regularization term that imposes a limit on these operator norms. So let’s try the first idea - of observing the operator norms.

Observing stability bounds in practice

We will do it with our beloved California Housing data-set that I use a lot in my blog posts, simply because it’s there on Colab. So let’s load it:

import pandas as pd

train_df = pd.read_csv('sample_data/california_housing_train.csv')
test_df = pd.read_csv('sample_data/california_housing_test.csv')

You may recall from our previous blog posts, that the dataset has four very skewed columns that we typically apply a log transformation to:

import numpy as np

skewed_columns = ['total_rooms', 'total_bedrooms', 'population', 'households']
train_df[skewed_columns] = train_df[skewed_columns].apply(np.log)
test_df[skewed_columns] = test_df[skewed_columns].apply(np.log)

Our final data preprocessing step is plain scaling:

from sklearn.preprocessing import StandardScaler

scaler = StandardScaler().set_output(transform='pandas')
train_scaled = scaler.fit_transform(train_df)
test_scaled = scaler.transform(test_df)

label_scale = scaler.scale_[-1]

We remember the scale of the last column, the label of the data-set, because we want our evaluation metrics in the original units of the label, not in the normalized units. Before training, let’s put our training data in PyTorch tensors - it will be more convenient:

from torch import as_tensor

def to_tensors(df):
    target = 'median_house_value'
    return (
        as_tensor(df.drop(target, axis=1).values), 
        as_tensor(df[target].values)
    )

X_train, y_train = to_tensors(train_scaled)
X_test, y_test = to_tensors(test_scaled)

num_features = X_train.shape[1]
n_train = len(X_train)

Alright! So now let’s write our training loop. Here is a fairly standard PyTorch loop for one epoch:

def train_epoch(
        device, net, optimizer, criterion, regularizer, X_batches, y_batches
    ):
    epoch_loss = torch.zeros(1).to(device)
    for x, y in zip(X_batches, y_batches):
        optimizer.zero_grad()
        loss = criterion(net(x), y)
        cost = loss + regularizer(net)
        cost.backward()

        with torch.no_grad():
            epoch_loss += loss * x.shape[0]
        optimizer.step()
    return (epoch_loss / n_train).cpu().item()

The regularizer will become useful later in this post - it’s just an additional penalty beyond the loss. And here is our pretty-standard training loop for the model, but with a twist: we yield intermediate results. Why? It’s convenient to work with - we fully decouple training code from reporting / plotting code:

def train_model_stream(
        net, criterion, *, n_epochs=200, batch_size=10, lr=1e-4, regularizer=None
    ):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    regularizer = regularizer or (lambda model: 0.) # by default - no reg.

    net.to(device)
    X_train_batches = X_train.to(device).split(batch_size)
    y_train_batches = y_train.to(device).split(batch_size)
    X_test_device = X_test.to(device)
    y_test_device = y_test.to(device)

    optimizer = torch.optim.AdamW(net.parameters(), lr=lr)
    for epoch in range(1, 1 + n_epochs):
        train_loss = train_epoch(
            device, net, optimizer, criterion, regularizer,
            X_train_batches, y_train_batches
        )

        with torch.no_grad():
            test_loss = criterion(net(X_test_device), y_test_device)
            test_loss = test_loss.cpu().item()

        yield {
            'step': epoch,
            'model': net,
            'train_error': math.sqrt(train_loss) * label_scale,
            'test_error': math.sqrt(test_loss) * label_scale,
        }

This is where we use the label_scale we previously stored - to report the error in the units of the original labels, not the normalized ones. Let’s try a few epochs, to see how it works:

model = MultivariateSpectral(num_features=num_features, dim=5, eigval_idx=2)
criterion = nn.MSELoss()
for event in train_model_stream(model, criterion, n_epochs=5):
    print(event['step'], event['train_error'], event['test_error'], sep='\t')

1	109072.75280327671	99536.88016316161
2	95295.4091119307	89540.08167748511
3	86628.42283455568	82857.71197495822
4	80906.37794419663	78535.01666073261
5	77319.691256485	75821.19027087984

OK - model appears to be training nicely. This trick of yielding lets us do interesting stuff - for example, we can create a new stream that yields train and test errors, together with the spectral norms of the feature matrices:

def add_spectral_norms(stream):
    for event in stream:
        model = event['model']
        with torch.no_grad():
            # remember - we're using only lower-triangular part of each A_i
            matrices_sym = \
                model.A.tril() + model.A.tril(diagonal=-1).transpose(-1, -2)
            norms = tla.matrix_norm(matrices_sym, ord=2)
            norms = norms.ravel().cpu().tolist()
        
        yield {
            'step': event['step'],
            'train_error': event['train_error'],
            'test_error': event['test_error'],
        } | {
            f'norm_{feature_name}': norm 
            for feature_name, norm in zip(feature_names, norms)
        }

Let’s try it out. This time we’ll use the rich library for pretty printing, since the regular Python print doesn’t produce a nice output. So here are 2 training epochs:

from rich.pretty import pprint

model = MultivariateSpectral(num_features=num_features, dim=5, eigval_idx=2)
criterion = nn.MSELoss()
for event in add_spectral_norms(train_model_stream(model, criterion, n_epochs=2)):
    pprint(event)

  {
    'step': 1,
    'train_error': 105523.77536281559,
    'test_error': 97731.75678823328,
    'norm_longitude': 0.12640513479709625,
    'norm_latitude': 0.1759399175643921,
    'norm_housing_median_age': 0.15713410079479218,
    'norm_total_rooms': 0.17977216839790344,
    'norm_total_bedrooms': 0.16544003784656525,
    'norm_population': 0.19817979633808136,
    'norm_households': 0.2670281231403351,
    'norm_median_income': 0.33458226919174194
}
{
    'step': 2,
    'train_error': 94575.6796415192,
    'test_error': 88766.7939206047,
    'norm_longitude': 0.15076042711734772,
    'norm_latitude': 0.1756560057401657,
    'norm_housing_median_age': 0.1784549206495285,
    'norm_total_rooms': 0.18125228583812714,
    'norm_total_bedrooms': 0.15172506868839264,
    'norm_population': 0.19934602081775665,
    'norm_households': 0.26696428656578064,
    'norm_median_income': 0.47850197553634644
}

Nice! So now we can iterate and do live-plotting of everything! This is a lengthy function with mostly boilerplate that plots two graphs - one with train/test errors, and another one with spectral norms of feature matrices. I added comments to make the code clear, but the principle is simple: we create empty plots, and gradually update them as new events arrive.

import matplotlib.pyplot as plt

def plot_progress(events, max_step):
    # create a plot with two axes - one for errors, one for norms
    fig, (err_ax, norm_ax) = plt.subplots(
        2, 1, figsize=(8, 8), layout='constrained'
    )

    # create empty line objects
    def plot_empty(ax, label):
        return ax.plot([], [], label=label)[0]

    line_dict = {
        'train_error': plot_empty(err_ax, 'train error'),
        'test_error': plot_empty(err_ax, 'test error'),
    } | {
        f'norm_{feature_name}': plot_empty(norm_ax, feature_name)
        for feature_name in feature_names
    }

    # setup axis properties
    err_ax.set_title("Error")
    norm_ax.set_title("Matrix norms")
    for ax in (err_ax, norm_ax):
        ax.set_xlabel("Step")
        ax.set_xlim(0, max_step)
        ax.grid(True)
        ax.legend()


    # display figure and obtain its handle
    h = display(fig, display_id=True)
    plt.close(fig)

    # iterate over events and update the plot
    min_test_error = float('inf')
    for event in events:
        step = event['step']
        min_test_error = min(min_test_error, event['test_error'])
        err_ax.set_title(f'Error (min test err = {min_test_error:.2f})')

        for key, line in line_dict.items():
            value = event[key]
            x, y = line.get_data(orig=True)
            line.set_data(np.append(x, step), np.append(y, value))

        for axs in (err_ax, norm_ax):
            axs.relim()
            axs.autoscale_view()

        fig.canvas.draw()
        h.update(fig)

Alright! Let’s use it to train a mid-eigenvalue model with \(5 \times 5\) matrices:

def live_plot_training(dim, n_epochs):
    model = MultivariateSpectral(
        num_features=num_features, dim=dim, eigval_idx=dim // 2
    )
    criterion = nn.MSELoss()
    events = add_spectral_norms(train_model_stream(
        model, criterion, n_epochs=n_epochs
    ))
    plot_progress(events, max_step=n_epochs)

live_plot_training(5, 500)

OK. We can see that the model is learning, and after 500 epochs we observe that the resulting model’s strongest three features are longitude, latitude, and population. What happens when we increase model size? Let’s try \(15 \times 15\) matrices:

live_plot_training(15, 500)

We see that the test loss decreases with the model size, and even though the ranking between features is slightly different, the three strongest features remain longitude, latitude, and population. But we also see something else - the matrix norms continue growing. Apparently, after 500 epochs, the model’s parameters do not appear to be converging. Perhaps a more thorough hyper-parameter tuning would help, I don’t know. But I chose a conservative option of a small learning rate and many epochs for a reason - to show that scaling model size improves performance, while keeping our model’s ability to be interpretable almost as if it was linear.

Let’s go even further up, to \(30 \times 30\) matrices:

live_plot_training(30, 500)

We see that the train and test errors go further down, and the three features previously at the top remain there. Again - scaling up improves performance, while keeping interpretability and computable robustness bounds.

So what we got here is really interesting! We have a model that is nonlinear and improves with scaling, while remaining interpretable in terms of feature sensitivity / importance, and we have an easy way to compute global sensitivity bounds (which can be loose).

As a reference, if you try fitting a gradient-boosted decision forest using XGBoost, you’ll observe a test error of approximately $48,000. So the eigenvalue model we see here isn’t close to what trees can achieve, but tree ensembles are often discontinuous and don’t come with simple global sensitivity/Lipschitz certificates in the same way. So it’s a tradeoff.

Sensitivity control

Another way we can use our understanding of the stability properties is to regularize the model by either imposing a bound on the maximum spectral norm, or adding a regularization term that penalizes the spectral norms, so our training code will be minimizing

This is where we shall use the regularizer parameter of our training function that I promised you:

def live_plot_reg_training(dim, n_epochs, reg_coef):
    model = MultivariateSpectral(
        num_features=num_features, dim=dim, eigval_idx=dim // 2
    )

    def penalty(net):
        matrices_sym = \
            net.A.tril() + net.A.tril(diagonal=-1).transpose(-1, -2)
        norms = tla.matrix_norm(matrices_sym, ord=2)
        return reg_coef * norms.sum()

    criterion = nn.MSELoss()
    events = add_spectral_norms(train_model_stream(
        model, criterion, n_epochs=n_epochs, regularizer=penalty
    ))
    plot_progress(events, max_step=n_epochs)

Let’s try it out with \(15 \times 15\) matrices:

live_plot_reg_training(15, 500, 1e-3)

We can see that the spectral norms are smaller than our previous attempt with \(15 \times 15\) matrices above, norm growth appears to stabilize, but performance appears similar. Just the gap between the top four features and the rest of the features became more pronounced - that’s the effect of delicate regularization. A larger regularization coefficient may even drive some of the matrices towards zero, similarly to \(\ell_1\) regularization in Lasso.

Imposing such a regularizer with standard PyTorch optimizers, rather than a dedicated optimizer, may not be the optimal (pun intended!) thing to do, and some of you can probably think of better ways. But that’s beside the point - the point is that we can, in principle, regularize the spectral norm to control the model’s sensitivity to feature perturbations. And that is quite powerful.

So now after we’ve seen plenty of stuff - it’s time for a recap.

Summary

We saw that matrix eigenvalues let us find a nice sweet-spot between several opposing forces - performance, robustness, and interpretability. Beyond just models for tabular data, this nice idea can also be employed for another use case we haven’t yet discussed - ensembling. There, too, we care about the ensemble’s prediction to behave “sensibly” w.r.t the predictions of the individual models, and there too we may care about robustness and interpretability. So it’s nice to have a learnable ensembling technique that both improves with scaling, but remains robust and somewhat interpretable.

We will study other mathematical properties in future posts that will let us understand on a deeper level what kind of information we can elicit from those models, but as of now we have a slightly more urgent concern: training is slow. We need many epochs, and each epoch is expensive. This makes experimentation hard - our feedback loop is slow as well. So this is something we shall try to address in the next post!

References

Intro

When trying to model a complicated relationship between features, our go-to architectures are typically either neural networks or decision trees. They are well-established, well-studied, and have an abundance of software for training them. So why not?

But sometimes we have some additional requirements. Maybe we want our model to be an increasing / decreasing / convex / concave function of one or more feature. Perhaps we want to measure diminishing returns, and we want it to be both increasing and concave. Or maybe we care about the sensitivity of our model to noise, and want to certify its Lipschitz constant: what happens to the prediction if we slightly change the input?

In the recent interview with Ilya Sutskever there was one interesting insight - perhaps we need neurons to do more compute than they do now. And it immediately rang a bell, and returned me to my Ph.D days - what would be this “more compute” that is useful? Well, maybe a neuron can solve a small optimization problem! So in this post we shall explore an idea that most optimization researchers are familiar with - the eigenvalues of a matrix are solutions to optimization problems. So what can one such neuron do? Turns out quite a lot! This is exactly what we shall explore in this post.

Why eigenvalues? Well, because of a unique combination of three properties. They can model fairly complicated functions, we have a lot of theoretical machinery to reason about them, and we can stand on the shoulders of giants and reuse the vast talent and resources invested over decades in their reliable computation.

The notebook for reproducing all the results is here. Feel free to deploy it to Colab and play with it.

Univariate functions

Let’s begin our adventures with a simple case - function of one variable. Suppose we’re given symmetric matrices, \(\mathbf{A}\) and \(\mathbf{B}\), and we define the function:

where \(\lambda_k(\cdot)\) is the \(k\)-th smallest eigenvalue of the given matrix. In general, we will be interested in two things - what kind of functions can we represent, and whether we learn such functions from data.

Our exploration begins with plotting, so let’s implement such a function in a vectorized manner - it will take a vector of inputs, and produce a corresponding vector of outputs. Turns out SciPy has excellent tools just for that:

import scipy.linalg as sla
import numpy as np

def univariate_spectral(A, B, k, xs):
  """Computes the vector y[i] = λₖ(A + B * xs[i])."""
  
  # support negative eigenvalue indices,
  #  e.g., k=-1 is the largest eigenvalue
  k = k % A.shape[0]
  
  # create a batch of matrices, one for each entry in xs
  mats = A + B * xs[..., np.newaxis, np.newaxis]
  
  # compute the k-th eigenvalue of each matrix
  return sla.eigvalsh(mats, subset_by_index=(k, k)).squeeze()

Note - we don’t have to pass symmetric matrices. eigvalsh treats its input as symmetric/Hermitian and reads only one triangle (the lower triangle by default). So what do the functions look like? Let’s take two \(3 \times 3\) matrices \(\mathbf{A}\) and \(\mathbf{B}\) and plot the three functions corresponding to the three eigenvalues:

import matplotlib.pyplot as plt

def plot_eigenfunctions(A, B, n_rows, n_cols):
  """Plots λₖ(A + B * x) on a grid layout"""
  fig, axs = plt.subplots(n_rows, n_cols, figsize=(3 * n_cols, 2 * n_rows), layout='constrained')
  plot_xs = np.linspace(-5, 5, 1000)
  for k, ax in enumerate(axs.ravel()):
    ax.plot(plot_xs, univariate_spectral(A, B, k, plot_xs))
    ax.set_title(f'$\\lambda_{k}(A + B * x)$')
  plt.show()  


np.random.seed(42)
A = np.random.randn(3, 3)
B = np.random.randn(3, 3)
plot_eigenfunctions(A, B, 1, 3)

Interesting. The smallest eigenvalue is a concave function, the largest is convex, and the middle appears arbitrary. What happens if we take two \(9 \times 9\) matrices?

np.random.seed(42)
A = np.random.randn(9, 9)
B = np.random.randn(9, 9)
plot_eigenfunctions(A, B, 3, 3)

Again, smallest eigenvalue is concave, largest is convex, and the eigenvalues in between are typically neither convex nor concave. Coincidence?

Some of you probably know it is not a coincidence at all. Indeed, recalling elementary results from linear algebra, the smallest eigenvalue of a matrix can be alternatively written as a minimum of quadratic functions:

These functions may be quadratic in \(\mathbf{x}\), but they are linear in the matrix \(\mathbf{P}\). And a minimum of linear functions is concave. Similarly, the largest eigenvalue is:

and a maximum of linear functions is convex. So it means our “univariate neuron”, at least for \(\lambda_1\) and \(\lambda_n\), is indeed solving an optimization problem!

But what about the other eigenvalues? There’s a more general version of the above, somewhat less known in the ML community - the Courant-Fischer theorem, for characterizing the \(k\)-th smallest eigenvalue:

This one appears a bit hairy - so lets dissect it. It can be thought of a game between two players, one chooses \(\mathbf{V}\) and the other one chooses a unit vector \(\mathbf{x}\) that is orthogonal to the rows of \(\mathbf{V}\). The “outer” player aims to maximize their reward, and the “inner” one aims to harm the outer as much as possible.

Consequently, \(\lambda_k\) in general is also solution for an optimization problem. And the farther \(k\) is from the extremities of \(k=1\) and \(k=n\), the “crazier” function it models. If we go back to the plots and look, we can even see that the functions are not necessarily smooth. Indeed, they have “kinks”.

To get a more complete picture, let’s see what happens if we make \(\mathbf{B}\) be not just a symmetric matrix, but a positive-semidefinite one, meaning all its eigenvalues are either positive or zero. It’s easy to make one - just create a random matrix and zero-out all its negative eigenvalues:

def make_psd(B):
  eigvals, eigvecs = sla.eigh(B)
  eigvals_pos = np.maximum(0, eigvals)
  return eigvecs @ np.diag(eigvals_pos) @ eigvecs.T

Now let’s plot!

np.random.seed(42)
A = np.random.randn(9, 9)
B = np.random.randn(9, 9)
plot_eigenfunctions(A, make_psd(B), 3, 3)

That’s interesting - all functions are increasing! What if we take a negative-definite \(\mathbf{B}\)?

np.random.seed(42)
A = np.random.randn(9, 9)
B = np.random.randn(9, 9)
plot_eigenfunctions(A, -make_psd(B), 3, 3)

All are decreasing!

Again, this is not a coincidence. Turns out eigenvalue functions are monotone - if we take a matrix \(\mathbf{A}\) and add the matrix \(x \mathbf{B}\) whose eigenvalues are all non-negative, the entire spectrum of eigenvalues increases. So larger \(x\) results in larger eigenvalues, and vice versa, and we obtain an increasing function. The opposite happens when all eigenvalues of \(\mathbf{B}\) are nonpositive.

Beyond univariate functions

To understand what happens beyond the univariate case, let’s look at functions of two variables given three symmetric matrices \(\mathbf{A}, \mathbf{B}, \mathbf{C}\):

For plotting to be convenient, let’s first compute such \(f\) in a vectorized manner:

def bivariate_spectral(A, B, C, k, xs, ys):
  """Computes the vector z[i] = λₖ(A + B * xs[i] + C * ys[i])."""
  
  # support negative eigenvalue indices,
  #  e.g., k=-1 is the largest eigenvalue
  k = k % A.shape[0]
  
  # create a batch of matrices, one for each point (xs[i], ys[i])
  mats = (
      A + B * xs[..., np.newaxis, np.newaxis] 
        + C * ys[..., np.newaxis, np.newaxis]
  )
  
  # compute the k-th eigenvalue of each matrix
  return sla.eigvalsh(mats, subset_by_index=(k, k)).squeeze()

Now we can create surface plots of all eigenvalue functions given the three matrices. Here is the plotting function - nothing special, just a grid of 3D surface plots:

def plot_eigenfunctions_2d(A, B, C, n_rows, n_cols):
  """Plots z = λₖ(A + B * x + C * y) on a grid layout"""
  fig, axs = plt.subplots(
      n_rows, n_cols, figsize=(4 * n_cols, 3 * n_rows),
      subplot_kw={"projection": "3d"}, layout='constrained'
  )
  
  plot_xs = np.linspace(-5, 5, 50)
  grid_xs, grid_ys = np.meshgrid(plot_xs, plot_xs)

  k = A.shape[0]
  for k, ax in enumerate(axs.ravel()[:k]):
    grid_zs = bivariate_spectral(A, B, C, k, grid_xs, grid_ys)
    ax.plot_surface(grid_xs, grid_ys, grid_zs, cmap='viridis')
    ax.set_title(f'$\\lambda_{1 + k}(A + B * x + C * y)$')
  plt.show()  

We’re all set. Let’s plot three eigenvalue functions corresponding to random \(3 \times 3\) matrices:

np.random.seed(42)
A = np.random.randn(3, 3)
B = np.random.randn(3, 3)
C = np.random.randn(3, 3)
plot_eigenfunctions_2d(A, B, C, 1, 3)

As expected, the smallest eigenvalue produces a concave function, the largest produces a convex one, and the mid eigenvalue produces some arbitrary shape. What if we increase the matrix size to \(9 \times 9\)?

np.random.seed(42)
A = np.random.randn(9, 9)
B = np.random.randn(9, 9)
C = np.random.randn(9, 9)
plot_eigenfunctions_2d(A, B, C, 3, 3)

Similar things happen. We can express more “complicated” functions, where the mid eigenvalue \(\lambda_5\) has the “craziest” shape, whereas the smallest eigenvalue produces a concave function and the largest yields a convex one.

I assume you can guess what happens when \(\mathbf{B}\) is negative semi-definite, and \(\mathbf{C}\) is positive semi-definite:

np.random.seed(42)
A = np.random.randn(9, 9)
B = np.random.randn(9, 9)
C = np.random.randn(9, 9)
plot_eigenfunctions_2d(A, -make_psd(B), make_psd(C), 3, 3)

As expected, all functions are decreasing in \(x\), and increasing in \(y\). Moreover, the smallest eigenvalue yields a concave function, whereas the largest yields a convex one.

Now why would it be interesting for ML people? Because we can learn the matrices from data! In fact, we can do even more than that - we can predict them. For example, consider a model for predicting the probability of winning a government contract given the contract document and a bid. We use an encoder to transform the document into a symmetric matrix \(\mathbf{A}\) and a positive semi-definite matrix \(\mathbf{B}\), and just compute \(\lambda_k(\mathbf{A} + \mathrm{bid} \cdot \mathbf{B})\), as illustrated below:

Now we can pass the score to a sigmoid function to obtain a probability that, by design, is nondecreasing in the bid!

You want a more complex model of two numerical inputs \(x\) and \(y\), but it must increase in \(x\) and decrease in \(y\)? No problem! Just encode your other features into an arbitrary \(\mathbf{A}\), a positive semi-definite \(\mathbf{B}\), and a negative semi-definite \(\mathbf{C}\), and compute \(\lambda_k(\mathbf{A} + x \mathbf{B} + y \mathbf{C})\).

Training

But can we actually back-propagate through \(\lambda_k\) to learn our encoder? Well, it turns out we can - everything we need is already in PyTorch. In this post, for the sake of the demonstration, we shall do something extremely simple - use a model whose parameters are just \(\mathbf{A}\), \(\mathbf{B}\) and \(\mathbf{C}\), and learn them from samples of a concave function of two variables.

So here is our simple concave function:

def f(x, y):
    return -np.log(np.exp(x-1) + np.exp(y+0.5) + np.exp(-x-y+0.5))

Let’s take a look at its graph and its level sets:

def plot_bivariate(X, Y, Z, show_rmse=False):
    fig = plt.figure(figsize=(10, 5))
    ax = fig.add_subplot(1, 2, 1, projection='3d')
    ax.plot_surface(X, Y, Z)
    ax = fig.add_subplot(1, 2, 2)
    ax.contour(X, Y, Z, levels=20)
    if show_rmse:
        rmse = np.sqrt(np.mean(np.square(Z - f(X, Y))))
        fig.suptitle(f'RMSE = {rmse:.4f}')

x = np.linspace(-3, 3, 300)
y = np.linspace(-3, 3, 300)
X, Y = np.meshgrid(x, y)
plot_bivariate(X, Y, f(X, Y))

It indeed appears to be a concave function with some non-trivial shape. The show_rmse flag will be used later when we plot a fitted model and would like to show its approximation error.

To train a model we shall need two ingredients: training data and a PyTorch module. So let’s generate some training data:

from torch.utils.data import TensorDataset
from torch import as_tensor

def make_trainset(n_train: int = 200, train_noise: float = 0.2):
    np.random.seed(42)
    x = np.random.uniform(-3, 3, n_train)
    y = np.random.uniform(-3, 3, n_train)
    z = f(x, y) + train_noise * np.random.randn(n_train)
    return x, y, z

x_train, y_train, z_train = make_trainset()
ds = TensorDataset(
    as_tensor(x_train), as_tensor(y_train), as_tensor(z_train)
)

Here is what it looks like:

fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.scatter(x_train, y_train, z_train, color='red')
ax.plot_surface(X, Y, f(X, Y), alpha=0.2)

Assume that we know we are learning some real-world phenomenon that must be concave. Then our PyTorch model should use the smallest eigenvalue function. This is just a PyTorch re-implementation of the above - it is slightly more generic than we need now to support other eigenvalue indices, not just the smallest. Just like with SciPy, PyTorch reads only one triangle of each matrix (the lower triangle by default), so we don’t need to explicitly symmetrize them.

import torch
import torch.linalg as tla
from torch import nn

class BivariateSpectral(nn.Module):
    def __init__(self, *, dim: int, eigval_idx: int):
        super().__init__()
        self.eigval_idx = eigval_idx % dim  # modulo - to support negative idx
        self.A = nn.Parameter(torch.randn(dim, dim))
        self.B = nn.Parameter(torch.randn(dim, dim))
        self.C = nn.Parameter(torch.randn(dim, dim))

    def forward(self, x, y):
        # create a batch of matrices, one for each point (x[i], y[i])
        mats = (
            self.A + self.B * x[..., None, None] + self.C * y[..., None, None]
        )

        # compute the eigenvalues
        eigvals = tla.eigvalsh(mats)
        return eigvals[..., self.eigval_idx]

Alright! Let’s train our smallest eigenvalue model! Below is a pretty standard PyTorch training loop:

from torch.utils.data import DataLoader
import math

def train_model(
        model: nn.Module, n_epochs: int = 500, batch_size = 5, lr=1e-3
    ):
    print_every = n_epochs // 10
    dl = DataLoader(ds, batch_size=batch_size, shuffle=True)

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for epoch in range(1, 1 + n_epochs):
        epoch_loss = 0.
        for x, y, z in dl:
            optimizer.zero_grad()
            loss = loss_fn(model(x, y), z)
            loss.backward()

            with torch.no_grad():
                epoch_loss += loss.item()
            optimizer.step()

        if epoch == n_epochs or epoch % print_every == 0:
            train_rmse = math.sqrt(epoch_loss / len(ds))
            print(f'Epoch {epoch}, train RMSE: {train_rmse: .4f}')

    return model

Alright. Let’s try training a concave model with \(3 \times 3\) matrices, and plot it:

model = train_model(BivariateSpectral(dim=3, eigval_idx=0))
with torch.no_grad():
    Z = model(as_tensor(X), as_tensor(Y)).numpy()
plot_bivariate(X, Y, Z, show_rmse=True)

The output:

Epoch 50, train RMSE:  0.0952
Epoch 100, train RMSE:  0.0910
Epoch 150, train RMSE:  0.0884
Epoch 200, train RMSE:  0.0866
Epoch 250, train RMSE:  0.0859
Epoch 300, train RMSE:  0.0862
Epoch 350, train RMSE:  0.0860
Epoch 400, train RMSE:  0.0858
Epoch 450, train RMSE:  0.0858
Epoch 500, train RMSE:  0.0856

Not bad! Back-propagation appears to be working - the training loss went down, and the plot shows a the model learned something close to the truth.

Now let’s try something interesting - matrices of size \(20 \times 20\). Note, that even one such matrix has more entries than the size of the training set. So such a model is heavily over-parametrized. To make sure we’re converging to the smallest loss I did some tuning, and we need more epochs and a lower learning rate:

model = train_model(
    BivariateSpectral(dim=20, eigval_idx=0), n_epochs=2000, lr=1e-4
)
with torch.no_grad():
    Z = model(as_tensor(X), as_tensor(Y)).numpy()
plot_bivariate(X, Y, Z, show_rmse=True)

Epoch 200, train RMSE:  4.1648
Epoch 400, train RMSE:  2.0389
Epoch 600, train RMSE:  0.5819
Epoch 800, train RMSE:  0.0964
Epoch 1000, train RMSE:  0.0826
Epoch 1200, train RMSE:  0.0811
Epoch 1400, train RMSE:  0.0802
Epoch 1600, train RMSE:  0.0796
Epoch 1800, train RMSE:  0.0793
Epoch 2000, train RMSE:  0.0792

Whoa! This is interesting! Even though the model appears over-parametrized - it does not memorize the training set! Why? Well, one possible explanation is that concavity is a strong inductive bias: with noisy samples, a concave function typically cannot interpolate all points, so the best achievable MSE stays positive. Optimization effects and implicit regularization may also play a role. Also, we see that the model is not “crazy” - its shape still resembles the true one, even if it’s a bit different.

So let’s put our conjecture to a test, and try to fit the mid eigenvalue. This is not a function that is concave by design, and as we saw can have many “crazy” shapes.

model = train_model(
    BivariateSpectral(dim=20, eigval_idx=10), n_epochs=2000, lr=1e-4
)
with torch.no_grad():
    Z = model(as_tensor(X), as_tensor(Y)).numpy()
plot_bivariate(X, Y, Z, show_rmse=True)

Epoch 200, train RMSE:  0.1979
Epoch 400, train RMSE:  0.1031
Epoch 600, train RMSE:  0.0876
Epoch 800, train RMSE:  0.0822
Epoch 1000, train RMSE:  0.0798
Epoch 1200, train RMSE:  0.0786
Epoch 1400, train RMSE:  0.0778
Epoch 1600, train RMSE:  0.0768
Epoch 1800, train RMSE:  0.0764
Epoch 2000, train RMSE:  0.0760

Training error is slightly lower, but apparently does not go down to zero! What happens if we increase the model’s representation power by increasing the size of the matrices? Let’s try the middle eigenvalue of \(40 \times 40\) matrices:

model = train_model(
    BivariateSpectral(dim=40, eigval_idx=20), n_epochs=2000, lr=1e-4
)
with torch.no_grad():
    Z = model(as_tensor(X), as_tensor(Y)).numpy()
plot_bivariate(X, Y, Z, show_rmse=True)

Epoch 200, train RMSE:  0.1032
Epoch 400, train RMSE:  0.0710
Epoch 600, train RMSE:  0.0668
Epoch 800, train RMSE:  0.0650
Epoch 1000, train RMSE:  0.0639
Epoch 1200, train RMSE:  0.0624
Epoch 1400, train RMSE:  0.0617
Epoch 1600, train RMSE:  0.0609
Epoch 1800, train RMSE:  0.0605
Epoch 2000, train RMSE:  0.0598

Training error indeed goes down, so our model gained expressive power. The model has 2460 parameters, much more than our 200 training samples. Yet, the training error does not go to zero, the test error does not explode, and the plot still resembles the true function! So perhaps there is some other property these functions have “by design”, even if we don’t impose shape by using extreme eigenvalues?

Well, there are apparently some interesting properties of these eigenvalue functions worth investigating. So now it would be a good time to do a recap and summarise what we saw, because it appears to be quite a lot to grasp.

Summary

We saw here something that most linear algebra and optimization researchers have known for a long time - eigenvalues of symmetric matrices can be used to model functions of various desired shapes. If we know that the real-world phenomenon we aim to learn has some shape, it might be a good idea to model it that way by design. It makes our model sane at inference, because it behaves as we would expect, and in shape-constrained settings it can also improve training behavior (at least in this toy example).

Some bibliographic notes. The idea of monotonicity in neural networks was studied before. One prominent line of work is that of deep lattice networks¹, and mixing monotonicity with convexity and concavity was studied by Runje and Shankaranarayana². And despite the fact that what I showed here has been known for a long time, the idea of applying this knowledge to construct a learnable model from generic eigenvalue functions is, to the best of my knowledge, quite recent, and appeared in the 2025 paper by Cook et al.³. Their work also proved these models to be universal approximators. Universal approximation is an asymptotic expressivity statement: it does not guarantee interpolation at a fixed size, nor does it say anything about optimization. So it does not contradict our observations here. Moreover, universality is not that important, especially when we care about a specific shape. In fact, in these cases we actually do not want a universal family - we want a family respecting our shape constraints.

Now with this toolbox in mind, we can try to explore other things:

• How do we build a model that learns positive semidefinite matrices in PyTorch to model increasing / decreasing functions?
• Can we mix and match? For example - build a model that is by design increasing in \(x, z\), decreasing in \(y\), and jointly concave in \((x, y)\)?
• Do we need fully dense matrices? Maybe low-rank / sparse / banded matrices already have enough expressive power?
• What exactly are the properties of eigenvalue functions? What is their representation power? How do eigenvalue functions compare to just plain old ReLU MLPs?
• What are the theoretical interpretations of these eigenvalue functions? Perhaps we can think of them as deep neural networks? Or maybe something else?

In the next posts we shall explore some of these aspects, and maybe more. So stay tuned!

References

Yes, I’m making a joke of the tendency to put the words “attention” and “alignment” in any ML paper 😎. Now let’s see how this provocative title is related to our adventures in the land of polynomial features.

The Legendre polynomial basis serverd us well in recent posts about polynomial features. One interesting thing we saw in the series is that its orthogonality is, in some sense informativeness. This is because it orthogonal bases produce features, and hence each basis function, in some sense, carries information that the other basis functions do not. Of course, we all like informative features. So I’d like to devote this post to studying it a bit deeper.

But the Legendre basis is informative in this sense only if our features are uniformly distributed. But real data isn’t uniformly distributed. So in this post I’d like to discuss two ways in which can deal with this practical issue. The associated notebook for reproducing all results is here.

Orthogonality = informativeness

So that we all are on the same page, let’s recall why orthogonal bases produce uncorrelated features. Recall, the two polynomials \(P_i\) and \(P_j\) defined on \([-1, 1]\) are orthogonal if

just like two orthogonal vectors - their inner product is zero. But here the inner product is an integral rather than a sum. But an integral is also an expectation, and empirical averages approximate expectations. So if our data points \(x_1, \dots, x_n\) are approximately uniform in \([-1, 1]\), then

Hence, any column in the data-set the model observes during training is uncorelated to the other columns coming from the same orthogonal basis, and thus in some sense carry information that the other columns do not have.

Informativeness, of course, is not the only important trait of a good basis for non-linear features. In fact, even the norms of the orthogonal basis are important, as well as other traits. I advise reading the well-written and enlightening paper¹ by Muthukumar et. al for more on this. But here, in this post, we focus mainly on orthogonality as informativeness.

Weighted orthogonality

Going back to our linear algebra classes, inner products come in many forms. Given a vector of weights \(\mathbf{w} > 0\), we can define a weighted inner product:

The contribution of every two components at index \(i\) is weighted by the weight \(w_i\).

Similarly, given a weight function \(w(x)>0\) integrable over the domain \(D\), we can define a weighted inner product between two functions on that domain:

The contribution at each point \(x\) to the integral is weighted by the weight \(w(x)\). The Legendre basis, for example, is orthogonal on \(D = [-1, 1]\) according to the uniform weight \(w(x) = 1\).

Now let’s see what does it mean in terms of informativeness. Suppose without loss of generality that \(w(x)\) is normalized such that \(\int_{D} w(x) = 1\). If it’s not, we can always divide it by its integral. Now, it can be thought of as PDF of some probability distribution over \(x\), and therefore the inner product is just an expectation. Therefore if our data points \(x_1, \dots, x_n\) come from that distribution, then

Consequently, if \(f\) and \(g\) are orthogonal w.r.t our inner product, the two features generated by \(f(x)\) and \(g(x)\) are uncorrlated, or informative.

So ideally we would like to devise orthogonal bases w.r.t the PDF of our data distribution. The differential equations community has been designing orthogonal bases w.r.t various weights for a long time, and have come up with plenty of methods and an enormous body of literature. But there are two extremely simple tricks we can adopt from them for ML, which I’d like to discuss in this pose. Both are also discussed in the context of approximation theory and differential equations in the recent survey paper by Shen and Wang². One of them appears extremely intuitive, easy to implement, and is indeed useful in practice. The other one requires some careful math before coding, harder to apply in practice, and I’d like to present it and discuss when I believe it may be useful to.

The mapping trick

Let’s focus on the Legendre basis that is orthogonal on \([-1, 1]\). Instead of min-max scaling, which we did in previous posts, suppose we use some invertible and differentiable function \(\phi: D \to [-1, 1]\) that maps our feature from its original domain. In terms of the raw feature, our basis functions are

Are they orthogonal? Well, in some sense, they are. Using the change of variable \(y = \phi(x)\), we know from high-school calculus that :

The conclusion is simple - mapping with \(\phi\) results in orthogonal functions weighted by \(\phi'\). In particular, if \(\phi\) is a CDF of some distribution, then using the basis \(Q_0, Q_1, ...\) will result in uncorrelated features!

So what mapping should we use? If we know or can estimate the CDF \(W\) of our feature, we should use

Indeed, it maps to \([-1, 1]\), and the derivative of this mapping is twice the PDF. Just what we need.

We can, of course, attempt to do mathematical acrobatics to extend this to non-differentiable CDF functions \(W\), but this is not a paper, just a blog post. In fact, we will use a non-differentiable CDF in our code, without doing the math. This idea aligns with our intuition at the intro - mapping using a “uniformizing” transformation before computing Legendre polynomials produces an orthogonal basis w.r.t the original raw feature.

A small simulation

We shall sample data from some distributions, and use the above mapping to transform it before computing the Legendre vandermonde matrix. Then, we shall inspect the correlation between columns. Here is a function that accepts a scipy.stats distribution object, and computes the correlation matrix:

import numpy as np

def simulate_correlation(dist, degree=20, n_samples=10000):
    samples = dist.rvs(size=n_samples)
    mapped = 2 * dist.cdf(samples) - 1
    vander = np.polynomial.legendre.legvander(mapped)
    return np.corrcoef(vander.T)