Reflecting to optimise | Magnus Ross

This is nothing to be proud of, but I have never really studied optimisation in depth. Oh sure, I know my Adam from my AdaGrad and I even used L-BFGS one time, but when people start talking about dual spaces and convergence for L∞L^\infty continuous functions, I tend to glaze over a bit. For some reason I think in my head a lot of it seemed a bit old fashioned, whatever that means; why do I need to learn about optimising constrained convex functions in the gradient-descent-for-everything-wait-actually-just-use-a-pre-trained-model era? Isn’t that what they used to, like, make the cheapest possible diet? Boring! Well, today on the blog, I’d like to talk a bit about my optimisation blind spot in the context of an interesting problem I have been working on related to protein binder design. The first way you think to do something is rarely the best; here we’re going to discuss a concrete example of that. If you, like me, are not an optimisation expert, then I hope you’ll learn something useful from this post. If you are, then feel free to have a good old laugh at my ignorance! Setup Ok so what’s the setup? Let’s say we have a categorical probability distribution with kk categories where the probability of each category is given by the vector x∈Rk\mathbf{x}\in\mathbb{R}^k. There are some constraints on x\mathbf{x} for it to be valid:

the probabilities must be normalised: ∑i=1kxi=1\sum^{k}_{i=1} x_i = 1, and, the probabilities must be greater than 0: xi≥0  ∀i∈{1,⋯ ,k}x_i \geq 0 \ \ \forall i \in \{1, \cdots, k\}.

We have a non-convex function ff which takes in our probability vector, and spits out a real number. We want to find the x\mathbf{x} that minimises this function x∗=argminx∈Δf(x)\mathbf{x}^{*} = \text{argmin}_{\mathbf{x}\in\Delta}f(\mathbf{x}).

In our setup, we will assume we can compute the gradient ∇xf\nabla_\mathbf{x} f, but evaluation of the gradient, and indeed the function itself, is computationally expensive.

Protein related aside This is a simplified version of the problem of hallucination for de-novo binder design where x\mathbf{x} represents a distribution over k=20k=20 amino acids, and ff is a folding model, like Alphafold which takes a sequence of LL of these amino acid distributions, called a position specific scoring matrix (PSSM), as input. We want to find the sequence that gives the best fold (according to some metric like ipSAE). In this case the input to ff is now a matrix X∈Rk×L\mathbf{X} \in \mathbb{R}^{k \times L} where each column is a probability vector.

A first attempt When seeing a problem like this, where we need to optimise something with constraints, my first instinct is to try and re-parameterise the problem so that I don’t have to worry about them, and we can just use all the “normal” methods. We basically want to rewrite x\mathbf{x} as a function of some other parameters with no constraints, then we can just optimise those. In this case, we can write xi=softmax(ℓ)i=eℓi∑j=1keℓj \mathbf{x}_i = \text{softmax}(\mathbf{\ell})_i = \frac{e^{\ell_i}}{\sum^{k}_{j=1} e^{\ell_j}}

where we call ℓ∈Rk\ell\in \mathbb{R}^{k} the logits. Note that for any ℓ\mathbf{\ell}, the output x\mathbf{x} is guaranteed to be a valid probability vector obeying the constraints we set out earlier. Perfect! Problem solved! Now we can just find ℓ∗=argminℓ∈Rkf(softmax(ℓ))\mathbf{\ell}^{*} = \text{argmin}_{\mathbf{\ell}\in \mathbb{R}^{k}}f(\text{softmax}(\ell)) by taking gradients w.r.t. ℓ\ell and running some gradient descent method, and get x∗=softmax(ℓ∗)\mathbf{x}^* = \text{softmax}(\ell^*).

This all makes sense, and would have been case closed for me in the past, but it turns out there are other ways, and those ways are interesting! The simplex Let’s take a step back and think a bit more about the structure of this problem. It turns out there is a special name for the space of x\mathbf{x}’s that obey the constraints set out at the start: the probability simplex Δk−1\Delta^{k-1}, so we can write x∈Δk−1\mathbf{x} \in \Delta^{k-1}. The probability simplex is the k−1k-1 dimensional region of space where both non-negativity and normalisation are satisfied. Let’s take the concrete example of k=3k=3 with categories A,B,CA,B,C. The 2-simplex is a triangular region oriented perpendicular to the (1,1,1)(1, 1, 1) direction, it looks something like this:

2-simplex in 3D and 2D.

On the left we show the region in the full k=3k=3 dimensions, and on the right, we show the 2D “top down” view. The vertices of the simplex are the points where we have one category with certainty. On the faces, one of the categories doesn’t contribute (e.g. on the A/C face, B doesn’t contribute). Our optimisation problem boils down to finding the smallest value of ff that lies on this simplex. You’re projecting Let’s think about what happens if we just try the simplest possible thing we can think of: we start at a point xt\mathbf{x}_t on the simplex, take the gradient ∇xf(xt)\nabla_\mathbf{x} f(\mathbf{x}_t) and try to run a step of gradient descent: xt+1=xt−η∇xf(xt)\mathbf{x}_{t+1} = \mathbf{x}_t - \eta\nabla_\mathbf{x} f(\mathbf{x}_{t}). There are approximately four things that can happen:

We step “above” or “below” the simplex region, which happens if our gradient is not perpendicular to 1=(1,1,1)\mathbf{1}=(1, 1, 1).

In this case we would be adding to all the probabilities at once, which isn’t allowed because it breaks normalisation. Our gradient is perpendicular to 1\mathbf{1}, but we step “outside” the triangular region, which corresponds to one of the coordinates becoming negative. We really mess up, and do both 1. and 2. We do neither 1. or 2. and stay on the simplex 🥳

So in general, we’re going to end up off the simplex, in an invalid region. It turns out we can actually avoid case 1. pretty easily by just computing the component of the gradient that is perpendicular to 1\mathbf{1}, and stepping in that direction. We can do this by subtracting the mean of the gradient components.1 From now on we just assume we always make the gradients centred. This leaves the only problem with applying gradient descent to the raw probabilities as case 2., where we step outside the valid region. Here’s how that might look

One step of PGD on the 2-simplex.

with the black arrow representing the raw gradient descent step. We need to find a valid point in order to take the next step. To do this we compute the projection of the step onto the simplex, which is to say we find the nearest point on the simplex to the desired raw gradient step. This projection operation is represented by the red dotted arrow, with the final step given by the black dotted arrow. This process is called projected gradient descent (PGD) and is, in a way, conceptually simpler than our initial reparameterisation solution. If we want to optimise xx just run GD and correct it any time we go wrong. One step of PGD can be written as

xt+1=ΠΔk−1[xt−η∇xf(xt)] \mathbf{x}_{t+1} = \Pi_{\Delta^{k-1}}[\mathbf{x}_t - \eta\nabla_\mathbf{x} f(\mathbf{x}_{t})]

where ΠΔk−1\Pi_{\Delta^{k-1}} is the projection operator onto the k−1k-1 simplex. Computing the projection is actually not trivial, and we won’t go into details here, but there are ways to do it in O(k)\mathcal{O}(k) time by using a couple of cool tricks.

You can read about it in this classy ICML paper from ‘08. Basically it’s very fast relative to the cost of computing the gradient. It’s worth noting a few things that aren’t necessarily intuitive from our simple 2D case above when we scale to higher dimensions. Firstly, as the dimensionality of the simplex increases, the probability we step off the simplex with a gradient step becomes larger and larger; more points inside the simplex are close to an edge. Whenever we step off the simplex, we get projected back onto a face, where at least one coordinate is set to zero. In higher dimensions it becomes increasingly likely we end up projected on lower and lower dimensional faces, resulting in sparser steps. If we run a quick simulation and take a randomly directed step from one of the vertices, then in 2D the probability we step back inside the simplex is 0.170.17, for k=9k=9 it’s about 10−710^{-7}, and for k=20k=20 it’s essentially 0. Basically what all this amounts to is that PGD in high dimensions has a tendency to produce sparse solutions, which get sparser as optimisation proceeds. This is interesting and useful to bear in mind. Looking into the mirror Applying a softmax to get probabilities from unconstrained real numbers is of course a very common operation in machine learning, and so the reparameterisation in the first section feels natural, but it’s worth taking a step back and thinking about how we can justify this. Why not use a different function that maintains non-negativity/normalisation? We could do all sorts of fun things here. We’ve got to remember that since ff is non-convex, using different transformations is likely going to give us a totally different answer for the optimum, because we’re probably not going to find the global optimum. Because of that our choice of transformation is actually very important. To add some mathematical flesh to this choice, we’re going to need some extra machinery, in this case that’s something that’s intriguingly called mirror descent. If you’re good at maths you can read these lecture notes (chapter 17), which are very clear. I’m going to try to transmit the salient points. What’s actually going on when we do standard gradient descent?

When we take a gradient step, we are making the assumption that the local neighbourhood of our current iterate xt\mathbf{x}_t can be well approximated by a linear function (the first order Taylor approximation). When we take a step we want to minimise this approximation. If we just did that however, we’re going to end up at infinity, because a linear function has no minimum, so we need to make sure we stay in the neighbourhood. We can do this by introducing an additional penalty to our minimisation: the L2 distance from the iterate. So at each step we want to find argminx∈Rk[f(xt)+η⟨∇f(xt),x⟩+12∣∣x−xt∣∣2]. \text{argmin}_{\mathbf{x}\in \mathbb{R}^k}[f(\mathbf{x}_t) + \eta \langle \nabla f(\mathbf{x}_t), \mathbf{x} \rangle + \frac{1}{2} ||\mathbf{x} - \mathbf{x}_t||^2 ]. If you differentiate this w.r.t. x\mathbf{x} then set it to zero to find the minimum, then you get that η⋅∇f(xt)+(x−xt)=0  ⟹  xt+1=xt−η⋅∇f(xt),\eta \cdot \nabla f(\mathbf{x}_t) + (\mathbf{x} - \mathbf{x}_t) = 0 \implies \mathbf{x}_{t+1} = \mathbf{x}_t - \eta \cdot \nabla f(\mathbf{x}_t),so the optimum for each step gives us standard gradient descent. If you’re switched-on, you might be thinking: wait? Why the L2 norm? What’s the justification for that? Well dear reader, that would be a fair question, and actually the beauty is we can (subject to some rules) choose different penalties here, and generate different algorithms. The class of penalty we can choose from here are called Bregman divergences.