My training loss plateaued and wouldn’t budge. Obviously I’d screwed something up. I tried every hyperparameter combination, rewrote my loss function, spent days assuming I’d made some stupid mistake. Because it’s always user error.

This time, it wasn’t. It was a niche PyTorch bug that forced me through layers of abstraction I normally never think about: optimizer internals, memory layouts, dispatch systems, kernel implementations. Taught me more about the framework than years of using it.

I had a surprisingly fun time with this bug hunt and wrote up the whole investigation step-by-step, explaining framework internals as they become necessary to crack the case. If you enjoy debugging mysteries or find that tracking down bugs teaches you more than docs ever could, this might resonate. 🕵️‍♀️

Debugging post-mortems sometimes make me worry I wouldn’t have been smart enough to figure them out myself. So I structured this walkthrough to show the reasoning behind each step: what clues suggested each move, why I tested that hypothesis, why certain results pointed where they did. While the investigation took time and persistence, it didn’t require any particular expertise or wizardry— just observation and willingness to keep digging. I’ve included background knowledge exactly when you need it to understand the next step—think of it as an excuse to learn (or re-learn) PyTorch internals through a real problem. If you’d prefer to jump straight to reproducing the bug yourself, check out the minimal reproduction script and walkthrough on GitHub. Otherwise, join me on the investigation!

Table of Contents: 🤔 The Mystery: A Plateauing Loss…… 🔎 Isolating the Problem…… 💻 Device-Specific Differences…… ⌺ Tensor Memory Layouts…… 💔 Identifying the Broken Operations……. 🍎 Inside the Kernel Implementation…… 🕵️‍♀️ Case Closed

The Mystery: A Plateauing Loss

Training loss plateaued way too early. This felt like a standard hyperparameter issue- but I’d trained this same architecture on similar data with similar hyperparameters countless times and hit much lower losses.

What had changed? Those runs were months old. I tried reproducing them exactly, but couldn’t pin down the exact environment—the codebase had evolved through multiple projects, refactors, and dependency updates. Without a clean “before vs after,” I had to debug forward.

The architecture itself is straightforward: a two-layer sparse autoencoder (encoder –> sparse hidden layer –> decoder). However, it has some training quirks the could be potential culprits: the hidden layer uses TopK sparsity, where only the k largest activations remain (others are zeroed); the training process includes some manual gradient adjustments (gradient clipping for stability and modifications to decoder weight gradients); there’s an auxiliary loss term to encourage feature activation.

Even though I thought my initial hyperparameters were already well-tested, I tried everything: varied learning rates, tested different schedules, tried different k values and hidden dimensions, adjusted the auxiliary loss coefficients.

Nothing made a difference.

Meanwhile, my actual research sat on hold while I was stuck second-guessing everything: was my code broken? My data corrupted? And the creeping doubt- I’ve been doing ML for years, why can’t I make a simple two-layer autoencoder train properly?

The model was small enough that I was training on my MacBook (using the Apple Silicon GPU) and simple enough I could actually inspect every parameter. So after the standard checks turned up nothing, I started looking at the weights directly.

I visualized the weights at initialization and after the first few training steps. The decoder weights were updating- values shifting, gradients being applied, nothing crazy. But the encoder weights… weren’t updating at all. No NaNs, no suspicious patterns… they just… weren’t changing. They stayed exactly at their initialized values, down to the last decimal place.

Both layers participate in the same forward and backward pass. Why would one update and the other freeze completely?

Isolating the Problem

Are Gradients Flowing?

First check: are gradients even making it back to the encoder? The TopK sparsity should make gradients sparse—only the k activated features get gradients through backprop, the rest are zeroed. But maybe I messed up the implementation so that no encoder gradients flow at all? Or the manual gradient adjustments I was making somehow blocked everything?

After loss.backward(), the gradient statistics were:

The encoder gradients were there- and they were pretty big (as intended for my dataset)! And they were sparse (majority zeros) which was also expected, but there were still plenty of non-zero gradients. So gradients are definitely being calculated.

Is It the Optimizer?

Since the gradients exist but weights aren’t updating, the optimizer must be doing something wrong. Testing with a simpler optimizer, stochastic gradient descent (SGD):

# Manual SGD update
with torch.no_grad():
    model.encoder.weight -= 0.001 * model.encoder.weight.grad
# Encoder weights change! ✓

# Torch SGD update
sgd_optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
sgd_optimizer.step()
# Encoder weights change! ✓

# But with Adam...
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
optimizer.step()
# Encoder weights don't change! ✗

How Adam Works

To understand what might be breaking, I need to understand what Adam actually does differently from simple gradient descent.

Understanding Adam's Algorithm (click to collapse if familiar)

Problems with Vanilla SGD

SGD updates all parameters the same way:

# SGD: one learning rate for everything
param = param - learning_rate * gradient

This has a few problems:SGD has other problems too (hence all the optimizer research), but these are the ones Adam addresses.

1. Different parameters need different learning rates. Some parameters might consistently get gradients around 1000 while others get 0.01. With SGD’s fixed learning rate, you’re stuck: either you move too slowly on small gradients or you overshoot wildly on large ones.

2. The learning rate needs to change over time. Early in training, you want big steps to explore the space. Later, you need tiny steps to settle into a minimum. SGD requires manually decaying the learning rate on a schedule.

Adam’s Solution: Adaptive Learning Rates via Gradient Magnitude Tracking

Adam maintains two pieces of state per parameter and uses two hyperparameters to control how these states evolve:

State variables (initialized to zero for each parameter):

• exp_avg: Running average of gradients (first moment)
• exp_avg_sq: Running average of squared gradients (second moment)

Hyperparameters (typically beta_1=0.9, beta_2=0.999):

• beta_1: Decay rate for first moment (momentum)
• beta_2: Decay rate for second moment (gradient magnitude history)

Here’s the simplified algorithm:

Initialize state (done once per parameter)

exp_avg = zeros_like(param)
exp_avg_sq = zeros_like(param)
step = 0

Each training step:

# Update moments with exponential moving averages
exp_avg = beta_1 * exp_avg + (1 - beta_1) * grad
exp_avg_sq = beta_2 * exp_avg_sq + (1 - beta_2) * grad**2

# Update step count
# (It effectively starts at 1 to avoid division by zero in bias correction)
step += 1

# Bias correction
exp_avg_corrected = exp_avg / (1 - beta_1**step)
exp_avg_sq_corrected = exp_avg_sq / (1 - beta_2**step)

# Adaptive parameter update
param = param - lr * exp_avg_corrected / (sqrt(exp_avg_sq_corrected) + ε)

What Each Moment Does:

• First moment (exp_avg): Smooths out noisy gradients by averaging recent directions—like momentum in physics. When gradients oscillate (+10, -10, +8, -9…), the positive and negative values cancel out, revealing there’s no consistent direction. Beta_1=0.9 means “keep 90% of old momentum, add 10% of new gradient.” This smoothed momentum is what gets multiplied by the learning rate in the parameter update: lr * exp_avg.

• Second moment (exp_avg_sq): Tracks typical gradient magnitude for each parameter by averaging squared gradients. Squaring removes the +/- sign (both +10 and -10 become 100), preventing cancellation. Beta_2=0.999 means “keep 99.9% of magnitude history, add 0.1% of new squared gradient.” This magnitude normalizes the momentum-based update: lr * exp_avg / sqrt(exp_avg_sq). Parameters with consistently large gradients get their updates scaled down (large denominator), while parameters with small gradients get boosted (small denominator). This is how Adam achieves adaptive per-parameter learning rates.

• Epsilon (ε=1e-8): Prevents division by zero.

Bias Correction:

Both moments start at zero, causing early estimates to be biased toward zero. The correction factor (1 - β**step) provides a large boost early to counteract this, effectively “warming up” the optimizer over the first ~1000-3000 steps. As training progresses, the correction approaches 1 and has negligible effect.

The second moment works similarly. Without correction, exp_avg_sq would be only 0.1% of gradient² at step 1, but bias correction restores it to the full value.

For a deeper dive into Adam’s design and intuition, as well as other optimizers that use momentum and adaptive learning rates (RMSprop, AdaGrad, etc.), check out Stanford’s CS231n notes on optimization.

Knowing what Adam should be doing, let’s look at the state it’s maintaining (those exp_avg and exp_avg_sq tensors that track momentum and variance) to see what it’s actually doing.

Examining Adam’s State

For our frozen encoder, the maximum values in each state tensor were:

Wait, WHAT?! The encoder’s exp_avg_sq is zero despite having momentum accumulated in exp_avg.

This feels mathematically impossible… The second moment (exp_avg_sq) is zero despite non-zero gradients. Since exp_avg_sq stores squared gradients, it should NEVER be zero if gradients are non-zero.

And if it truly were zero, we’d see massive weight updates.

param_update = lr * exp_avg / (sqrt(exp_avg_sq) + ε) 
             = 0.001 * 1.96e5 / (sqrt(0) + 1e-8)
             = 196 / 1e-8
             = 1.96e10  # <-- HUGE!

This would be huge! Yet we see NO updates… this paradox points to a deeper issue.

Testing Hypotheses

Could it be bias correction?

Adam uses bias correction to counteract zero initialization. Having previously encountered subtle training issues due to Adam bias initialization bugs, I wondered if the correction might be broken here. 💡If you haven't been hurt by a bias correction bug before, check out these examples to learn the importance of getting this step right!

Recall, the bias correction is simply making our effective beta values dependent on the step index, so if the issue has to do with bias correction, it might have some relation to our beta parameters or step index.

I tested with different beta values, at different steps, and even beta_2=0 (which bypasses the exponential average entirely, making exp_avg_sq = grad**2 directly). The encoder’s exp_avg_sq still stayed zero, making bias correction seem less likely as a culprit.

Plus, exp_avg updated correctly despite using the same bias correction mechanism. So maybe something else is preventing exp_avg_sq from updating.

Is it a precision issue?

My largest gradients were big (1e6), and squared that’s 1e12. While that is quite large, it shouldn’t overflow in float32. However, I’ve also been hurt by precision bugs beforeFloating point precision issues have a fun habit of causing silent failures/degradations like this one (where it completes but produces incorrect values). Always worth checking, even when it seems unlikely., so I had to try it anyway.

I moved everything to float64… AND IT STARTED WORKING!

After a few more minutes of spiraling You're probably not reading this for the mid-debugging-self-doubt, but every debugging adventure has a spiraling moment (at least for me) so feels disingenuous to skip this step. And maybe one of these theories could've actually been correct! , I realized something: when I switched to float64, I also had to switch from MPS (Apple Silicon GPU) to CPU, since MPS doesn’t support float64. I’d changed two variables at once.

Testing with float32 on CPU… the weights update!!

﹡ This is progress!!

﹡ Note to self… simpler explanations are more likely correct- even (and especially!) when LLMs confidently assert complicated theories that are hard to understand / verify

﹡ Now I just need to figure out why the bug only occurs with MPS

Device-Specific Differences

Why the Same Operation Behaves Differently on Different Chips

PyTorch’s device abstraction lets you write the same code and run it on CPUs, GPUs, and even Apple Silicon. It feels like the same computation is running everywhere — but under the hood, each device has its own entirely separate implementation.

When you call a tensor operation like matmul, PyTorch looks at the tensor’s metadata (e.g. device, dtype, shape) and dispatches to a specialized kernel: a device-specific, highly optimized implementation tailored for that particular hardware backend.

So when you write something like result = tensor_a @ tensor_b, you’re not invoking a universal multiply function. PyTorch uses the tensors’ metadata to select a device- and dtype-specific kernel that performs the actual computation.

Multiplying two tensors on the CPU uses a completely different kernel than on MPS or CUDA. Even on the same device, changing the dtype or layout can trigger a different kernel. PyTorch maintains a large set of these implementations to support all the combinations.

We’ll see exactly how this dispatch system works in C++ later when we dive into the source code. For now, the important point is: even with identical Python code different tensor metadata → different kernel code → different efficiency / bugs.

In my case, because I’m running this on my M3 MacBook Pro, I’ m using MPS (Metal Performance Shaders), which is the GPU backend for Apple Silicon. While it feels a bit crazy to assume that my training plateau is due to an internal kernel-level bug, it’s a bit less unreasonable with MPS as it’s newer and less mature than the CPU and CUDA backends. (And honestly, most people training/debugging ML models are not doing it on their MacBooks.)

Why Does Only the Encoder Hit This Bug?

The Adam bug appears when working with the encoder on MPS. What makes the encoder different from the decoder that would trigger different behavior?

I tested everything I could think of that might differentiate the two tensors:

• Different gradient scales
• Dense vs sparse gradient patterns
• Removing decoder-specific gradient transformations
• Making encoder and decoder gradients statistically identical

Nothing helped. Even when both tensors had similar gradient statistics, only the encoder’s exp_avg_sq stayed frozen. The difference wasn’t in the values of the tensor - something else about the encoder tensor itself was triggering the bug.

What properties does a PyTorch tensor even have? I asked Claude what attributes could differ between two tensors and checked them one-by-one:

Three differences! The encoder and decoder have different shapes (they’re transposes of each other)PyTorch's nn.Linear stores weights as [out_features, in_features], so the encoder (384→1536) has shape [1536, 384] and the decoder (1536→384) has shape [384, 1536]., different stride patterns, and different contiguity. These properties are all related (more on that below).

The shape difference itself can’t cause different behavior (PyTorch operations handle any shape). But contiguity? That’s a low-level memory detail that could be relevant. Maybe the MPS Adam bug only affects non-contiguous tensors? Worth a shot:

model.encoder.weight.data = model.encoder.weight.contiguous()
optimizer.step()
# Encoder updates!! ✓

IT WORKS! But why?

Tensor Memory Layouts

What Does “Contiguous” Even Mean?

Your computer’s memory is just a flat, 1D array of bytes, but tensors represent multi-dimensional grids. When you index tensor[i, j], PyTorch needs to find that element in the flat memory. The tensor’s stride tells it how to do this conversion (and the exact amount you jump between elements depends on the dtype and how much memory each element takes up).

Think of stride as navigation instructions: “to get from one row to the next, skip this many elements.” By default, memory is stored row-wise—each row is stored sequentially, then the next row comes after. If you read through a row, you skip over 1 element at a time; to go to the next row, you move row-length elements over. (This is why going across a row is faster than going down a column.)

However, the memory layout doesn’t have to match the logical layout we use to think about the tensor. We can change how the user views the tensor without moving any data! For example, when we run transpose (.T), we don’t need to move around any data—we just change the stride!

As we see in the images, reading all the elements row-by-row in the contiguous tensor is easy and linear, but the same row-wise pattern in the non-contiguous tensor is much jumpier. This jumping pattern makes the tensor “non-contiguous.”

While there’s only one way for a tensor to be contiguous (the “natural” layout), there are many ways to become non-contiguous. By default, tensors are initialized as contiguous, but operations like slicing (tensor[::2, :]), reshaping, and dimension reordering (permute) can all create different non-contiguous stride patterns.

Why design tensors this way? Wouldn’t it be simpler to always keep data in the “natural” contiguous layout? The answer is performance: by just adjusting the tensor’s metadata, operations like transpose, slice, and reshape can be nearly instant— no data movement or memory allocation required. Keeping everything contiguous would mean expensive copying every time you reorganize dimensions.

How My Encoder Became Non-Contiguous

Looking at the weight initialization code:

self.encoder.weight.data = self.decoder.weight.T.clone()

The .T creates a non-contiguous view, and .clone() preserves the stride pattern.

x_t = x.T  # Start with non-contiguous
y_noncontig = x_t.clone()              # Preserves non-contiguous (1.919ms)
y_contig = x_t.clone(memory_format=torch.contiguous_format)  # Force contiguous (4.401ms)

Okay so we now know this initialization is why only the encoder is non-contiguous, and thus why only the encoder has training issues!

While I could just call .contiguous() on my encoder, declare victory, and get back to the research this bug was blocking me from doing… I felt like I was just scratching the surface of this bug and I feared it would haunt me until I fully figured out WHAT happened and WHY.

Identifying the Broken Operations

What Operations Does Adam Use?

When Adam updates parameters, what operations does it perform? Let’s look at PyTorch’s Adam implementation.

Fair warning: this file is over 1000 lines! To find what we need, search for where exp_avg and exp_avg_sq are defined and updated.

Here are the critical lines (lines 101, 391-407):

# State initialization (line 101)
state["exp_avg"] = torch.zeros_like(param, memory_format=torch.preserve_format)
state["exp_avg_sq"] = torch.zeros_like(param, memory_format=torch.preserve_format)

# ... [300 lines of setup and parameter group handling] ...

# First moment update (line 391)
exp_avg.lerp_(grad, 1 - beta1)

# Second moment update (line 392)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

# ... [bias correction calculations] ...

# Parameter update (line 407)
param.addcdiv_(exp_avg, denom, value=-step_size)

Look at that initialization! memory_format=torch.preserve_format means the state tensors inherit their stride pattern from param. So when our encoder weight is non-contiguous, both exp_avg and exp_avg_sq are also non-contiguous.

But they’re BOTH non-contiguous - so why does only one break?

Well, while they both are computed via addition and multiplication, they don’t use the exact same operations to perform this. Any of these operations could be a suspect, so let’s test each one individually!

For operations like output.addcmul_(input1, input2), the output tensorIn PyTorch, when a function name ends with an underscore (like mul_), that indicates that it is performing an in-place operation to modify a tensor directly in memory. Just as different devices can distinct kernels, so can distinctions like these! is modified while input tensors are read from. In our case, we know the output tensor is non-contiguous, so let’s test if that is sufficient to cause our bug.

Testing the Broken Operations

Testing each Adam operation with non-contiguous output tensors on MPS:

Interestingly, input contiguity doesn’t matter, only the output! Whether grad, exp_avg, or denom are contiguous makes no difference. The bug is purely in how these kernels write to non-contiguous output buffers.

The broken operations aren’t producing zeros or NaNs. They’re simply not modifying the output tensor at all. This wasn’t immediately obvious since exp_avg_sq was initialized to zeros, making “stays at zero” and “never updates” look identical. But testing with a non-zero, non-contiguous output tensor confirms that after calling addcmul_ or addcdiv_, the values remain unchanged. No update happens.

Yet timing shows MPS is doing substantial work. Non-contiguous operations take >2x longer than contiguous ones, proving the kernels are computing something, yet those results never make it to the output tensor. On CPU, each of these operations work correctly regardless of memory layout. This is purely a MPS-specific bug.

With the broken operations identified, we can trace the complete chain of events that triggers our failure:

Putting the Pieces Together

# Creates non-contiguous encoder weight (stride: 1, 1536)
encoder.weight = decoder.weight.T.clone()

# Both state tensors inherit non-contiguous layout from param
state["exp_avg"] = zeros_like(param, memory_format=torch.preserve_format)
state["exp_avg_sq"] = zeros_like(param, memory_format=torch.preserve_format)

exp_avg.lerp_(grad, 1-beta_1)  # ✓ Works fine

exp_avg_sq.mul_(beta_2)                        # ✓ Works fine
exp_avg_sq.addcmul_(grad, grad, 1-beta_2)      # ✗ No update - stays zero!

# Should update param, does nothing, leading to silent failure
param.addcdiv_(exp_avg, denom, value=-step_size)  # ✗ No update!

If only exp_avg_sq.addcmul_() failed, the zero exp_avg_sq would produce massive weight explosions (update = lr × exp_avg / √(ε)), making the bug immediately obvious. But param.addcdiv_() also failed, producing no updates at all!

The second bug masked the first, creating a silent failure: the spookiest type of error. The model appeared to be learning (the decoder was training normally), but progress stalled because the encoder stayed frozen. A subtle plateau that looked exactly like a hyperparameter issue 🙃

While we have made so much progress on this case, we’re still not done yet!!

Inside the Kernel Implementation

To understand why some operations work and others don’t, I needed to look at PyTorch’s source code for the buggy kernels.

While I normally trace through a Python codebase by jumping to definitions in my IDE, that doesn’t work with tensor.addcmul_(). When you call this function, there’s no Python source code executing - instead, Python immediately jumps into compiled C++ code for performance. And since PyTorch ships this as a pre-compiled binary, I can’t see that C++ implementation.

Tensor& addcmul_(Tensor& self, const Tensor& tensor1, const Tensor& tensor2, const Scalar& value)

tensor.addcmul_(tensor1, tensor2, value=1.0)

And as we discussed earlier, PyTorch dispatches based on tensor metadata, so there isn’t just one implementation - there are device-specific kernels for CPU, CUDA, MPS, etc. Since my PyTorch installation just has the compiled binary files, to investigate the actual implementations, we need to clone PyTorch’s repository.

PyTorch’s Dispatch System

All kernels are listed in an operation registry - a YAML file that maps operation names (like addcmul_) to their tensor-specific C++ implementations. In practice, when PyTorch is compiled (normally done before you install it), this registry is used to automatically generate hundreds of scripts that do the actual dispatching based on the patterns described here, but if we just want to understand what kernel our tensor is calling, we can look through the registry.

Searching for “addcmul_” in the registry native_functions.yaml:

- func: addcmul_(Tensor(a!) self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor(a!)
  # our addcmul_ function just points us to the yaml for addcmul.out
  structured_delegate: addcmul.out

# The function addcmul_ points to:
- func: addcmul.out(...)
  dispatch:
    CPU, CUDA: addcmul_out
    MPS: addcmul_out_mps  # Different function for MPS!

Now that we have the device-specific operation names, we can search them in the PyTorch repo within the mps implementations, and we find our implementation for addcmul_out_mps in PointwiseOps.mm. Upon a first skim of the code, I realized I had no clue how to read the MPS codebase. There were too many unknown variables and constructs, and I wasn’t sure what to look for in this implementation. I’d written a CUDA kernel before, and was pretty good with C about a decade ago, but as turns out, neither of those helped here :(

Comparing Broken vs Working Implementations

Rather than trying to decode unfamiliar code in isolation, I’d find something similar that works correctly and compare the two. mul_ was the perfect comparison since both are simple element-wise in-place operations. The registry pointed me to binaryOpTensor in BinaryOps.mm.

Now I had my comparison:

• Broken: addc_mul_div_out_mps in PointwiseOps.mm (used by addcmul_)
• Working: binaryOpTensor in BinaryOps.mm (used by mul_)

I opened both side-by-side, scanning specifically for differences in how they handle the output tensor. My experiments had already narrowed the search: I knew both operations were computing something (timing proved that), so the bug had to be in how results get written back to non-contiguous outputs. Look for anything related to contiguity checks or special output handling.

Broken version (addcmul_):

static void addc_mul_div_out_mps(..., Tensor& output, ...) {
  // ... setup code ...
  Placeholder outputPlaceholder = Placeholder(output);
  runMPSGraph(...);
  // That's it - no additional handling
}

Working version (mul_):

static void binaryOpTensor(..., Tensor& output, ...) {
  // ... setup code ...
  
  bool needsCopyToOutput = !output.is_contiguous();
  if (needsCopyToOutput) {
    // Create temporary contiguous tensor
    output = at::empty(...);
  }
  
  Placeholder outputPlaceholder = Placeholder(output);
  runMPSGraph(...);
  
  if (needsCopyToOutput) {
    output_.copy_(output);  // Copy results back!
  }
}

The working version explicitly checks !output.is_contiguous() and adds extra handling: it creates a temporary contiguous tensor, runs the operation, then copies results back. The broken version just passes the output directly to Placeholder and calls it a day.

But this raises a new question: if non-contiguous memory layouts need this kind of explicit handling, why doesn’t addcmul just crash or throw an error instead of silently failing?

The Memory Conversion Problem

The answer lies in understanding what Placeholder does. PyTorch tensors and Metal (Apple’s GPU framework) use different memory formats, so PyTorch needs a converter when running operations on Apple Silicon. Placeholder handles this conversion - it takes PyTorch tensors and wraps them in Metal-compatible buffers, handles different data types, manages memory layouts, and sets up the compute pipeline.

For most tensors, this conversion is straightforward. But for non-contiguous tensors, Metal can’t work with the scattered memory layout directly. Looking at the Placeholder code:

if (!src.is_contiguous()) {
    _tensor = src.clone(MemoryFormat::Contiguous);  // Create contiguous copy
    srcBuf = getMTLBufferStorage(_tensor);          // Point Metal to the copy
}

When Placeholder encounters a non-contiguous tensor, it automatically creates a contiguous copy and points Metal to that copy instead. This happens transparently - the broken kernels have no idea they’re working with a temporary.

This automatic copying is perfect for input tensors - Metal reads from the copy, computation proceeds normally, and nobody cares what happens to the temporary afterward.

But it’s disastrous for output tensors where the goal is in-place editing. The computation succeeds and writes results to the temporary copy, but those results never make it back to the original tensor that’s supposed to be updated.

The Complete Bug Mechanism

The broken kernels work perfectly with contiguous tensors and silently fail with non-contiguous ones. The working kernels detect this situation and add an explicit copy-back step to move results from the temporary to the original tensor.

The Fix

Understanding the bug made the solution clear - apply the same pattern that working kernels use:

I tested this locally and it worked! The encoder weights finally updated and the model trained successfully 🎉🎉

You can see the complete reproduction, debugging experiments, fix at https://github.com/ElanaPearl/pytorch-mps-noncontiguous-bug.

Case Closed

A Lesson in Version Control

While editing a Python package just involves installing your locally editable version of the code instead of the default package, to test my PyTorch fix, I had to re-build it all locally, which was more work than expected and also made me acutely aware that this whole time I was working on PyTorch v2.2.1I was working on a research codebase with dependency conflicts that blocked upgrading PyTorch. Common enough situation, but lesson learned: always check versions early in debugging, even if you can't immediately update! (as this fact made it difficult to build and I had to downgrade things like CMake and deal with weird version conflicts to even build this older PyTorch).

Checking the latest version revealed the bug was already fixed in v2.4, patched by an ML engineer at Apple last year using almost the exact same approach I’d used.The official fix uses slightly different syntax; but the same core pattern: detect non-contiguous output, create a contiguous temporary buffer, perform the computation, then copy results back to the original tensor. This updated code even informed me that in macOS 15+, MPS now handles non-contiguous tensors natively! In macOS 15, Apple added native strided array support to MPSGraph via the arrayView API (see WWDC 2024 session at timestamp 13:41). Instead of the gather-scatter workaround, Metal can now read/write directly from non-contiguous memory using stride metadata. This means on macOS 15+, PyTorch can skip the manual copy workarounds entirely. The performance gap between contiguous and non-contiguous tensors is now much smaller, though contiguous is still faster due to better cache utilization.

The Pattern Strikes Again

While writing this up, I added some more tests for my kernel fix to confirm it really worked, and one of the tests failed! I looked into it more and realized I’d stumbled upon the same failure pattern in the random_ operation (in the most up-to-date PyTorch this time!)

Turns out, all random in-place operations (normal_, uniform_, exponential_, random_, bernoulli_) silently fail when called on non-contiguous tensors on MPS.

x = torch.zeros(10, 10).T  # Non-contiguous
x.normal_()  # Should fill with random values
print(x.max())  # Prints 0.0 - the operation silently failed!

Yet again, the operations complete without error, but the tensor remains unchanged—the kernel computes random values into a temporary contiguous buffer but never copies them back.

Having just traced through this exact bug pattern, I recognized it immediately and knew exactly how to fix it. Filed an Issue and made a PR applying the same solution.

I suspect there are other similar bugs lying around, as none of these fixes actually address the underlying quirk that the Placeholder abstraction itself is problematic when used with output tensors.

The core issue: Placeholder’s constructor silently creates a temporary contiguous copy for non-contiguous tensors, but it has no way to know if it’s wrapping an input (where the copy is fine- we just read from it) or an output (where the copy is broken- results get written to it then lost). This means every single operation that uses Placeholder for outputs must manually implement the same workaround pattern or else it has this silent failure:

// Every MPS operation must remember to do this:
bool needsCopy = !output.is_contiguous();
Tensor temp = needsCopy ? at::empty(...) : output;
@autoreleasepool {
    Placeholder p(temp);
    runGraph();
}
if (needsCopy)
  output.copy_(temp);

This is a leaky abstractionA "leaky abstraction" is when an abstraction that's supposed to hide implementation details forces you to understand and work around those details anyway. Placeholder is supposed to abstract Metal buffer management, but its internal copying leaks through, forcing every caller to manually handle non-contiguous outputs. See Joel Spolsky's The Law of Leaky Abstractions for the canonical explanation.: the internal implementation detail that “Placeholder makes temporary copies” has leaked out to every caller, making it each operation’s responsibility to work around. A better design would be:

• Placeholder knows input vs output: Pass a flag so Placeholder can handle the copy-back itself
• Separate abstractions: Different wrapper types for inputs (InputPlaceholder) and outputs (OutputPlaceholder)
• Make the temporary explicit: Don’t hide the copy inside Placeholder—make callers explicitly create and manage contiguous temporaries (this is what I used in the fixes for addcmul_/addcdiv_/the random ops)

The good news: macOS 15+ Metal now handles non-contiguous tensors natively, making this entire issue obsolete for newer systems. But for anyone on older macOS versions or maintaining PyTorch’s MPS backend, this abstraction continues to cause issues.

So ideally, the Placeholder class would be redesigned to handle output tensors correctly by default, but given that the hardware is moving to handle this natively anyway, the pragmatic fix is probably just to audit and patch the remaining operations using the established pattern.

Practical Takeaways for Your Code

Performance Considerations

Even with the code fixes, non-contiguous tensors on MPS involve: Allocate temporary buffer -> Copy to contiguous layout -> Compute -> Copy back. Making tensors contiguous once at initialization avoids thousands of copies during training! And even if your OS can avoid making this temporary contiguous copy, it is still slower to operate on non-contiguous memory if you will be using it many times.

When to Call .contiguous()

# When to call .contiguous() - General Principles

# 1. After operations that change memory layout:
x = tensor.transpose(0, 1)  # Non-contiguous
x = tensor.view(-1)          # Might fail if non-contiguous!
x = x.contiguous().view(-1)  # Safe

# 2. Before operations that might not handle strides:
# - Custom CUDA/Metal kernels  
# - Newer backend features
# - Operations that failed mysteriously on certain devices

# 3. For performance on repeated operations:
weights = init_weights().T   # Used in every forward pass
weights = weights.contiguous()  # Pay copy cost once, not every iteration

# But don't overuse it!
x = x + y  # Creates new contiguous tensor anyway
x = x.contiguous()  # Unnecessary copy!

For MPS specifically: If on macOS <15, make sure all your parameters are contiguous!

What I Learned

Isolate to specific, measurable symptoms. The most standard advice and for such good reason. Everything got easier once I had a concrete target: “exp_avg_sq stays at zero” is infinitely more debuggable than “the loss plateaus mysteriously.” Once I had a specific symptom, I could strip away components and test the minimal case that triggered it.

When debugging tensor issues, check metadata not just values. I was checking for NaNs, visualizing weights, inspecting gradients—all focused on the numbers inside tensors. The actual problem was the tensor’s stride pattern. Device, dtype, contiguity, memory layout—these aren’t just performance details, they can cause silent correctness bugs. tensor.is_contiguous() is now part of my debugging checklist.

When I’m confused, I might have changed two things—or there might be two bugs. Switching to fp64 “fixed” it, but I’d also switched from MPS to CPU. Untangling that revealed the real culprit. And exp_avg_sq staying zero should have caused explosions, but the parameter update also failed—one bug perfectly masked the other.

Documentation makes more sense when I need it. I’d skimmed PyTorch internals docs before and nothing stuck—dispatch systems, stride patterns, kernel implementations all felt overwhelming. But once I had to understand how addcmul_ dispatches to MPS kernels, everything clicked. Now PyTorch feels less like a black box. And when I hit the random ops bug weeks later, I wasn’t intimidated—I knew exactly how to trace through the source.

Explore the system before exploring the code. When I needed to debug addcmul_out_mps in unfamiliar MPS code, I ran experiments first: which operations fail? Do they run at all? What triggers the bug? By the time I opened the source, I knew to compare addcmul_ (broken) against mul_ (working) and scan specifically for differences in output handling. Without that context, I’d have been lost in Objective-C++ with no idea what mattered. Also LLMs were very helpful with unfamiliar constructs like MPSGraphTensor or @autoreleasepool, although they’re still less reliable with MPS than more documented frameworks.

Write post-mortems– even for yourself. Forcing myself to explain why I tried each debugging step was as educational as the original investigation. It’s like experience replay in RL: you explore many failed paths, find one that works, then replay that successful trajectory to reinforce the policy. Writing it down builds pattern recognition—when I’m in “situation A”, what hypotheses are worth trying? I’ve written lower-effort debugging debriefs before, but making this one readable for an external audience forced me to articulate why each step made sense, deepening my understanding of what actually worked.

What started as a frustrating research roadblock became a surprisingly fun & educational detour. It forced a closer look at things normally taken for granted: Adam’s momentum mechanics, stride patterns, kernel dispatch. Understanding why each operation behaved differently revealed more about PyTorch’s architecture than typical usage ever does.

If you made it this far, thanks for joining! Hope you had fun and/or learned something & happy debugging!

Special thanks to Nicholas Joseph, Ben Kuhn, Nelson Elhage and Alex Tamkin for giving feedback on this 💜

Who should read this

Do you want to understand exactly how AlphaFold3 works? The architecture is quite complicated and the description in the paper can be overwhelming, so we made a much more friendly (but just as detailed!) visual walkthrough.

This is mostly written for an ML audience and multiple points assume familiarity with the steps of attention. If you’re rusty, see Jay Alammar’s The Illustrated Transformer for a thorough visual explanation. That post is one of the best explanations of a model architecture at the level of individual matrix operations and also the inspiration for the diagrams and naming.

There are already many great explanations of the motivation for protein structure prediction, the CASP competition, model failure modes, debates about evaluations, implications for biotech, etc. so we don’t focus on any of that. Instead we explore the how.

How are these molecules represented in the model and what are all of the operations that convert them into a predicted structure?

This is probably more exhaustive than most people are looking for, but if you want to understand all the details and you like learning via diagrams, this should help :)

Architecture Overview

We’ll start by pointing out that goals of the model are a bit different than previous AlphaFold models: instead of just predicting the structure of individual protein sequences (AF2) or protein complexes (AF-multimeter), it predicts the structure of a protein, optionally complexed with other proteins, nucleic acids, or small molecules, all from sequence alone. So while previous AF models only had to represent sequences of standard amino acids, AF3 has to represent more complex input types, and thus there is a more complex featurization/tokenization scheme. Tokenization is described in its own section, but for now just know that when we say “token” it either represents a single amino acid (for proteins), nucleotide (for DNA/RNA), or an individual atom if that atom is not part of a standard amino acid/nucleotide.

Interactive Table of Contents

Full architecture. If you click on any part of the architecture, it will take you to that section of the post. If you resize the page, you might need to refresh to keep the interactive part working. (Diagram modified from AF3 paper)

The model can be broken down into 3 main sections:

1. Input Preparation The user provides sequences of some molecules to predict structures for and these need to be embedded into numerical tensors. Furthermore, the model retrieves a collection of other molecules that are presumed to have similar structures to the user-provided molecules. The input preparation step identifies these molecules and also embeds these as their own tensors.
2. Representation learning Given the Single and Pair tensors created in section 1, we use many variants of attention to update these representations.
3. Structure prediction We use these improved representations, and the original inputs created in section 1 to predict the structure using conditional diffusion.

We also have additional sections describing 4. the loss function, confidence heads, and other relevant training details and 5. some thoughts on the model from an ML trends perspective.

Notes on the variables and diagrams

Throughout the model a protein complex is represented in two primary forms: the “single” representation which represents all the tokens in our protein complex, and a “pair” representation which represents the relationships (e.g. distance, potential interactions) between all pairs of amino acids / atoms in the complex. Each of these can be represented at an atom-level or a token-level, and will always be shown with these names (as established in the AF3 paper) and colors:

• The diagrams abstract away the model weights and only visualize how the shapes of activations change
• The activation tensors are always labeled with the dimension names used in the paper and the sizes of the diagrams vaguely aim to follow when these dimensions grow/shrink. The hidden dimension names usually start with "c" for "channel". For reference the main dimensions used are c_z=128, c_m=64, c_atom=128, c_atompair=16, c_token=768, c_s=384.
• Whenever possible, the names above the tensors in this (and every) diagram match the names of the tensors use in the AF3 supplement. Typically, a tensor maintains its name as it goes through the model. However, in some cases, we use different names to distinguish between versions of a tensor at different stages of processing. For example, in the atom-level single representation, c represents the initial atom-level single representation while q represents the updated version of this representation as it progresses through the Atom Transformer.
• We also ignore most of the LayerNorms for simplicity but they are used everywhere.

1. Input Preparation

The actual input a user provides to AF3 is the sequence of one protein and optionally additional molecules. The goal of this section is to convert these sequences into a series of 6 tensors that will be used as the input to the main trunk of the model as outlined in this diagram. These tensors are s, our token-level single representation, z, our token-level pair representation, q, our atom-level single representation, p, our atom-level pair representation, m, our MSA representation, and t, our template representation.

This section contains:

• Tokenization describes how molecules are tokenized and clarifyies the difference between atom-level and token-level
• Retrieval (Create MSA and Templates) expalains why and how we include additional inputs to the model. It creates our MSA (m) and structure templates (t).
• Create Atom-Level Representations creates our first atom-level representations q (single) and p (pair) and includes information about generated conformers of the molecules.
• Update Atom-Level Representations (Atom Transformer) is the main “Input Embedder” block, also called the “Atom Transformer”, which gets repreated 3 times and updates the atom-level single representation (q). The building blocks introduced here (Adaptive LayerNorm, Attention with Pair Bias, Conditioned Gating, and Conditioned Transition) are also relevant later in the model.
• Aggregate Atom-Level -> Token-Level takes our atom-level representations (q, p) and aggregates all the atoms that at part of multi-atom tokens to create token-level representations s (single) and z (pair) and includes information from the MSA (m) and any user-provided information about known bonds that involve ligands.

Tokenization

See where this fits into the full architecture

In AF2, as the model only represented proteins with a fixed set of amino acids, each amino acid was represented with its own token. This is maintained in AF3, but additional tokens are also introduced for the additional molecule types that AF3 can handle:

• Standard amino acid: 1 token (as per AF2)
• Standard nucleotide: 1 token
• Non-standard amino acids or nucleotides (methylated nucleotide, amino acid with post-translational modification, etc.): 1 token per atom
• Other molecules: 1 token per atom

As a result, we can think of some tokens (like those for amino acids) as being associated with multiple atoms, while other tokens (like those for an atom in a ligand) are associated with only a single atom. So, while a protein with 35 standard amino acids (likely > 600 atoms) would be represented by 35 tokens, a ligand with 35 atoms would also be represented by 35 tokens.

Retrieval (Create MSA and Templates)

See where this fits into the full architecture

One of the key early steps in AF3 is something akin to Retrieval Augmented Generation RAG in language models. We find similar sequences to our protein and RNA sequences of interest (collected into a multiple sequence alignment, “MSA”), and any structures related to those (called the “templates”), then include them as additional inputs to the model called m and t, respectively.

(Image from AF2)

So how are these sequences and structures retrieved?

First, a genetic search is done searching for any protein or RNA chains that resemble any input protein or RNA chains. This does not involve any training and relies upon existing Hidden Markov Model (HMM) based methodsSpecifically, they use jackhmmer, HHBlits, and nhmmer to scan multiple protein databases and RNA databases for relevant hits. Then these sequences are aligned to each other to construct an MSA with N_MSA sequences. As the computational complexity of the model scales with N_MSA they limit this to N_MSA < 2¹⁴. Typically, MSAs are constructed from individual protein chains but, as described in AF-multimer, instead of just concatenating the separate MSAs together into a block diagonal matrix, certain chains from the same species can be 'paired' as described here. This way, the MSA does not have to be as large and sparse, and evolutionary information can be learned about relationships between chains.

Then, for each protein chain, they use another HMM-based method (hmmsearch) to find sequences in the Protein Data Bank (PDB) that resemble the constructed MSA. The highest quality structures are selected and up to 4 of these are sampled to be included as "templates".

The only new part of these retrieval steps compared to AF-multimer is the fact that we now do this retrieval for RNA sequences in addition to protein sequences. Note that this is not traditionally called “retrieval” as the practice of using structural templates to guide protein structure modeling has been common practice in the field of homology modeling long before the term RAG existed. However, even though AlphaFold doesn’t explicitly refer to this process as retrieval, it does quite resemble what has now been popularized as RAG.

How do we represent these templates?

From our template search, we have a 3D structure for each of our templates and information about which tokens are in which chains. First, the euclidean distances between all pairs of tokens in a given template are calculated. For tokens associated with multiple atoms, a representative “center atom” is used to calculate distances. This would be the C_ɑ atom for amino acids and C¹’ atom for standard nucleotides.

Highlighting "center atoms" in single-token building blocks

This generates a N_token x N_token matrix for each template. However, instead of representing each distance as a numerical value, the distances get discretized into a “distogram” (a histogram of distances).Specifically, the values are binned into 38 bins between 3.15A and 50.75A and there's 1 additional bin for any distances bigger than that.

To each distogram, we then append metadata about which chain In molecular complexes, a chain refers to a distinct molecule or part of a molecule. This can be a protein chain (a sequence of amino acids), a DNA or RNA chain (a sequence of nucleotides), or other biomolecules. AlphaFold uses chain information to differentiate between parts of a complex, helping it predict how these parts interact to form the overall structure each token belongs to, whether this token was resolved in the crystal structure, and information about local distances within each amino acid. We then mask out this matrix such that we only look at distances within each chain (e.g., we ignore the distances between chain A and chain B) as they “make no attempt to select templates… to gain information about inter-chain interactions”‘ It is not specified why, but note that while there is no inter-chain interactions in the templates, they do incorporate them the MSA construction. .

Create Atom-Level Representations

See where this fits into the full architecture

To create q, our atom-level single representation, we need to pull all our atom-level features. The first step is to calculate a “reference conformer” for each amino acid, nucleotide, and ligand. While we do not yet know the structure of the entire complex, we have strong priors on the local structures of each individual component. A conformer (short for conformational isomer) is a 3D arrangement of atoms in a molecule that is generated by sampling rotations about single bonds. Each amino acid has a “standard” conformer which is just one of the low-energy conformations this amino acid can exist in, which can be retrieved through a look-up. However, each small molecule requires its own conformation generation. These are generated with RDKit’s ETKDGv3, an algorithm that combines experimental data and torsion angle preferences to produce 3D conformers.

Then we concatenate the information from this conformer (relative location) with each atom’s charge, atomic number, and other identifiers. Matrix c stores this information for all the atoms in our sequencesIn the AF3 supplement, the atom-level matrices (c, and q) are typically referred to in their vector forms (e.g. c_l or c_m), where l and m are used to index atoms.. We then use c to initialize our atom-level pair representation p to store the relative distances between atoms. Because we only know reference distances within each token, we use a mask (v) to ensure this initial distance matrix only represents distances we’ve calculated in the conformer generation. We also include a linear embedding of the inverse square of the distances, add to it a projection of c_l and c_m, and update this with a few more linear layers with residual connectionsThe AF3 paper doesn't really clarify why this additional inverse distance step is performed or contain ablations for their effect of it; so, as with many of the steps we will discuss, we can only assume they were empirically shown to be useful.In the AF3 supplement, the p tensor is typically referred to in its vector form p_l,m (where this represents the relationship between atom l and atom m)..

Finally, we make a copy of our atom-level single representation, calling this copy q. This matrix q is what we will be updating going forward, but c does get saved and used later.

Update Atom-Level Representations (Atom Transformer)

See where this fits into the full architecture

Having generated q (representation of all the atoms) and p (representation of each pair of atoms), we now want to update these representations based on other atoms nearby. Anytime AF3 applies attention at the atom-level, we use a module called the Atom Transformer. The atom transformer is a series of blocks that use attention to update q using both p and the original representation of q called c. As c does not get updated by the Attention Transformer, it can be thought of as a residual connection to the starting representation.

The Atom Transformer mostly follows a standard transformer structure using layer norm, attention, then an MLP transition. However, each step has been adapted to include additional input from c and p (including a secondary input here is sometimes referred to as “conditioning”.) There is also a ‘gating’ step between the attention and MLP blocks. Going through each of these 4 steps in more detail:

1. Adaptive LayerNorm

Adaptive LayerNorm (AdaNorm) is a variant of LayerNorm with one simple extension. Recall that for a given input matrix, traditional LayerNorm learns two parameters (a scaling factor gamma and a bias factor beta) that adjust the mean and standard deviation of each of the channels in our matrix. Instead of learning fixed parameters for gamma and beta, AdaNorm learns a function to generate gamma and beta adaptively based on the input matrix. However, instead of generating the parameters based on the input getting re-scaled (in the Atom Transformer this is q), a secondary input (c in the Atom Transformer) is used to predict the gamma and beta that re-scale the mean and standard deviation of q.

2. Attention with Pair Bias

Atom-Level Attention with Pair-Bias can be thought of as an extension of self-attention. Like in self-attention, the queries, keys, and values all come from the same 1D sequence (our single representation, q). However, there are 3 differences:

1. Pair-biasing: after the dot product of the queries and keys are calculated, a linear projection of the pair representation is added as a bias to scale the attention weights. Note that this operation does not involve any information from q being used to update p, just one way flow from the pair representation to q. The reasoning for this is that atoms that have a stronger pairwise relationship should attend to each other more strongly and p is effectively already encoding an attention map.

2. Gating: In addition to the queries, keys, and values, we create an additional projection of q that is passed through a sigmoid, to squash the values between 0 and 1. Our output is multiplied by this “gate” right before all the heads are re-combined. This effectively forces the model to ignore some of what it learned in this attention process. This type of gating appears frequently in AF3 and is discussed more in the ML-musings section. To briefly elaborate, because the model is constantly adding the outputs of each section to the residual stream, this gating mechanism can be thought of as the model’s way to specify what information does or does not get saved in this residual stream. It is presumably named a “gate” after the similar “gates” in LSTM which uses a sigmoid to learn a filter for what inputs get added to the running cell state.

3. Sparse attention:

Because the number of atoms can be much larger than the number of tokens, we do not run full attention at this step, rather, we use a type of sparse attention (called Sequence-local atom attention) in which the attention is effectively run in local groups where groups of 32 atoms at a time can all attend to 128 other atoms. Sparse attention patterns are more thoroughly described elsewhere on the internet.

3. Conditioned Gating

We apply another gate to our data, but this time the gate is generated from our origin atom-level single matrix, cAs with so many steps, it is unclear why it is done this way and what the benefit of conditioning on the original representation c does as opposed to learning the gate from the primary single representation q.

4. Conditioned Transition

This step is equivalent to the MLP layers in a transformer, and is called “conditioned” because the MLP is sandwiched in between Adaptive LayerNorm (Step 1 of Atom Transformer) and Conditional Gating (Step 3 of Atom Transformer) which both depend on c.

The only other piece of note in this section is that AF3 uses SwiGLU in the transition block instead of ReLU. The switch from ReLU → SwiGLU happened with AF2 → AF3 and has been a common change in many recent architectures so we visualize it here.

With a ReLU-based transition layer (as in AF2), we take the activations, project them up to 4x the size, apply a ReLU, then down-project them back to their original size. When using SwiGLU (in AF3), the input activation creates two intermediate up-projections, one of which goes through a swish non-linearity (improved variant of ReLU), then these are multiplied before down-projecting. The diagram below shows the differences:

Aggregate Atom-Level → Token-Level

See where this fits into the full architecture

While the data so far has all been stored at an atom-level, the representation learning section of AF3 from here onwards operates at the token-level. To create these token-level representations, we first project our atom-level representation to a larger dimension (c_atom=128, c_token=384). Then, we take the mean over all atoms assigned to the same token. Note that this only applies to the atoms associated with standard amino acids and nucleotides (by taking the mean across all atoms attached to the same token), while the rest remain unchangedThe AF3 paper describes these molecule types as having a representative atom per token (the center atom). Recall that this is the C_α atom for amino acids and C¹' atom for standard nucleotides. So while we mostly consider this reduced representation as "token space", we can also think of each token as representing a single atom (either a representative C_α/C¹' atom or an individual atom)..

Now that we are working in “token space”, we concatenate our token-level features and statistics from our MSA (where available)e.g., The amino acid type (dim = 32), distribution of amino acids at this position in our MSA (dim = 32), and the deletion mean at this token (dim = 1) from our MSA. Note that these values will be zero for ligand atoms not associated with an MSA.. This matrix, s^inputs, having grown a bit from these concatenations, is projected back down to c_token, and called s^init: the starting representation of our sequence that will be updated in the representation learning section. Note that s^init gets updated in the representation learning section, but s^inputs are saved to be used later in the structure prediction section.

Now that we have created s^init, our initialized single representation, the next step is to initialize our pair representation z^init. The pair representation is a three dimensional tensor, but it’s easiest to think of it as a heatmap-like 2D matrix with an implicit depth dimension of c_z=128 channels. So, entry z_i,j of our pair representation is a c_z dimensional vector meant to store information about the relationship between token i and token j in our token sequence. We have created an analogous atom-level matrix p, and we follow a similar process here at the token-level.

To initialize z_i,j, we use a linear projection to make the channel dimension of our sequence representation match that of the pair representation (384 → 128) and add the resulting s_i and s_j. To this, we add a relative positional encoding, p_i,jThis encoding consists of a^rel_pos, a one-hot encoding of the offset of the two token ids in token space (or set to a maximum of 65 if the two tokens are not on the same chain), a^rel_token, a one-hot encoding of the offset of the two token ids in token space (or set to a maximum of 65 if the tokens are part of different amino acids or nucleotides), and a^rel_chain, encoding the offset of the two chains the tokens are on. We project this concatenated encoding into the dimensionality of z too.. If the user has also specified particular bonds between tokens, those are linearly embedded here and added to that entry in the pair representation.

Now we’ve successfully created and embedded all of the inputs that will be used in the rest of our model:

For Step 2, we will set aside the atom-level representations (c, q, p) and focus on updating our token-level representations s and z in the next section (with the help of m and t).

2. Representation Learning

(Diagram modified from full AF3 architecture diagram)

This section is the majority of the model, often referred to as the “trunk”, as it is where most of the computation is done. We call it the representation learning section of the model, as the goal is to learn improved representations of our token-level “single” (s) and “pair” (z) tensors initialized above. Recall that we refer to the "single" sequence representations, these are not necessarily the sequence of one protein, but rather the concatenated sequence of all the atoms or tokens in our structure (which could contain multiple separate molecules).

This section contains:

1. Template module updates z using the structure templates t
2. MSA module first updates the MSA m, then adds it to the token-level pair representation z. In this section we spend significant time on two operations:
   • The Outer Product Mean enables m to influence z
   • MSA Row-wise Gated Self-Attention Using Only Pair Bias updates m based on z and is a simplified version of attention with pair-bias (intended for MSAs)
3. Pairformer updates s and z with geometry-inspired (triangle) attention. This section mostly describes the triangle operations (used extensively throughout both AF2 and AF3).
   • Why look at triangles? explains some intuition for the triangle operations
   • Triangle Updates and Triangle Attention both update z using methods similar to self-attention, but inspired by the triangle inequality
   • Single Attention With Pair Bias updates s based on z and is the token-level equivalent of attention with pair-bias (intended for single sequences)

Each individual block is repeated multiple times, and then the output of the whole section is fed back into itself again as input and the process is repeated (this is called recycling).

Template Module

See where this fits into the full architecture

Each template (N_templates=2 in the diagram) goes through a linear projection and is added together with a linear projection of our pair representation (z). This newly combined matrix goes through a series of operations called the Pairformer Stack (described in depth later). Finally, all of the templates are averaged together and go through - you guessed it - another linear layer.This is both called the template module and template embedder depending on where you look in the AF3 supplement, but they seem to just refer to the same thing. Interestingly, this last linear layer has a ReLU as the non-linearity which wouldn’t be particularly notable except for the fact that it is one of only two places ReLU is used as the non-linearity in AF3. As always, can only hypothesize as to why this was selected.

MSA Module

See where this fits into the full architecture

Architecture of MSA Module. {Diagram from AF3}

This module greatly resembles what was called “Evoformer” in AF2, and the goal of it is to simultaneously improve the MSA and pair representations. It does a series of operations independently on these two representations then also enables cross-talk between them.

The first step is to subsample the rows of the MSA, rather than use all rows of the MSA previously generated (which could be up to 16k), then add a projected version of our single representation to this subsampled MSA.

Outer Product Mean

Next, we take the MSA representation and incorporate it into the pair representation via the “Outer Product Mean”. Comparing two columns of the MSA reveals information about the relationships between two positions in the sequence (e.g. how correlated are these two positions in the sequence across evolution). For each pair of token indices i,j, we iterate over all evolutionary sequences, taking the outer product of m_s,i and m_s,j, then averaging these across all the evolutionary sequences. We then flatten this outer product, project it back down, and add this to the pair representation z_i,j (full details in diagram). While each outer product only compares values within a given sequence m_s, when we take the mean of these, that mixes information across sequences. This is the only point in the model where information is shared across evolutionary sequences. This is a significant change to reduce the computational complexity of the Evoformer in AF2.

Row-wise gated self-attention using only pair bias

Having updated the pair representation based on the MSA, the model next updates the MSA based on the pair representation. This specific update pattern is called row-wise gated self attention using only pair bias, and is a simplified version of self attention with pair bias, discussed in the Atom Transformer section, applied to every sequence (row) in the MSA independently. It is inspired by attention, but instead of using queries and keys to determine what other positions each token should attend to, we just use the existing relationships between tokens stored in our pair representation z.

In the pair representation, each z_i,j is a vector containing information about the relationship between tokens i and j. When the tensor z gets projected down to a matrix, each z_i,j vector becomes a scalar that can be used to determine how much token i should attend to token j. After applying row-wise softmax, these are now equivalent to attention scores, which are used to create a weighted average of the values as a typical attention map would.

Note that there is no information shared across the evolutionary sequences in the MSA as it is run independently for each row.

Updates to pair representation

The last step of the MSA module is to update the pair representation through a series of steps referred to as triangle updates and attention. These triangle operations are described below with Pairformer, where they are used again. There are also some transition blocks that use SwiGLU to up/down project the matrix as was done in the Atom Transformer.

Pairformer module

See where this fits into the full architecture

Diagram from AF3 supplement

Having used the templates and MSA to update our pair representation, we now ignore them for the rest of the model. Instead, only the updated pair representation (z) and single representation (s) enter the Pairformer and are used to update each other. As the transition blocks have already been described, this section focuses on the Triangle Updates and Triangle Attention, then briefly explains how the Single Attention with Pair Bias differs from the variant described earlier. These triangle-based layers were first introduced in AF2 are one of the pieces that not only remained in AF3, but now are even more present in the architecture, so they get quite a bit of attention.

Why look at triangles?

The guiding principle here is the idea of the triangle inequality: “the sum of any two sides of a triangle is greater than or equal to the third side”. Recall that each z_i,j in the pair tensor encodes the relationship between positions i and j in the sequence. While it does not literally encode the physical distances between pairs of tokens, let’s think about it for a moment as if it did. If we imagine that each z_i,j is the distance between two amino acids and we know z_i,j=1 and z_j,k=1. By the triangle inequality z_i,k cannot be larger than 2. Knowing two of the distances gives us a strong belief about what the third distance must be. The goal of triangle updates and triangle attention are to try to encode these geometric constraints into the model.

The triangle inequality is not enforced in the model but rather, it is encouraged through ensuring each position z_i,j is updated by looking at all possible triplets of positions (i,j,k) at a time. So z_i,j is updated based on z_j,k and z_i,k for all other atoms k. Because z represents the complex physical relationship between these tokens, rather than merely their distance, these relationships can be directional. So for z_i,j, we also want to encourage consistency with z_k,i and z_k,j for all atoms k. If we think of the atoms as a graph, with z as a directed adjacency matrix, it makes sense that AlphaFold calls these “outgoing edges” and “incoming edges”.

Consider row i=0 of this adjacency matrix, and let’s say we want to update z_0,2, which has been highlighted in purple. The idea behind the update is that if we know the distances between 0→1 and 2→1, that gives us some constraints on what 0→2 can be. Similarly, if we know the distances between 0→3 and 2→3, this also gives us a constraint on 0→2. This would apply for all atoms k.

So, in the triangle updates and attention, we effectively look at all directed paths for 3 nodes in this graph (a.k.a triangles, hence the name!).

Triangle Updates

Having carefully looked at the triangle operations from a graph theory perspective, we can see how this is implemented with tensor operations. In the outgoing update, every position z_i,j in the pair representation gets updated independently based on a weighted combination of the other elements in the same row (z_i,j), where the weighting of each z_i,k is based on the third element in its outgoing edge triangle (z_j,k).

Practically, we take three linear projections of z (called a, b, and g). To update z_i,j, we take an element-wise multiplication of row i from a and row j from b. We then sum over all these rows (different values of k), and gate with our g projection.

For the incoming update, we effectively do the same thing but flipping the rows with the columns, so to update z_i,j we take a weighted sum of the other elements in the same column (z_k,j), where the weighting of each z_k,j is based on the third element in its outgoing edge triangle (z_k,i). After creating the same linear projections, we take an element-wise multiplication of column i from a and column j from b, and sum over all the rows of this matrix. You’ll find that these operations exactly mirror the graph-theory adjacency view described above.

Triangle Attention

After our two triangle update steps, we also update each z_i,j using triangle attention for the outgoing edges and triangle attention for the incoming edges. The AF3 paper refers to the “outgoing edges” as attention “around starting node” and “incoming edges” as attention “around ending node”.

To build up to triangle attention, it can be helpful to start with typical self-attention over a 1D sequence. Recall that queries, keys, and values are all transformations of the original 1D sequence. An attention variant called axial attention extends this to matrices by applying independent 1D self-attention over the different axes of a 2D matrix (the rows, then the columns). Triangle attention adds the triangle principle we discussed earlier to this, updating z_i,j by incorporating z_i,k and z_j,k for all atoms k. Specifically, in the “starting node” case, to calculate the attention scores along row i (to determine how much z_i,j should be influenced by z_i,k), we do a query-key comparison between z_i,j and z_i,k as usual, then bias the attention based on z_j,k as is shown above.

For the “ending node” case, we again swap rows for columns. For z_i,j, the keys and values will both come from column i of z, while the bias will come from column j. So, when comparing the query z_i,j with the key z_k,i, we bias that attention score based on z_k,j. Then, once we have attention scores over all k, we use our values vectors from column i.

Single Attention with Pair Bias

Now that we’ve updated our pair representation with these four triangle steps, we pass the pair representation through a Transition block as described above. Finally, we want to update our single representation (s) using this new updated pair representation (z), so we will use single attention with pair bias, pictured below. This is identical to Single Attention with Pair Bias describedFor reference, in the AF3 supplement, Single Attention with Pair Bias is also referred to as "Attention Pair Bias" in the Atom Transformer section, but at the token-level. As it operates on the token-level, it uses full attention as opposed to the block-wise sparse pattern used when operating at the atom-level.

We repeat the Pairformer for 48 blocks, eventually creating s^trunk and z^trunk.

3. Structure Prediction

Basics of Diffusion

Now, with these refined representations, we are ready to use s and z to predict the structure of our complex. One of the changes introduced in AF3 is that entire structure prediction is based on atom-level diffusion. Existing posts more thoroughly explain the intuition and math for diffusion, but the basic idea of a Diffusion Model is to start with real data, add random noise to your data, then train a model to predict what noise was added. Noise is iteratively added to the data over a series of T timesteps to create a sequence of T variants of each datapoint. We call the original data point x_t=0 and the fully noised version x_t=T. During training, at timestep t, the model is given the x_t and predicts what noise was added between x_t-1 and x_t. We take a gradient step on the predicted noise added compared to the actual noise that had been added.

Then, at inference time, we simply start with random noise, which is equivalent to x_t=T. For every time step, we predict the noise the model thinks has been added, and remove that predicted noise. After a pre-specified number of timesteps, we end up with a fully “denoised” datapoint that should resemble the original data from our dataset.

Conditional Diffusion lets the model ‘condition’ these de-noising predictions on some input. Practically this means that for each step of the model, it takes three inputs:

1. The current noisy iteration of our generation
2. A representation of the current time step we are at
3. The information we want to condition on (this could be a caption for an image to generate, or properties for a protein).

As a result, the final generation is not just a random example that resembles the training data distribution, but should specifically match the information represented by this conditioning vector.

With AF3, the data we learn to de-noise is a matrix x with the x,y,z coordinates of all the atoms in our sequences. During training, we add Gaussian noise to these coordinates until they are effectively fully random. Then at inference time, we start with random coordinates. At each time step, we first randomly rotate and translate our entire predicted complex. This data-augmentation teaches the model that any rotation and translation of our complex is equally valid, and replaces the much more complicated Invariant Point Attention used in AF2. AF2 had developed a complicated architecture called Invariant Point Attention meant to enforce equivariance to translations and rotations. This led to a vigorous debate over the importance of IPA in AF2's success. In AF3, this is dropped in favor of a much simpler approach: applying random rotations and translations as data-augmentations to help the model learn such equivariances naturally. So here we simply randomly rotate all atoms' coordinates around the center of our current generation (the mean over all atoms' coordinates), and randomly sample a translation in each dimension (x, y, and z) from a N(0,1) Gaussian. It appears from the algorithm that the translation is universal, that is the same translation is applied to every atom in our current generation. This type of data augmentation was popularize with CNNs but in the past few years, equivariant architectures like IPA have been considered an more efficient and elegant approach to solve the same problem. Thus, when AF3 replaced equivariant attention with data-augmentation, it sparked a lot of internet discussions. We then add a small amount of noise to the coordinates to encourage more heterogeneous generations.It benefits us for the model to generate several slightly different variations. At inference time, we can score each using our confidence head, and return only the generation with the highest score. Finally, we predict a de-noising step using the Diffusion Module. We cover this module in more detail below:

Data (coordinates) to get de-noised

Diffusion Module

See where this fits into the full architecture

In each de-noising diffusion step, we condition our prediction on multiple representations of the input sequences:

• the outputs of the trunk (our post-Pairformer updated s and z, now called s^trunk and z^trunk)
• the initial atom and token-level representations of the sequence created in the input embedder that have not gone through the trunk (s^inputs, c^inputs)

The AF3 paper breaks down its diffusion process into 4 steps that involve moving from tokens to atoms, back to tokens, and back to atoms:

1. Prepare token-level conditioning tensors
2. Prepare atom-level conditioning tensors, update them using the Atom Transformer, and aggregate them back to token-level
3. Apply attention at the token-level, and project back to atoms
4. Apply attention at the atom-level to predict atom-level noise updates

1. Prepare token-level conditioning tensors

To initialize our token-level conditioning representation, we concatenate z^trunk to the relative positional encodings then project this larger representation back down and pass it through several residual-connection transition blocks.

Similarly, for our token-level single representation, we concatenate the very first representation of the input created at the start of the model (s^inputs) and our current representation (s^trunk), then project it back down to its original size. We then create a Fourier embedding based on the current diffusion time stepMore specifically, the amount of noise associated with this timestep in the Noise Schedule, add that to our single representation, and pass that combination through several Transition blocks. By including the diffusion time step in the conditioning input here, it ensures the model is aware of the timestep in the diffusion process when making de-noising predictions, and so predicts the right scale of noise to remove for this timestep.

2. Prepare atom-level tensors, apply atom-level attention, and aggregate back to token-level

At this point, our conditioning vectors are storing information at a per-token level, but we want to also run attention at the atom-level. To address this, we take our initial atom-level representations of the input created in the Embedding section (c and p), and update them based on the current token-level representations, to create atom-level conditioning tensors.

Next, we scale the atom’s current coordinates (x) by the variance of the data, effectively creating “dimensionless” coordinates with unit variance (called r). We then update q based on r such that q is now aware of the atom’s current location. Finally, we update q with the Atom Transformer (which also takes the pair representation as input), and aggregate the atoms back to tokens as we’ve previously seen.Recall from the input preparation section that Atom Transformer runs sparse attention over the atoms, and all steps (layer norm, attention, gating) are conditioned on the conditioning tensor c.

At the end of this step, we return

• q: updated atom representation after incorporating information about the atom’s current coordinates
• a: token-level aggregated form of q, capturing coordinates and sequence information
• c: atom representation for conditioning based on the trunk
• p: our updated atom-pair representation for conditioning

3. Apply attention at the token-level

The goal of this step is to apply attention to update our token-level representation of the atom coordinates and sequence information, a. This step uses the Diffusion Transformer visualized during input preparation, which mirrors the Atom Transformer but for tokens.

4. Apply attention at the atom-level to predict atom-level noise updates

Now, we return to atom space. We use our updated a (token-level representations based on current “center atom” locations) to update q (atom-level representation of all atoms based on current location) using the Atom Transformer. As was done in step 3, we broadcast our tokens representation to match the number of atoms we started with (selectively duplicating the tokens that represent multiple atoms), and run the Atom Transformer. Most importantly, one last linear layer maps this atom-level representation q back to R³. This is the key step: we’ve used all these conditioning representations to generate coordinate updates r^update for all atoms. Now, because we generated these in the “dimensionless” space r_l, we carefully re-scaleThis careful scaling involves both the variance of our data, and the noise schedule based on our current timestep, so that our updates are smaller and smaller as we get deeper into the de-noising process. the updates from r^update to their form with non-unit variance, x^update, and apply the updates to x_l.

With that, we’ve completed our tour through the main architecture of AlphaFold 3! Now we provide some additional information about the loss function, auxiliary confidence heads, and training details.

4. Loss Function and Other Training Details

Loss function and confidence heads

L_loss = L_distogram * α_distogram + L_diffusion * α_diffusion + L_confidence * α_confidence

The loss is a weighted sum of 3 terms:

• L_distogram which evaluates the accuracy of the predicted distogram at a token-level
• L_diffusion which evaluates the accuracy of the predicted distogram at an atom-level. It looks at all pairwise distances then includes additional terms to prioritize distances between nearby atoms and atoms involved in protein-ligand bonds.
• L_confidence which evaluates the model’s self-awareness about which structures are likely to be inaccurate

L_distogram

The output of our model is atom-level coordinates, which can easily be used to create an atom-level distogramRecall how the distograms were initially created by binning pairwise distances between atoms. However, this loss evaluates a token-level distogram. To get the xyz coordinates for tokens, we just use the coordinate of the “center atom”. As these distogram distances are categorical, the predicted distogram is then compared to the true distogram via cross entropy.

L_diffusion

The diffusion loss itself is a weighted sum of three terms each computed over the atom positions, additionally scaled by the amount of noiset^{^}, the sampled noise level for the current time step, and σ_data, the variance of the data which scales the amount of noise at each time step added at the current time step:

L_diffusion = (L_MSE + L_bond * α_bond) * (t̂² + σ_data²)/(t̂+σ_data)² + L_{smooth_lddt}

• L_MSE is a version of the distogram loss we just discussed, but over all atoms rather just “center atoms” (and with DNA, RNA, and ligand atoms upweighted). Additionally, it looks at the mean squared error between positions, rather than binning them into a distogram.
• L_bond aims to ensure the accuracy of bond lengths for protein-ligand bonds by adding an additional MSE loss on the difference in predicted and ground-truth distograms for atom-pairs that are part of protein-ligand bonds.There are various stages of training and α_bond is set to 0 in the initial stages, so this term is only introduced later.
• L_{smooth_LDDT} (smoothed local distance difference test) is yet another variant of the distogram loss that tries to capture the accuracy of local distances. An atom-pair’s predicted distance “passes the test” if it is within a given threshold of the atom-pair’s true distance. To make this metric smooth and differentiable, we pass the difference between predicted and ground-truth distograms through a sigmoid centered on the test’s threshold. We can think of this as generating a probability (between 0 and 1) that this atom-pair passes the test. We take the average of four “tests’’ with increasingly tight thresholds (4, 2, 1, and .5 Å). Using this loss encourages the model to reduce the probability of failing each test. Finally, to make the test “local”, we ignore the loss for an atom-pair if that atom-pair’s ground truth distance is large, as we only want the model to focus on accurately predicting an atom’s distances to nearby atomsSpecifically, for an atom-pair l,m, we ignore the loss for l and m if l and m are more than 30 Å away if atom m is part of a nucleotide. We ignore the loss for l and m if they are more than 15 Å away from each other and m is not a nucleotide (so is part of a protein or ligand)..

L_confidence

The goal of this loss is not to improve the accuracy of the structure, but rather to teach the model to predict its own accuracy. This loss is a weighted sum of 4 terms that each correspond to a method of evaluating the quality of a predicted structure:

L_confidence = L_plDDT + L_PDE + L_resolved + L_PAE * α_PAE

• lDDT Atom-level “local distance difference test”, capturing the expected accuracy of an atom’s predicted distances to nearby atoms.

• PAE Predicted alignment error between token i’s predicted and the ground-truth positions. We first rotate and translate the predicted token i and ground-truth token i into the frame of token j. That is, if we assume for a moment token j is in exactly its ground-truth position, we predict how close token i is to where it should be, based on its relation to token j.

• PDE Predicted distance error between tokens, capturing the accuracy of predicted differences between all pairs of tokens.

• Experimentally resolved prediction The model predicts which atoms were experimentally resolved (not every atom is experimentally resolved in every crystal structure).

To get these confidence losses for each of the metrics, AF3 predicts values for these error metrics, then these error metrics are calculated on the predicted structure, and the loss is based on the difference between these two. So even if the structure is really incorrect and the PAE is high, if the predicted PAE is also high, the L_pae will be low.

These confidence predictions are generated mid-way through the diffusion process. At a selected diffusion step t, the predicted coordinates r^t are used to update the single and pair representations created in the representation learning trunk. The predicted errors are then calculated from linear projections of the updated pair representation (for PAE and PDE) or this updated single representation (pLDDT and experimentally resolved). Then, the actual error metrics are calculated based on the same generated atom coordinates (process described below, if interested) for comparison.

While these terms are included in the confidence head loss, gradients from these terms are only used to update the confidence prediction heads and do not affect the rest of the model.

Other Training Details

Now that the architecture is covered, the last pieces are some of the additional training details.

Recycling

As introduced in AF2, AF3 recycles its weights; that is, rather than making the model deeper, the model weights are re-used and inputs are run through the modules multiple times to continually improve the representations. Diffusion inherently uses recycling at inference time, as the model is trained to incorporate the timestep information and use the same model weights for every time step.

Cross-distillation

AF3 uses a mix of synthetic training data generated by itself (via self-distillation) but also by AF2, via cross-distillation. Specifically, the authors note that, by switching to the diffusion-based generative module, the model stopped producing the characteristic “spaghetti” regions that allowed users of AF2 to visually identify low-confidence and likely disordered regions. Just visually looking at the diffusion-based generations, all regions appeared equally high in confidence, making it more difficult to identify potential hallucinations.

To solve this problem, they included generations from AF2 and AF-Multimer in the training data for AF3, allowing the model to learn that, when AF2 was not confident in its prediction, it should output these unfolded regions and to “instruct” AF3 to do the same.Nucleic acids and small molecules in distillation datasets had to be removed as they could not be processed by AF2 and AF-multimer. However, once previous models generated new predicted structures, and these structures got aligned to the originals, the removed molecules were added back in. If adding these back in created new atom clashes, the whole structure was excluded, to avoid accidentally teaching the model to accept clashes.

(Diagram from AF3 paper)

Cropping and Training Stages

While no part of the model has an explicit restriction on the length of the input sequences, the memory and compute requirements increase significantly with sequence length (recall the multiple O(N_tokens³ operations)). Thus, for efficiency the proteins get randomly cropped. As introduced in AF-multimer, because we want to model the interactions between multiple chains, the random cropping needs to include all of these. They use 3 methods for cropping and all 3 of these are used in different proportions depending on the training data (ex: PDB crystal structure vs disordered PDB complex vs distillation, etc.)

• Contiguous cropping: Contiguous sequences of amino acids are selected for each chain
• Spatial cropping: Amino acids are selected based on distance to a reference atom (typically this atom is part of a specific chain or binding interface of interest)
• Spatial interface cropping: Similar to spatial cropping, but based on distances to atoms that specifically at a binding interface.

While a model trained on random crops of 384 can be applied to longer sequences, to improve the model’s ability to handle these sequences, it is iteratively fine-tuned on larger sequence lengths. The mix of datasets and other training details is also varied in each training stage as is shown in the table below.

(Table from AF3 supplement)

Clashing

The authors note that AF3’s loss does not include a clash penalty for overlapping atoms. While switching to a diffusion-based structure module means the model could in theory predict two atoms to be in the same location, this seems to be minimal after training. That said, AF3 does employ a clashing penalty when ranking generated structures.

Batch sizes

Although the diffusion process sounds quite involved, it is still significantly less computationally expensive than the trunk of the model. Thus, the AF3 authors found that it is more efficient from a training perspective to expand the batch size of the model after the trunk. So for each input structure, it gets run through the embedding and trunk, then 48 independent data-augmented versions of the structure are applied, and these 48 structures are all trained in parallel.

That’s it for the training process! There are some other small details but this is probably already more than you need, and if you’ve made it this far, the rest should be easy to pick up from reading the AF3 supplement.

ML Musings

Having walked so thoroughly through the architecture of AF3 and its comparisons to AF2, it is interesting how the choices made by the authors fit into broader Machine Learning trends.

AlphaFold as Retrieval-Augmented Generation

At the time AF2 was released, it was not common to include retrievals from the training set at inference time. In the case of AF, utilizing an MSA and template search. MSA-based methods were being used for protein modeling, but this type of retrieval was less used in other areas of Deep Learning (i.e., ResNets do not embed relevant training images at inference time when classifying a new image in Computer Vision, for example). Although AF3 reduces the emphasis on the MSA compared to AF2 (it is no longer operated on and updated in the 48 blocks of the Evoformer/Pairformer), they still incorporate both the MSA and templates, even as other protein prediction models such as ESMFold have dropped retrieval in favor of fully parametric inference.

Interestingly, some of the largest and most successful Deep Learning models now often include similar additional information at inference time. While the details of the retrieval systems are not always disclosed, Large Language Models routinely use Retrieval Augmented Generation systems such as a traditional web search at inference time to orient the model toward relevant information (even if that information was likely already in its training data) that should guide inference. It will be interesting to see how the use of directly relevant examples at inference time develops in the future.

Pair-Bias Attention

One of the major components of AF2 that is even more present in AF3 is Pair-Bias Attention. That is, attention where the queries, keys, and values all originate from the same source (like in self-attention), but where there is a bias term added to the attention map from another source. This effectively acts as a light-touch version of information sharing, without full cross-attention. Pair-Bias Attention appears in almost every module. While this type of attention is now used in other protein modeling architectures, we have not seen this particular type of cross-biasing used in other fields (although that does not mean it hasn’t been done!). Perhaps it only works well here because the pair-representation is naturally analogous to a self-attention map already, but is an intriguing alternative to pure self or pure cross-attention.

Self-supervised training

Self-supervised models like ESM have been able to achieve impressive results at predicting protein structure by replacing the MSA embedding with a “probabilistic MSA” using self-supervised pre-training. In AF2, the model had an additional task that predicted masked tokens from the MSA, achieving a similar self-supervision, but that was removed with AF3. We have not seen commentary from the authors on why they did not use any self-supervised language modeling pre-training approach on the MSA, and in fact decreased the compute used to process the MSA. Three possible reasons self-supervised learning is not used to initialize the MSA embeddings are 1) they viewed the massive pre-training phase as a suboptimal use of compute 2) they tried it and found that including a small MSA module outperformed pre-trained embeddings and was worth the additional inference-time cost or 3) utilizing a mix of pre-trained embeddings for amino acid tokens and randomly initialized embeddings for DNA/RNA/ligands would not be compatible or underperformed fully supervised training on their hybrid atom-token structure. By focusing on self-supervision tasks, models in the ESM family are also much more simple than AF3 (although they don't handle DNA/RNA/ligands and have slightly different goals.) It is interesting to watch that as some models aim to maximize architectural simplicity, AlphaFold remains this complicated!

Classification vs. Regression

As in AF2, AF3 continues to use a mix of MSE and binned classification losses. The classification components are interesting as, if the model predicts a distogram bin that is only off-by-one, it gets no “credit” for being close rather than way off. It is unclear what informed this design decision, but perhaps the authors found the gradients to be more stable than working with several different MSE losses, and perhaps the per-atom losses saw so many gradient steps that the additional signal from a continuous loss would not have proven beneficial.

Similarities to Recurrent Architectures (e.g LSTMs)

AF3’s architecture incorporates several design elements reminiscent of recurrent neural networks that are not typically found in traditional transformers:

• Extensive Gating: AF3 uses gating mechanisms throughout its architecture to control information flow in the residual stream. This is more akin to the gating in LSTMs or GRUs than the standard feed-forward nature of normal transformer layers.
• Iterative Processing with Weight Reuse: AF3 applies the same weights multiple times to progressively refine its predictions. This process, involving both recycling and the diffusion model, resembles how recurrent networks process sequential data over time steps using a shared set of weights. It differs from standard transformers, which typically make predictions in a single forward pass. This approach allows AF3 to iteratively improve its protein structure predictions without increasing the number of parameters.
• Adaptive Computation: The recycling is also similar to the iterative updating used in diffusion and quite related to the idea of adaptive compute time (ACT), originally introduced to dynamically determine how much compute to use for RNNs and more recently used in Mixture-of-Depths to achieve a similar goal with transformers. This contrasts with the fixed depth of standard transformers and theoretically would allow the model to apply more processing to challenging inputs.

It was shown in the AF2 ablations that the recycling was important, but there was little disucssion on the importance of gating. Presumably it helps with training stability as in LSTMs but it is interesting that it is so prevalent here yet not in many other transformer-based architectures.

Cross-distillation

The use of AF2 generations to re-introduce its distinctive style specifically for low-confidence regions is very interesting. If there is a lesson here, it may be the most practical of all: If your previous model is doing one specific thing better than your new model, you can try cross-distillation to get the best of both worlds!

If you’ve made it this far, thanks for reading and hope it was helpful!! If you have any questions / comments / corrections / feedback feel free to reach out to us on twitter (E, J) or email (E, J) !

Special thanks to Kristy Carpenter, Nicholas Joseph, Kyle Swanson, and Kara Liu for giving feedback on this 💜

Note: this was originally written for the Reverie Labs blog which got taken down after acquisition so now it’s re-posted here

This blog post discusses why it is important to visualize the latent space of chemical datasets, what makes UMAP a useful tool for this purpose, and how we use UMAP at Reverie Labs. As an example, we use the Blood Brain Barrier Permeability (BBBP) dataset from MoleculeNet for our visualizations and the code tutorial. This dataset has measurements from over 2000 unique compounds, many of which are approved drugs, each labeled as “permeable” or “not permeable”. We look into the details of how this dataset gets embedded by various dimensionality reduction methods and reveal some fascinating properties of UMAP.

For a walkthrough using this dataset of how to use UMAP to visualize chemical space, see this Colab notebook:

Motivation

A fundamental assumption behind most machine learning methods is that data are independent and identically distributed (IID). However, in drug discovery datasets, compounds are almost never sampled independently, as they are typically extracted from experiments for specific therapeutic programs. Measurements often follow the patterns of the drug development efforts that generate them. Any biases in the data-generation process can also sneak into the training and evaluation of models. In practice, this means that open source and industry datasets are often “clumpy”, consisting of measurements for compounds that are very similar to one another and non-uniformly cover chemical space. Visualizing the chemical space of a dataset does not solve these issues, but it helps us better understand them within the context of our datasets.

At Reverie Labs, we work with dozens of datasets spanning a range of different properties. When ingesting a new dataset, we begin with a systematized analysis that identifies biases and helps determine how best to prepare and use the dataset for modeling. As with most machine learning problems, we want to understand the distribution of the measured property we are modeling (What is the class balance within a classification dataset? Does our regression dataset have truncated measurements due to assay limits? etc.) We want to make sure the measurements look reasonable and compare their distributions in our training and test data. However, we cannot stop there. We must also look at the distribution of the chemical structures that the measurements are taken from because the compounds themselves are typically selected from a biased generation process.

Visualizing the distribution of our compounds in chemical space allows us to gauge how much we expect models trained on a given dataset to generalize to new chemical matter. These visualizations can help with manual inspection of Structure-Activity / Structure-Property Relationships (SAR/SPR), expose potential quirks or biases in the dataset, and reveal insights for how we might want to split the dataset into training and evaluation sets. These are many methods we can use for visualizing chemical space, but at Reverie, we have selected a default procedure involving UMAP that optimizes accuracy, speed, and ease-of-use.

Embeddings

That sounds great, but how do we actually visualize these distributions? The key here is that we need to embed our compounds into a low-dimensional vector form that can be easily interpreted (in this post we stick to 2D for simplicity’s sake but 3D would also work). Representing our molecules as 2048-bit Extended-Connectivity Fingerprints (ECFPs) gives us high dimensional vectors that we can then project into a 2-dimensional space for visualization. PCA and t-SNE are commonly used for this kind of dimensionality reduction, and have been used for many biology and chemistry purposes. Since usage of these tools has been thoroughly documented, here we focus on the utility of a more recent addition to the dimensionality reduction repertoire: UMAP.

Uniform Manifold Approximation and Projection (UMAP) constructs a high-dimensional graph representation of the entire dataset then tries to re-create a low dimensional version of this initial graph that maintains as much of the local and global structure as possible. This method is somewhat similar to t-SNE, but with a few key differences that lead to important advantages for our purposes. The technical differences between UMAP and t-SNE would take up a full blog post, so we will not detail them here but more details can be found in these resources: the original paper, Understanding UMAP or How Exactly UMAP Works.

The most relevant benefits UMAP provides us are speed, the ability to maintain some of the local / global structure of the data, and an easy interface for applying an embedding from one dataset to a different dataset.

Speed

To evaluate the speed of these methods, we compute ECFPs of various sizes and embed each using each PCA, t-SNE, and UMAP:

We can see that PCA is the most efficient, UMAP is slightly less efficient, and t-SNE is by far the least efficient. For small datasets, these time differences are not particularly relevant, but they really add up as the number of compounds grows. To be fair, we used the most generic implementations (PCA and t-SNE from sklearn) and there are more efficient variants of t-SNE. However, the UMAP Docs contain a similar performance analysis on MNIST, including a wider variety of the performant methods, and find similar results.

Additionally, the creators of UMAP recently released a new version of the algorithm, ParametricUMAP, that uses a neural network to reduce the dimensions of the graphical embedding, giving it even greater speed improvements. For this post we stick to the original, non-parametric UMAP, but if you want to optimize the speed of embedding new compounds into a pre-fit UMAP model, ParametricUMAP can be quite useful.

Local / Global Structure

If PCA is so fast, why don’t we just use that? Our goal is to understand both the local structure (organization of similar compounds) and global structure (organization of groups of different compounds) of complex datasets. Speed is important, but we also evaluate the methods according to how informative they are for those tasks. We can use the example BBBP dataset to explore the local and global structure of the embeddings created by the various methods:

Click these links for interactive versions of the plots [PCA, t-SNE, UMAP]

At a first glance, we notice some high-level differences in the overall structures of the embeddings. PCA looks like the intersection of two orthogonal lines, and the compounds within them are somewhat uniformly distributed. t-SNE has a flowery shape that resembles a 2D gaussian, with a variety of isolated clusters around the edges. UMAP has many disjoint, tight clusters that do not follow a specific pattern. Beyond surface-level observations, we cannot ascertain much more about these embeddings from these plots alone. To do that, we need to look more into the actual compounds that are represented in these images.

These links [PCA, t-SNE, UMAP] take you to pages with interactive views of these plots, which is how we typically examine our datasets. Exploring the data in an interactive form helps builds intuition for how the different algorithms organize the chemical space of the dataset. To help guide this interactive exploration, we’ve selected a few clusters from the UMAP plot and highlighted where each of their compounds are embedded in t-SNE and PCA:

We see that local areas of the embedding contain compounds that not only look similar, but also belong to the same drug class. The clusters we’ve selected here contain steroids, tetracycline antibiotics and β-lactam antibiotics. Diving into the details of these clusters and their respective drug types serves as a great case study through which we can better understand and compare the structure of the embeddings.

Case study

In the steroid cluster there are many compounds with the 4-ring system that is characteristic of steroids, and even some non-steroid compounds with a similar structure. Each method separates these compounds out from the rest, although UMAP appears to isolate them out and group them together the most strongly.

In a nearby but fully isolated cluster we find a collection of tetracycline antibiotics antibiotics. These compounds also contain a fused 4-ring system, however the tetracycline rings differ from the steroid rings due to their shape, relative positioning, and bond orders.

Comparing the embedding of the steroid and tetracycline clusters we see: 1) Steroids appear more isolated than the tetracyclines. If we assume the global distances between clusters are meaningful, this implies that the steroids are more unique from the rest of the dataset than the tetracycline are. 2) While the steroid compounds are spread out within their cluster, the tetracyclines are all embedded practically on top of each other. If we assume the spread of a cluster reflects the local diversity, this implies that the steroids are more diverse than the tetracyclines. 3) The two clusters are relatively near each other. If we assume the relationships between clusters are meaningful, this implies that these clusters are more similar to each-other than they are to the rest of the dataset.

To evaluate whether these assumptions about the global and local structure are true, we can examine the validity of the claims they imply. Should steroids be placed farther away from the rest of the compounds than the tetracyclines are? Both groups appear to stand out from the rest of the dataset, but it is difficult to manually compare their levels of uniqueness. Should tetracyclines and steroids actually be placed near each other? Given that both groups contain 4-fused ring systems, it seems reasonable for them to be located near each other but this similarity could just be superficial. To more quantitatively address these questions we look at the measured similarities between steroids, tetracyclines, and the rest of the compounds in the dataset. For each steroid compound, we calculate its average similarity to the other steroids, the tetracyclines, and every other compound. This is repeated with the tetracyclines. We define chemical similarity between compounds using Tanimoto similarity between ECFPs for the compounds, the same metric used to create the embeddings. The higher the Tanimoto similarity, the more similar to compounds are to one another.

These plots contextualize some of our observations about the steroid and tetracycline clusters:

1. The steroids and tetracyclines have roughly equivalent levels of similarity to the rest of the dataset (average similarity of 0.09 and 0.1 respectively). The distributions are not identical but their difference is not large enough to explain the extra isolation of the steroid cluster. This supports the general belief that these global distances are not always interpretable.

2. Tetracyclines are indeed more homogenous than steroids. Tetracyclines have, on average 0.15 higher tanimoto similarity with other tetracyclines than steroids do with other steroids. This means that the spread of intra-group distributions actually reflect the local chemical diversity and are not just a strange quirk of the UMAP embedding.

3. Tetracyclines and steroids have higher similarity to each other than they do to the rest of the dataset. This means that the visual relationship between these embedded clusters actually reflects a real relationship between the compounds that the clusters represent.

This reveals that the global structure of this dataset is not maintained through exact distances between groups of compounds, but rather the relationships between them. Local structure is expressed by maintaining the local diversity of groups through their intra-group distribution. There actually is an even more detailed level of local structure hidden in these embeddings that we’ll examine later, but given the information we have so far these are the main conclusions. The embedding of two groups of compounds does not prove anything about UMAP that we expect to hold up for all embeddings. But their examples highlight important patterns in how chemical datasets get embedded.

We can compare UMAP’s arrangement of the steroids and tetracyclines with PCA and t-SNE’s to look for differences in the local and global structures of these embeddings:

Based on these observations, PCA is not as successful at maintaining the structure of these groups within the dataset. Specifically, t-SNE and UMAP highlight the uniqueness and homogeneity of tetracyclines, whereas PCA spreads the tetracyclines out amidst various other scaffolds in an unidentifiable way. This again supports that, although PCA maintains a few key elements of the global structure, t-SNE and UMAP preserve the global and local structure more consistently throughout the dataset.

Differences between the embeddings are less noticeable when examining the steroids Each method embeds the steroids in a clearly identifiable, yet disperse cluster. UMAP’s steroid cluster is the most isolated but as discussed earlier, this extra separation is not particularly meaningful. Both PCA and UMAP place the chemically similar steroids and tetracyclines nearby each other while t-SNE does not. This seemingly implies t-SNE’s global structure is not as informative. However, t-SNE places the tetracyclines near the β-lactam antibiotics, which, as we will read below, actually makes sense.

The differences between the methods are most apparent when we examine the β-lactam antibiotics. We see the namesake β-lactam ring present in every compound and, in terms of global structure, all three methods separate out the β-lactam antibiotic compounds as significantly unique from the rest. Again, the placement in PCA is not as isolated as it is in the other methods but the cluster still stands out.

The β-lactam antibiotics are interesting because they give us a view into the level of local organization that t-SNE and UMAP have within the subclasses of this drug-type:

We have labelled each of the various subclasses of β-lactam antibiotics. They vary based on the particular details of the β-lactam ring system in a given compound. You don’t need to actually understand the differences between the various β-lactam subclasses but know they are fairly small. The importance of visualizing them is to highlight their placements in each of the embeddings.

PCA has all of the subtypes mixed together, which makes sense given that in PCA, the principal components we are visualizing are meant to have maximal variance. The nuances within a class of drugs are not particularly high variance so a 2-D PCA plot loses this local structure. On the other hand, t-SNE and UMAP both maintain local structures in the dataset by embedding the β-lactam antibiotics in a way that separates the compounds based on their subclass.

Zooming into the bottom of our UMAP plot where the β-lactam compounds are embedded, we can better examine the details of each subclass:

Not only are the β-lactam antibiotics contained in this section of the UMAP embedding, but the individual subclasses are each grouped together. If we look at the interactive versions of the plots (PCA, TSNE, UMAP), we see a similar phenomena in the t-SNE embedding. This ability to maintain both global structure and such specificity in the local structure is what makes these methods useful for easily exploring a dataset.

If we look even closer at the placement of the β-lactam antibiotics, we discover a fascinating property of our embedding. In the annotated plot above there are two main clusters of β-lactam antibiotics and one outlier compound, in the upper left corner, that has the substructure of a Penicillin (penam), yet appears to be located far away; in fact, it is placed in the tetracycline cluster.

The tetracycline antibiotics are structurally quite different from the β-lactam antibiotics, and yet they are embedded relatively nearby in both UMAP and t-SNE. We do not expect the distances between clusters to necessarily mean anything, but in this case, their relative positioning actually does.

The outlier compound, Penimocycline, contains both the substructure of the Penicillins (penam) and the substructure of the Tetracycline. It is actually classified as both types of antibiotics. When constructing a graph representation of this dataset, this compound is likely connected to both of these two well-connected subgraphs. This would link the two groups together, ultimately leading to their respective clusters of compounds being located near each-other in the final embedding.

To investigate this assumption, we can remove Penimocycline from our dataset and generate a new embedding. If this compound is functioning as a link between the tetracycline antibiotics and β-lactam antibiotics, embedding the dataset without it would break the connection between the groups and their positions would no longer be close to each other.

As hypothesized, in the new embedding on the right, we see that the tetracycline antibiotics are no longer placed near the β-lactam antibiotics. Tetracycline antibiotics are still near the steroids in each embedding variant, yet the induced separation between the tetracycline and β-lactam antibiotics has shifted much of the embedding. UMAP is non-deterministic so we re-ran these embeddings (both the original on the left and the modified dataset on the right) multiple times and this phenomena held up. This implies that the single compound actually is the influential node linking the β-lactam and tetracycline antibiotic clusters. It also reveals how the structure of UMAP is greatly influenced by individual compounds with strong connections between otherwise disconnected subgraphs of the dataset.

If you continue exploring the other areas of this dataset in the interactive links or Colab notebook provided you will also find collections of narcotics, sedatives, NSAIDs etc. Examining where specific compounds get placed in each of the embeddings helps explain the structural differences between the embeddings.

Hyperparameters

UMAP has several hyperparameters that give the user a bit more control over the structure of the final embedding based on their particular priorities:

• n_components is the dimensionality of the final embedding. To create visualizations in 2 dimensions we keep this fixed at 2 components.
• metric is the metric used to determine distance between points. Because we are comparing ECFPs, we use Jaccard distance (typically referred to as Tanimoto distance in cheminformatics).
• n_neighbors determines the prioritization of local versus global structure in the embedding. This value constrains the number of neighbors that a given compound has in the graph representation of the dataset. If n_neighbors is small then the embedding focuses on optimizing the distances between similar compounds to ensure the small differences between them are well represented. If n_neighbors is larger, then the distances between less similar compounds is prioritized.
• min_dist is the minimum distance between any two points. This affects the tightness of the embedding. The larger min_dist, the more spread out the compounds will be.

n_neighbors and min_dist can be tuned based on the dataset’s properties and the user’s preferences. The ideal settings may vary based on the dataset. These plots show how varying n_neighbors (along the rows) and min_dist (along the columns) can influence the embedding:

As these plots reveal, the spread of the embedding is quite dependent on the relationship between these two parameters. When min_dist is very small, compounds that are very similar to each-other are placed almost directly on top of each other, which makes it very easy to identify unique clusters. However, it is more difficult to decipher the actual quantity of compounds and the nuanced differences between them. When min_dist is large, it is easier to gauge the full spread of the compounds but more difficult to isolate specific clusters. When n_neighbors is small there are meaningful patterns within clusters but the global structure is less interpretable. As this value increases we see the relationships between the clusters becoming more noticeable as the clusters become less sparse. As with the t-SNE / UMAP comparison, there is no clear answer that one particular set of hyperparameters is always best. Depending on what the user is looking for, there are many great options.

Other Works

For a non-comprehensive list of examples of others using UMAP for chemistry and biology purposes see:

• Dimensionality reduction by UMAP to visualize physical and genetic interactions
• Wicked Fast Cheminformatics with NVIDIA
• One class classification as a practical approach for accelerating π–π co-crystal discovery
• Generative molecular design in low data regimes
• Focus Your Screening Library: Rapid Identification of Novel PDE2 Inhibitors with in silico Driven Library Prioritization and MicroScale Thermophoresis
• UMAP Visualization of SARS-CoV-2 Data in ChEMBL
• De novo design and Bioactivity Prediction of SARS-CoV-2 Main Protease Inhibitors using ULMFit

This year yet another alternative to t-SNE, TMAP has been developed. We haven’t extensively investigated that method yet, but it does seem promising.

Caveat

It is impossible to distill all of the complexities of chemical space into 2-dimensions and a lot of information gets lost in the process. Our low-dimensional embeddings can only be as good as their high-dimensional predecessors, the ECFPs. ECFPs are an imperfect, yet important, method for vectorizing molecules and using the distances between these sparse vectors as the basis for generating the underlying UMAP graphs means that our embeddings can only be as good as those distance metrics. Despite these flaws, we still find these plots to provide significant value to our design efforts.

Practical Uses for UMAP at Reverie Labs

We use UMAP in two main ways when examining a dataset:

1. Dataset Specific Embeddings: Examine the particular distribution of compounds within a specific dataset
2. Dataset-Agnostic Embeddings: Examine where the compounds of a dataset fit into our general embedding of global chemical space

We can also use these two embeddings as new lenses to view the distribution of measured properties and other physicochemical properties.

Dataset Specific Embeddings

To create a dataset-specific embedding, we fit a UMAP model on the molecules of the particular dataset we are investigating. All of the visualizations we have used so far are based on dataset-specific embeddings of our example BBBP dataset. As we saw when examining the local and global structure of the UMAP embedding, a dataset-specific plot provides a good understanding of the nuances within the dataset.

To add an extra dimension, we can color the plots based on any other relevant data we have on the compounds. Earlier we did this using the drug-types of certain compounds, but we can just as easily visualize the compounds colored by the date-of-synthesis, the measured property (in our case blood brain barrier permeability), physicochemical properties, or the dataset-split.

This can help answer questions such as: Are there clusters of compounds that have not been actively developed in years? Are all of the most potent compounds in one area of chemical space? If splitting the dataset for model training and evaluation, do certain splitting methods lead to any particular artifacts?

Measured Properties

Here we color the points on the plot by the measured property of our dataset: whether the compounds can permeate the blood brain barrier or not. Many of the areas in this space contain exclusively permeable or exclusively impermeable compounds. This reveals that, within this dataset, there are certain types of compounds that are consistently permeable or not and the general scaffold of the compound is sufficient to determine the permeability of many compounds.

If we were to use machine learning to model this dataset, we might want to ensure that individual homogeneously-labeled scaffolds are not split between the training and test sets as that could misleadingly inflate performance metrics.

Physicochemical properties

Using Molecular Weight as our example physicochemical property, we observe that UMAP groups the heaviest compounds all together as part of one cluster, which contains many large macrocyclic compounds that are all impermeable.

Visualizing the Molecular Weight (MW) of the permeable/impermeable compounds we see a similar phenomena in which the majority of the heavy compounds (MW > 500) are impermeable. This aligns with traditional assumptions that compounds heavier than 500 Da will struggle to permeate the blood brain barrier.

The UMAP plot is useful in quickly identifying where the compounds of a given molecular weight lie, what is the diversity of molecular weight within given clusters, and how these properties interplay with blood brain barrier permeability.

Dataset-Agnostic Embeddings

To create a dataset-agnostic embedding, we take advantage of the way that UMAP treats fitting and transforming a dataset as two separate steps. We start by fitting a UMAP model on a large in-house corpus of drug-like compounds that we consider to be representative of ‘drug-like chemical space’. Now, when we want to examine a new dataset, we load up this cached model and use it to transform the compounds of our dataset into this universal embedding space. With Dataset Specific embeddings, the dimensionality reduction method itself depends on the relationship between all the compounds in the dataset we seek to visualize. With Dataset-Agnostic embeddings, the dimensionality reduction method is treated as fixed so the location of a given compound is agnostic to the other compounds in the dataset.

This allows us to examine multiple datasets all with a consistent frame of reference. We can visualize if the data is diverse and covers a wide area in this drug-like space or if it is contained to a few specific areas.

Here we visualize both the original embedding of our global chemical space compounds used to fit the general UMAP model, and a Dataset-Agnostic embedding of the BBBP dataset created with this fixed model. The compounds of each are colored by their cluster assignment in this embedded space. To get these cluster assignments we have pre-fit a clustering model on the original global chemical space compounds in their UMAP embedding. We then use this pre-trained clustering model on new datasets to quickly determine which of the global UMAP clusters our new compounds fit into. We can even calculate the percentage of global clusters covered by our new dataset to create a quick, quantitative heuristic of chemical space covered by the dataset.

As we saw in our case study, distances in UMAP space aren’t necessarily meaningful and global structure can be greatly influenced by individual compounds. Thus, clusters based off of these embedded distances should be interpreted with caution. However, they can definitely help with the visual examination of the dataset, and can establish a quick, quantitative approximation of chemical space coverage. When combining these cluster heuristics with the visuals of the global embedding, we can start to understand the diversity of a given chemical dataset.

Final Thoughts

UMAP is useful because it is easy and quick to create local and global embeddings. This allows us to treat this analysis as a standard piece of our internal pipeline for cleaning, analyzing, and preparing new datasets for analysis and machine learning. Overall UMAP seems to be a great alternative to the more popular methods for dataset embedding, and it would be exciting to see more examples of groups using it for chemical data.

This post highlighted UMAP’s value for exploring the chemical space of datasets relevant to drug design. The t-SNE vs. UMAP debate is still greatly contested, but we don’t aim to use this example to prove that UMAP is fundamentally better than t-SNE (although many have argued this, some have refuted it, and others have even refuted the refutation!) Similarly, even though PCA is less useful for our particular goals, PCA also has many advantages over its nonlinear alternatives that make it very useful for other purposes. Ultimately, there are many great dimensionality reduction tools to choose from, but we hope that this post has helped to put UMAP on your map.