This is a Systematization paper (specifically a handbook chapter) with a strong secondary Method projection.

Its primary goal is to serve as a “users’ guide” to the Embedded-Atom Method (EAM). The text organizes existing knowledge:

• It traces the physical origins of EAM from Density Functional Theory (DFT) and Effective Medium Theory.
• It synthesizes “closely related methods” (Second Moment Approximation, Glue Model), showing they are mathematically equivalent or very similar to EAM.
• It provides a pedagogical, step-by-step methodology for fitting potentials to experimental data.

Motivation: Bridging the Gap Between DFT and Pair Potentials

The primary motivation is to bridge the gap between accurate, expensive electronic structure calculations and fast, inaccurate pair potentials.

• Computational Efficiency: First-principles methods scale as $O(N^3)$ or worse, limiting simulations to $<100$ atoms (in 1994). Pair potentials scale as $O(N)$ and fail to capture essential many-body physics of metals.
• Physical Accuracy: Simple pair potentials cannot accurately model metallic defects; they predict zero Cauchy pressure ($C_{12} - C_{44} = 0$) and equate vacancy formation energy to cohesive energy, both of which are incorrect for transition metals.
• Practical Utility: There was a need for a clear guide on how to construct and apply these potentials for large-scale simulations ($10^6+$ atoms) of fracture and defects.

Novelty: A Unified Framework and Robust Fitting Recipe

As a review chapter, the novelty lies in the synthesis and the specific, reproducible recipe for potential construction. Central to this synthesis is the core EAM energy functional:

$$E_{\text{tot}} = \sum_i \left( F(\bar{\rho}_i) + \frac{1}{2} \sum_{j \neq i} \phi(r_{ij}) \right)$$

where the total energy $E_{\text{tot}}$ depends on embedding an atom $i$ into a local background electron density $\bar{\rho}_i = \sum_{j \neq i} \rho(r_{ij})$, plus a repulsive pair interaction $\phi(r_{ij})$.

• Unified Framework: It explicitly maps the “Second Moment Approximation” (Tight Binding) and the “Glue Model” onto the fundamental EAM framework above, clarifying that they differ primarily in terminology or specific functional choices (e.g., square root embedding functions).
• Cross-Potential Fitting Recipe: It details a robust method for fitting alloy potentials (specifically Ni-Al-B) by using “transformation invariance”, scaling the density and shifting the embedding function to fit alloy properties without disturbing pure element fits.
• Specific Parameters: It publishes optimized potential parameters for Ni, Al, and B that accurately reproduce properties like the Boron interstitial preference in $\text{Ni}_3\text{Al}$.

Validation: Computational Benchmarks and Simulations

The “experiments” described are computational validations and simulations using the fitted Ni-Al-B potential:

1. Potential Fitting:

   • Pure elements (Ni, Al) were fitted to elastic constants, vacancy formation energies, and diatomic data. The Ni fit achieved $\chi_{\text{rms}} = 0.75%$ and Al achieved $\chi_{\text{rms}} = 3.85%$.
   • Boron was fitted using hypothetical crystal structures (fcc, bcc) calculated via LMTO (Linear Muffin-Tin Orbital) since experimental data for fcc B does not exist.

2. Molecular Statics (Validation):

   • Surface Relaxation: Demonstrated that EAM captures the oscillatory relaxation of atomic layers near a free surface, a many-body effect that pair potentials fail to capture.
   • Defect Energetics: Calculated formation energies for Boron interstitials in $\text{Ni}_3\text{Al}$. Found the 6Ni-octahedral site is most stable ($-4.59$ eV relative to an isolated B atom and unperturbed crystal), followed by the 4Ni-2Al octahedral site ($-3.65$ eV) and the 3Ni-1Al tetrahedral site ($-2.99$ eV), consistent with channeling experiments.

3. Molecular Dynamics (Application):

   • Grain Boundary (GB) Cleavage: Simulated the fracture of a (210) tilt grain boundary in $\text{Ni}_3\text{Al}$ at a strain rate of $5 \times 10^{10}$ s$^{-1}$.
   • Comparison: Compared pure $\text{Ni}_3\text{Al}$ boundaries vs. those doped with Boron and substitutional Nickel.

Key Outcomes: EAM Efficiency and Boron Strengthening

• EAM Efficiency: Confirmed that EAM scales linearly with atom count ($N$), requiring only 2-5 times the computational work of pair potentials.
• Boron Strengthening Mechanism: The simulations suggested that Boron segregates to grain boundaries and, specifically when co-segregated with Ni, significantly increases cohesion.
  • The maximum stress for the enriched boundary was approximately 22 GPa, compared to approximately 19 GPa for the clean boundary.
  • The B-doped boundary required approximately 44% more work to cleave than the undoped boundary.
  • The fracture mode shifted from cleaving along the GB to failure in the bulk.
• Grain Boundary Segregation: Molecular statics calculations found B interstitial energies at the GB as low as $-6.9$ eV, compared to $-4.59$ eV in the bulk, consistent with experimental observations of boron segregation to grain boundaries.
• Limitations: The author concludes that while EAM is excellent for metals, it lacks the angular dependence required for strongly covalent materials (like $\text{MoSi}_2$) or directional bonding.

Reproducibility Details

The chapter provides nearly all details required to implement the described potential from scratch.

Data

• Experimental/Reference Data: Used for fitting the cost function $\chi_{\text{rms}}$.
  • Pure Elements: Lattice constants ($a_0$), cohesive energy ($E_{\text{coh}}$), bulk modulus ($B$), elastic constants ($C_{11}, C_{12}, C_{44}$), vacancy formation energy ($E_{\text{vac}}^f$), and diatomic bond length/strength ($R_e, D_e$).
  • Alloys: Heat of solution and defect energies (APB, SISF) for $\text{Ni}_3\text{Al}$.
  • Hypothetical Data: LMTO first-principles data used for unobserved phases (e.g., fcc Boron, B2 NiB) to constrain the fit.

Algorithms

• Component Functions:
  • Pair Potential $\phi(r)$: Morse potential form: $$\phi(r) = D_M {1 - \exp[-\alpha_M(r - R_M)]}^2 - D_M$$
  • Density Function $\rho(r)$: Modified hydrogenic 4s orbital: $$\rho(r) = r^6(e^{-\beta r} + 2^9 e^{-2\beta r})$$
  • Embedding Function $F(\bar{\rho})$: Derived numerically to force the crystal energy to match the “Universal Energy Relation” (Rose et al.) as a function of lattice constant.
• Fitting Strategy:
  • Smooth Cutoff: A polynomial smoothing function ($h_{\text{smooth}}$) applied at $r_{\text{cut}}$ to ensure continuous derivatives.
  • Simplex Algorithm: Used to optimize parameters ($D_M, R_M, \alpha_M, \beta, r_{\text{cut}}$).
  • Alloy Invariance: Used transformations $F’(\rho) = F(\rho) + g\rho$ and $\rho’(r) = s\rho(r)$ to fit cross-potentials without altering pure-element properties.

Models

• Parameters: The text provides the exact optimized parameters for the Ni-Al-B potential in Table 2 (Pure elements) and Table 5 (Cross-potentials).
  • Example Ni parameters: $D_M=1.5335$ eV, $\alpha_M=1.7728$ Å$^{-1}$, $r_{\text{cut}}=4.7895$ Å.

Hardware

• 1994 Context: Mentions that simulations of $10^6$ atoms were possible on the “fastest computers available”.
• Scaling: Explicitly notes computational work scales as $O(N)$, roughly 2-5x slower than pair potentials.

Paper Information

Citation: Voter, A. F. (1994). Chapter 4: The Embedded-Atom Method. In Intermetallic Compounds: Vol. 1, Principles, edited by J. H. Westbrook and R. L. Fleischer. John Wiley & Sons Ltd.

Publication: Intermetallic Compounds: Vol. 1, Principles (1994)

@incollection{voterEmbeddedAtomMethod1994,
  title = {The Embedded-Atom Method},
  author = {Voter, Arthur F.},
  booktitle = {Intermetallic Compounds: Vol. 1, Principles},
  editor = {Westbrook, J. H. and Fleischer, R. L.},
  year = {1994},
  publisher = {John Wiley & Sons Ltd},
  pages = {77--90},
  chapter = {4}
}

Additional Resources:

• NIST Interatomic Potentials Repository (Modern repository often hosting EAM files)
• Original EAM Paper (1984)
• EAM Review (1993)

A natural question follows: “What if I use more samples? Won’t that make it better?”

Using more samples improves performance when paired with the correct objective function. Averaging the loss over $k$ samples yields minimal gains. Changing the objective itself is where the real gain comes from. This post explores the difference between a “multi-sample VAE” and the Importance Weighted Autoencoder (IWAE), a model that uses the same architecture as a VAE but is trained with a different objective that optimizes a tighter bound on the log-likelihood.

All ideas here are based on the fantastic paper: “Importance Weighted Autoencoders” by Burda, Grosse, and Salakhutdinov.

The Two Ways to Use $k$ Samples

Let’s say we have our encoder $q(h|x)$ and decoder $p(x,h)$. We decide to use $k=5$ samples instead of $k=1$. We have two main options for how to calculate our loss.

Option 1: The “Multi-Sample VAE” (The Naive Way)

This is the most straightforward idea. For each input $x$ in our batch:

1. Draw 5 samples ($h_1, …, h_5$) from $q(h|x)$.
2. Calculate the standard VAE $\mathcal{L}_1$ loss for each sample.
3. Average these 5 losses together.

This is an average of logs. As the IWAE paper shows experimentally, this approach gives you a more stable gradient, but the final performance (in terms of log-likelihood) is “only slightly” better. You’re paying a 5x computational cost for a marginal gain because you’re still optimizing the same “loose” $\mathcal{L}_1$ bound.

Option 2: The Importance Weighted Autoencoder (IWAE) (The Right Way)

The IWAE takes a different approach. For each input $x$:

1. Draw 5 samples ($h_1, …, h_5$) from $q(h|x)$.
2. Calculate an “importance weight” $w_i$ for each sample.
3. Average these 5 weights together.
4. Take the logarithm of that average.

This is a log of an average, and the difference matters.

The Math: Average-of-Logs vs. Log-of-Averages

Let’s make this concrete. The standard VAE $\mathcal{L}_1$ objective is:

$$ \mathcal{L}_1(x) = \mathbb{E} _{h\sim q(h|x)} \left[ \log \frac{p(x,h)}{q(h|x)} \right] $$

A multi-sample VAE simply gets a better estimate of this same value:

$$ \mathcal{L} _{\text{VAE}, k}(x) \approx \frac{1}{k} \sum _{i=1}^{k} \log w_i \quad \text{where} \quad w_i = \frac{p(x,h_i)}{q(h_i|x)} $$

The IWAE objective, $\mathcal{L}_k$, is fundamentally different:

$$ \mathcal{L} _k (x) = \mathbb{E} _{h_1..h_k \sim q(h|x)} \left[ \log \left( \frac{1}{k} \sum _{i=1}^{k} \frac{p(x,h_i)}{q(h_i|x)} \right) \right] $$

In practice, we estimate this with a single Monte Carlo sample (of $k$ latents):

$$ \mathcal{L} _k (x) \approx \log \left( \frac{1}{k} \sum _{i=1}^{k} w_i \right) $$

Because the logarithm is a concave function, Jensen’s Inequality tells us that the “log of an average” is always greater than or equal to the “average of logs.”

$$ \mathcal{L}_k(x) \ge \mathcal{L}_1(x) $$

This means the IWAE is optimizing a strictly tighter lower bound on the true log-likelihood of the data.

Why Does This “Log-of-Average” Matter?

This mathematical property provides two practical benefits.

1. Better Density Estimation

Because $\mathcal{L}_k$ is a tighter bound on the true $p(x)$, optimizing it pushes the model to learn a much better generative distribution. The paper shows that IWAEs achieve “significantly higher log-likelihoods” than VAEs.

2. Richer Latent Representations

This is the most interesting part. The standard VAE $\mathcal{L}_1$ objective “harshly penalizes” the model if its one sample $h$ is a poor explanation for $x$. This pressure forces the recognition network $q(h|x)$ to be “overly simplified” to avoid bad samples, which can lead to many latent dimensions becoming inactive (the paper reports the number of “active units” per model).

The IWAE objective is more flexible. It only needs one of the $k$ samples to be good. This “increased flexibility” allows the model to learn far more complex posterior distributions and “richer latent space representations.” The paper’s experiments confirm this, showing IWAEs learn to use many more “active units” in their latent space.

What This Looks Like in Code (PyTorch)

The implementation difference makes this crystal clear.

First, the “k-sample” trick: for a batch x of shape [B, D] and k=5 samples, we repeat x to get x_repeated of shape [B*k, D]. We do all our forward passes on this large tensor.

VAE (Multi-Sample, k > 1) Loss

Here, we can still use the analytical KL divergence, which is a big simplification.

# x_repeated has shape [B*k, 784]
# mu, logvar have shape [B*k, latent_dim]
# recon_x has shape [B*k, 784]

# recon_loss_all shape: [B*k]
recon_loss_all = F.binary_cross_entropy(recon_x, x_repeated, reduction='none').sum(dim=1)

# kl_loss_all shape: [B*k]
# We use the simple, analytical KL term!
kl_loss_all = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)

# total_loss_all shape: [B*k]
total_loss_all = recon_loss_all + kl_loss_all

# --- The Key Step ---
# Just average all B*k losses. This is the "average of logs".
loss = total_loss_all.mean()

IWAE (k > 1) Loss

Here, we must compute the exact log-probabilities of the specific samples we drew.

# Helper function to compute log-prob of a sample from a Gaussian
def log_prob_gaussian(sample, mu, logvar):
    const = -0.5 * sample.shape[-1] * torch.log(2 * torch.tensor(math.pi))
    log_det = -0.5 * torch.sum(logvar, dim=-1)
    log_exp = -0.5 * torch.sum((sample - mu)**2 / torch.exp(logvar), dim=-1)
    return const + log_det + log_exp

# --- Get the 3 log-prob components ---
# x_repeated, recon_x, z_samples, mu_repeated, logvar_repeated
# all have a first dimension of [B*k]

# 1. log p(x|h_i): Log-Reconstruction Probability
# log_p_x_given_h shape: [B*k]
log_p_x_given_h = -F.binary_cross_entropy(recon_x, x_repeated, reduction='none').sum(dim=1)

# 2. log p(h_i): Log-Prior Probability (under N(0, I))
# log_p_h shape: [B*k]
log_p_h = log_prob_gaussian(z_samples, 0.0, 0.0) # mu=0, logvar=0

# 3. log q(h_i|x): Log-Encoder Probability
# log_q_h_given_x shape: [B*k]
log_q_h_given_x = log_prob_gaussian(z_samples, mu_repeated, logvar_repeated)

# --- The Key Step ---
# Combine to get the log-importance-weight
# log_w shape: [B*k]
log_w = log_p_x_given_h + log_p_h - log_q_h_given_x

# Reshape to [B, k] to group samples by their original input
log_w_matrix = log_w.view(B, k) # B is original batch size

# --- Apply the IWAE Objective (Log-Sum-Exp Trick) ---
# This is the "log of the average"
# log( (1/k) * sum(exp(log_w)) ) = logsumexp(log_w) - log(k)
log_iwae_bound_per_x = torch.logsumexp(log_w_matrix, dim=1) - math.log(k)

# The objective is to MAXIMIZE this bound, so the loss is its negative
loss = -log_iwae_bound_per_x.mean()

The Critical Implementation Detail

Notice the key difference in the final step:

• VAE: loss = total_loss_all.mean() average of individual losses
• IWAE: loss = -torch.logsumexp(log_w_matrix, dim=1).mean() log of averaged weights

This seemingly small change implements the fundamental mathematical difference between optimizing an “average of logs” versus a “log of averages.”

When to Use Each Approach

The Computational Trade-off

Both approaches scale linearly with $k$. If you use $k=5$ samples, you’re doing roughly 5x the computation. The question is whether you get 5x the benefit.

For multi-sample VAEs, the answer is usually “no”. You get more stable gradients but only marginal performance improvements.

For IWAEs, the answer is often “yes”. You get meaningfully better log-likelihoods and richer latent representations that can be worth the computational cost.

Conclusion

The next time you use more samples with your VAE, switch to the IWAE objective to get the full benefit of the computational cost of $k > 1$.

The mathematical insight is simple but powerful: Jensen’s Inequality tells us that the “log of an average” is always greater than or equal to the “average of logs.” By optimizing this tighter bound, IWAEs achieve better density estimation and learn richer latent representations than standard VAEs.

The implementation requires computing exact log-probabilities to evaluate the specific samples. The result is a fundamentally more powerful model using the exact same architecture.

Want to dive deeper? Check out the original IWAE paper for experimental results and theoretical analysis, or explore my VAE tutorial for hands-on implementation details.

Decades of chemical research, breakthroughs in medicine, and novel materials are archived in journals, patents, and textbooks. A huge portion of this knowledge is stored as images, a format inaccessible to standard computational tools. This imposes challenges for both data retrieval and leveraging modern computational tools to analyze and predict chemical properties, inefficiencies that compound across the literature: knowledge locked in image form is invisible to search, mining, and downstream model training.

This is the central challenge that Optical Chemical Structure Recognition (OCSR) aims to solve. At its heart, OCSR is to chemistry what OCR (Optical Character Recognition) is to text: a technology that teaches computers to extract chemical information directly from 2D diagrams of molecules. It’s the bridge between a picture of a molecule and a machine-readable format like SMILES (Simplified Molecular Input Line Entry System) that can be stored, searched, and used to power new discoveries.

Teaching a computer to read a chemical structure requires specialized techniques.

The Complexity of Chemical Graphs

Recognizing a molecule requires specialized techniques that extend standard Optical Character Recognition (OCR). A molecule is a graph: a collection of atoms (nodes) connected by bonds (edges).

(While this simplified view excludes complex structures like coordination compounds and polymers, it provides a highly effective starting point for this discussion.)

An OCSR system must overcome several hurdles:

• Varying Styles: Chemical drawings vary widely across publications. Bond lengths, angles, and fonts can differ dramatically from one document to another.

• Image Quality: Older documents might be scanned at low resolutions, containing noise, blur, or other artifacts that make interpretation difficult.

• Structural Complexity: From simple rings to sprawling polymers and complex Markush structures (common in patents to represent a whole family of related compounds), the variety is immense.

The Evolution of OCSR

The quest to automate this process has evolved significantly, moving from brittle, hand-coded systems to sophisticated AI that can learn from data.

Act 1: The Rule-Based Pioneers (OCR-1.0)

The first OCSR systems, developed in the early 1990s, represent what we can now call the “OCR-1.0” era. Tools like Kekulé, and later open-source solutions like OSRA and MolVec, operated like meticulous draftsmen. Their approach was methodical:

1. Vectorize the Image: Convert the pixel-based image into a collection of lines and shapes
2. Identify Components: Use a set of hard-coded rules to classify these components. “This thick line is a wedge bond.” “This group of pixels is the letter ‘O’.”
3. Reconstruct the Graph: Piece together the identified atoms and bonds into a coherent molecular graph

This rule-based approach was a real first step but brittle. It struggled with the messiness of real-world documents and was expensive to maintain because each new style or error required new rules.

Additionally, they were designed as interactive tools to assist human experts in digitizing chemical structures. There was always the assumption that a human would review and correct the output.

As a concrete case-study, consider the (reproduced) results from MolParser:

Table 2. Comparison of our method with existing OCSR models. We report the accuracy. We use bold to indicate the best performance and underline to denote the second-best performance. *: re-implemented results. †: results from original publications.

In this table, we see that the rule-based methods (OSRA, MolVec, Imago) perform reasonably well on cleaner datasets like USPTO and UoB but falter on more challenging ones like JPO and ColoredBG. Modern AI-based methods (MolGrapher, DECIMER, MolScribe, MolParser) improve most on the hardest benchmarks (like JPO and ColoredBG), especially when fine-tuned on real data, while the rule-based tools still do reasonably well on cleaner sets like USPTO and UoB.

Act 2: The AI Fork in the Road (2010s-2020s)

The rise of deep learning in the 2010s brought new paradigms that could learn from data. Here, the field split into two distinct paths.

Path A: The Rise of the Specialists (Graph-Based AI)

Some models replaced the hard-coded rules with AI components. Systems like MolGrapher and MolScribe use a two-stage process:

• Atom Detection: A neural network first identifies all the atoms in the image
• Bond Prediction: A second process then predicts the connections (bonds) between those atoms to form the final graph

These are highly specialized tools, trained specifically for the task of building a molecular graph.

Path B: The Rise of the Generalists (LVLMs)

Another, more direct method treats OCSR as an image captioning task. This approach aligns with the broader trend of Large Vision-Language Models (LVLMs): massive, general-purpose AIs like GPT-4V. Models like DECIMER and MolParser look at a molecular image and directly generate its textual representation, most commonly a SMILES string. This direct, end-to-end approach is powerful, though it requires enormous datasets to train effectively.

The Next Frontier: The OCR-2.0 Vision (2024+)

Recently, a proposal has emerged that charts a third path forward: OCR-2.0. This vision, proposed by Wei et al. in 2024, argues for a new class of models that combine the best of both worlds. An OCR-2.0 model should be:

1. End-to-End: A single, unified model that simplifies maintenance
2. Efficient & Low-Cost: A specialized, highly efficient perception engine. The paper argues that using a giant LVLM for a pure recognition task is often inefficient
3. Versatile: Capable of handling diverse artificial optical signals

The flagship model for this theory is GOT (General OCR Theory). It’s a single, unified model that can read an image and output structured text for a wide variety of inputs. It can translate a molecular diagram into a SMILES string, transcribe sheet music into musical notation, parse a bar chart into a data table, and describe a geometric shape using code.

This demonstrates that OCSR can be integrated into broader systems for processing human visual information. The same OCR-2.0 philosophy extends beyond chemistry: GutenOCR, for instance, applies grounded vision-language modeling to general document OCR, producing both text transcriptions and bounding-box outputs from a single model.

Pushing the Boundaries of Recognition

OCR-2.0 models like GOT push for breadth, and other state-of-the-art research deepens the depth of understanding for the uniquely complex task of chemical recognition.

Deepening Reasoning with a “Visual Chain of Thought”

The GTR-Mol-VLM model makes recognition more intelligent by mimicking how a person might analyze a complex diagram. The model traverses the molecule step-by-step, predicting an atom, then its bond, then the next atom, and so on. This “Visual Chain of Thought” improves accuracy, especially for complex molecules. It also faithfully recognizes abbreviations like “Ph” as single units, better representing the source image.

Deepening Application with “Visual Fingerprinting”

Subgrapher rethinks the end goal. Many applications (like searching a patent database) require only the identification of specific molecular features. Subgrapher detects key functional groups and backbones directly from the image and creates a visual fingerprint. This approach mirrors identifying a person by key features (“has glasses, a mustache”), making it well-suited to finding matches in a large set.

Why It Matters

The evolution of OCSR directly enables practical scientific advancements. This technology is a critical enabler for the future of science.

Searching Past Knowledge

OCSR digitizes decades of research from patents and journals, making it searchable and accessible for data mining. Imagine being able to search through every molecule ever published with a simple query. Or consider the practical impact: pharmaceutical companies can now automatically scan thousands of patent documents to ensure their new drug candidates don’t infringe existing intellectual property, a process that previously required substantial manual review by patent analysts.

Accelerating Drug Discovery

By extracting vast datasets of molecules, scientists can train AI models to predict drug efficacy and toxicity, speeding up the discovery pipeline. The more molecular data we can digitize, the better our predictive models become.

Building Universal Document Intelligence

OCSR contributes to building AI systems capable of processing complex human documents. A scientific paper is a mix of text, equations, charts, tables, and molecular diagrams. Unified OCR-2.0 models are the key to making all of this knowledge searchable holistically.

Looking Forward

The goal is a loop where scientific knowledge, regardless of how it is stored, can be fed back into systems that read, search, and reason over it.

From the rule-based systems of the 1990s to today’s models that read many printed diagrams reliably (though hard cases like low-quality scans and Markush structures remain open), OCSR has improved a great deal. As accuracy, efficiency, and breadth improve, more of the chemical literature becomes machine-readable.

This entire process begins with teaching a computer how to read a picture.

I ran into this recently while debugging a generative model. Sometimes the grammar of the string provides the clue as to what is going wrong. Other times, actually seeing the molecule is what helps. I had a terminal full of generated strings and needed to verify their structures visually. I needed a streamlined way to generate these images locally. A lightweight script turns that text into a properly formatted image directly from the terminal.

SMILES vs. SELFIES

There are two primary string representations you will encounter in modern cheminformatics:

1. SMILES: The industry standard. It uses simple rules (C for carbon, = for double bonds, parentheses for branches). It is compact and machine-parseable. However, random SMILES strings are often invalid (e.g., unclosed rings or invalid valences).
2. SELFIES: Designed specifically for machine learning. It is a robust representation where every string corresponds to a valid molecular graph. This makes it ideal for generative models. Note that it is more verbose than SMILES.

I often need to visualize both formats. Let’s build a single, robust Python tool to handle them.

The Quick Win: Native RDKit

If you just need a quick image from a SMILES string and don’t care about the image dimensions or adding a legend, RDKit can do this in three lines:

from rdkit import Chem
from rdkit.Chem import Draw

mol = Chem.MolFromSmiles("CCO")
Draw.MolToFile(mol, "ethanol.png")

The native RDKit method is fast for quick checks. However, custom rendering provides necessary control over image dimensions, formula subscripts, and handling multiple input formats like SELFIES.

Building a Custom Renderer for Precise Control

Let’s build a fuller tool using RDKit for chemical processing, the selfies library for decoding, and PIL for image manipulation.

Core Dependencies

import selfies as sf
from rdkit import Chem
from rdkit.Chem import Draw, rdDepictor, rdMolDescriptors
from PIL import Image, ImageDraw, ImageFont

RDKit handles the chemical logic, selfies translates SELFIES to SMILES, and PIL gives us fine control over the final image appearance.

The Main Conversion Function

Here is the core conversion logic. Notice the Python type hints on the signature.

def string_to_png(mol_string: str, output_file: str, size: int = 500, is_selfies: bool = False) -> None:
    """Generates a 2D molecule image with a chemical formula legend from SMILES or SELFIES."""

    # Decode SELFIES to SMILES if necessary
    if is_selfies:
        try:
            smiles = sf.decoder(mol_string)
        except Exception as e:
            raise ValueError(f"Invalid SELFIES string: {mol_string}") from e
    else:
        smiles = mol_string

    mol = Chem.MolFromSmiles(smiles)
    if not mol:
        raise ValueError(f"Could not generate molecule from SMILES: {smiles}")

    # Generate 2D coordinates and formula
    rdDepictor.Compute2DCoords(mol)
    formula = rdMolDescriptors.CalcMolFormula(mol)

    # Render the molecule
    img = Draw.MolToImage(mol, size=(size, size)).convert("RGBA")

    # Create a canvas with extra space at the bottom for the legend
    legend_height = int(size * 0.1)
    canvas = Image.new("RGBA", (size, size + legend_height), "white")
    canvas.paste(img, (0, 0))

    draw = ImageDraw.Draw(canvas)

    # Define dynamic font sizes
    font_reg = get_font(int(size * 0.03))
    font_sub = get_font(int(size * 0.02))

    # Draw the legend
    x = int(size * 0.02)
    y = size + int(size * 0.02)

    # Draw "Formula: " label
    draw.text((x, y), "Formula: ", fill="black", font=font_reg)
    x += draw.textlength("Formula: ", font=font_reg)

    # Draw formula with subscript handling for numbers
    for char in formula:
        # Use smaller font and lower y-offset for numbers (subscripts)
        font = font_sub if char.isdigit() else font_reg
        y_offset = int(size * 0.005) if char.isdigit() else 0

        draw.text((x, y + y_offset), char, fill="black", font=font)
        x += draw.textlength(char, font=font)

    # Draw original string
    label = "SELFIES" if is_selfies else "SMILES"
    draw.text((x, y), f" | {label}: {mol_string}", fill="black", font=font_reg)

    canvas.save(output_file)
    print(f"Saved: {output_file}")

This function handles everything: SELFIES decoding, validation, coordinate generation, image creation, and legend drawing.

Font Handling

We need a helper to handle fonts robustly across systems:

def get_font(size: int, font_name: str = "arial.ttf"):
    """Attempts to load a TTF font, falls back to default if unavailable."""
    try:
        return ImageFont.truetype(font_name, size)
    except IOError:
        return ImageFont.load_default()

Examples in Action

Let’s see the tool in action with some common molecules, comparing the SMILES and SELFIES inputs.

Simple Molecules

Aromatic Compounds

Complex Pharmaceuticals

Going Further: Vector Graphics (SVG)

Use vector graphics (SVG/PDF) for true publication-quality figures. Vector graphics scale infinitely without pixelation.

RDKit handles this natively with rdMolDraw2D:

from rdkit import Chem
from rdkit.Chem.Draw import rdMolDraw2D

def string_to_svg(mol_string: str, output_file: str, size: int = 500, is_selfies: bool = False) -> None:
    """Generates a 2D molecule SVG image."""
    if is_selfies:
        try:
            mol_string = sf.decoder(mol_string)
        except Exception as e:
            raise ValueError(f"Invalid SELFIES string: {mol_string}") from e

    mol = Chem.MolFromSmiles(mol_string)
    if not mol:
        raise ValueError(f"Invalid string: {mol_string}")

    rdDepictor.Compute2DCoords(mol)

    d = rdMolDraw2D.MolDraw2DSVG(size, size)
    d.DrawMolecule(mol)
    d.FinishDrawing()

    with open(output_file, "w") as f:
        f.write(d.GetDrawingText())
    print(f"Saved: {output_file}")

This provides a perfect vector image. Note that this method omits the custom PIL-based legend. Choose the right tool for the job: PNG for quick checks and slides, SVG for journal submissions.

Command-Line Interface

The tool uses Python’s standard argparse library for the command-line interface. It automatically detects if you want an SVG based on the file extension and includes a --selfies flag.

# Basic SMILES usage
python mol2img.py "CCO" -o ethanol.png

# SELFIES usage
python mol2img.py "[C][C][O]" -o ethanol.png --selfies

# Generate SVG for publication
python mol2img.py "CCO" -o ethanol.svg

Download the Complete Script

You can copy the complete mol2img.py script directly from the code block below. For a fuller version with an SVG fallback, type hints, and batch (grid) rendering, see the Molecular String Renderer project.

Installation and Setup

Before using the script, install the required dependencies:

pip install rdkit pillow selfies

Complete Script

import argparse
import sys
import os
import selfies as sf
from rdkit import Chem
from rdkit.Chem import Draw, rdDepictor, rdMolDescriptors
from rdkit.Chem.Draw import rdMolDraw2D
from PIL import Image, ImageDraw, ImageFont

def get_font(size: int, font_name: str = "arial.ttf"):
    """Attempts to load a TTF font, falls back to default if unavailable."""
    try:
        return ImageFont.truetype(font_name, size)
    except IOError:
        return ImageFont.load_default()

def string_to_svg(mol_string: str, output_file: str, size: int = 500, is_selfies: bool = False) -> None:
    """Generates a 2D molecule SVG image."""
    if is_selfies:
        try:
            mol_string = sf.decoder(mol_string)
        except Exception as e:
            raise ValueError(f"Invalid SELFIES string: {mol_string}") from e

    mol = Chem.MolFromSmiles(mol_string)
    if not mol:
        raise ValueError(f"Invalid string: {mol_string}")

    rdDepictor.Compute2DCoords(mol)

    d = rdMolDraw2D.MolDraw2DSVG(size, size)
    d.DrawMolecule(mol)
    d.FinishDrawing()

    with open(output_file, "w") as f:
        f.write(d.GetDrawingText())
    print(f"Saved: {output_file}")

def string_to_png(mol_string: str, output_file: str, size: int = 500, is_selfies: bool = False) -> None:
    """Generates a 2D molecule image with a chemical formula legend."""
    if is_selfies:
        try:
            smiles = sf.decoder(mol_string)
        except Exception as e:
            raise ValueError(f"Invalid SELFIES string: {mol_string}") from e
    else:
        smiles = mol_string

    mol = Chem.MolFromSmiles(smiles)
    if not mol:
        raise ValueError(f"Could not generate molecule from string: {mol_string}")

    # Generate 2D coordinates and formula
    rdDepictor.Compute2DCoords(mol)
    formula = rdMolDescriptors.CalcMolFormula(mol)

    # Render the molecule
    img = Draw.MolToImage(mol, size=(size, size)).convert("RGBA")

    # Create a canvas with extra space at the bottom for the legend
    legend_height = int(size * 0.1)
    canvas = Image.new("RGBA", (size, size + legend_height), "white")
    canvas.paste(img, (0, 0))

    draw = ImageDraw.Draw(canvas)

    # Define dynamic font sizes
    font_reg = get_font(int(size * 0.03))
    font_sub = get_font(int(size * 0.02))

    # Draw the legend
    x = int(size * 0.02)
    y = size + int(size * 0.02)

    # Draw "Formula: " label
    draw.text((x, y), "Formula: ", fill="black", font=font_reg)
    x += draw.textlength("Formula: ", font=font_reg)

    # Draw formula with subscript handling for numbers
    for char in formula:
        # Use smaller font and lower y-offset for numbers (subscripts)
        font = font_sub if char.isdigit() else font_reg
        y_offset = int(size * 0.005) if char.isdigit() else 0

        draw.text((x, y + y_offset), char, fill="black", font=font)
        x += draw.textlength(char, font=font)

    # Draw original string
    label = "SELFIES" if is_selfies else "SMILES"
    draw.text((x, y), f" | {label}: {mol_string}", fill="black", font=font_reg)

    canvas.save(output_file)
    print(f"Saved: {output_file}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert a SMILES or SELFIES string to a 2D molecular image.")
    parser.add_argument("string", help="The molecular string to convert")
    parser.add_argument("-o", "--output", default="molecule.png", help="Output filename (default: molecule.png)")
    parser.add_argument("--size", type=int, default=500, help="Image width/height in pixels (default: 500)")
    parser.add_argument("--svg", action="store_true", help="Force SVG output (overrides filename extension)")
    parser.add_argument("--selfies", action="store_true", help="Treat the input string as SELFIES.")

    args = parser.parse_args()

    try:
        # Determine format based on flag or file extension
        is_svg = args.svg or args.output.lower().endswith(".svg")

        if is_svg:
            # Ensure extension is correct if not present
            if not args.output.lower().endswith(".svg"):
                args.output = os.path.splitext(args.output)[0] + ".svg"
            string_to_svg(args.string, args.output, args.size, args.selfies)
        else:
            string_to_png(args.string, args.output, args.size, args.selfies)

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

In the early days of computing, generating random numbers was a significant computational challenge. In a landmark 1951 paper, mathematician John von Neumann detailed various “cooking recipes” for producing and using random numbers on machines like the ENIAC. While much of the paper focuses on generating uniform random digits, he also described ingenious methods for generating numbers from more complex, non-uniform probability distributions.

One of the most fundamental needs in scientific simulation (from modeling radioactive decay to calculating particle free-paths in molecular dynamics) is sampling from an exponential distribution with probability density function:

$$f(x) = e^{-x} \quad \text{for } x \ge 0$$

Today’s standard approach is elegant and direct, but it requires computing a natural logarithm (a computationally expensive operation on early hardware). To sidestep this limitation, von Neumann described a fascinating alternative that uses only basic comparisons, resembling what he called “a well known game of chance Twenty-One, or Black Jack.”

In this post, we’ll explore both methods: the modern inverse transform approach and von Neumann’s ingenious comparison-based algorithm. We’ll implement them in Python, verify their correctness, and compare their performance, empirically testing the trade-offs von Neumann identified nearly 75 years ago.

Method 1: The Standard Approach (Inverse Transform Sampling)

The most common method for sampling from a given distribution is inverse transform sampling. This method relies on a fundamental principle: if you have a uniform random variable $U$ on the interval (0, 1), you can transform it into a random variable $X$ with any desired cumulative distribution function (CDF) $F(x)$ by applying:

$$X = F^{-1}(U)$$

For the exponential distribution, the CDF is $F(x) = 1 - e^{-x}$. To find the inverse, we set $U = 1 - e^{-X}$ and solve for $X$:

$$ \begin{align} e^{-X} &= 1 - U \ -X &= \ln(1 - U) \ X &= -\ln(1 - U) \end{align} $$

Here’s a useful simplification: since $U$ is uniformly distributed on (0, 1), the quantity $(1 - U)$ is also uniformly distributed on (0, 1). Therefore, we can use the simpler formula:

$$X = -\ln(U)$$

This gives us an efficient method for generating exponentially distributed numbers, provided the logarithm function is computationally accessible.

Python Implementation

Here’s a straightforward implementation using NumPy:

import numpy as np

def exponential_inverse_transform(n_samples=1):
    """
    Generate samples from an exponential distribution using inverse transform sampling.

    Args:
        n_samples (int): Number of samples to generate.

    Returns:
        np.ndarray: Array of exponentially distributed samples.
    """
    # Generate uniform random numbers
    U = np.random.rand(n_samples)

    # Apply the inverse transform
    X = -np.log(U)

    return X

# Generate 100,000 samples for testing
n_samples = 100000
inverse_samples = exponential_inverse_transform(n_samples)

Method 2: Von Neumann’s Ingenious Trick (Rejection Sampling)

Von Neumann proposed a clever alternative that avoids transcendental functions entirely. His procedure, which he noted “resembles a well known game of chance Twenty-One, or Black Jack,” generates sequences of uniform random numbers and accepts or rejects them based on simple comparison rules.

The algorithm works as follows to generate a single exponential sample $X$:

1. Initialize: Start with an integer offset k = 0, which will form the integer part of the final result.

2. Generate a trial sequence:

   • Generate uniform random numbers $Y_1, Y_2, Y_3, \ldots$ from (0, 1)
   • Find the smallest integer n such that the sequence is no longer strictly decreasing
   • That is, find n where $Y_1 > Y_2 > \cdots > Y_n$ but $Y_n \leq Y_{n+1}$

3. Accept or reject:

   • If n is odd: Accept the trial. Return $X = Y_1 + k$ and terminate.
   • If n is even: Reject the trial. Increment k by 1 and start a new trial.

This process is guaranteed to terminate and produces samples that follow the exponential distribution exactly. As von Neumann elegantly put it, the machine has “in effect computed a logarithm by performing only discriminations on the relative magnitude of numbers.”

Python Implementation

This implementation requires more careful state management due to the nested trial structure:

import numpy as np

def exponential_von_neumann(n_samples=1):
    """
    Generate samples from an exponential distribution using von Neumann's
    comparison-based rejection sampling method.

    Args:
        n_samples (int): Number of samples to generate.

    Returns:
        tuple[np.ndarray, float]: Array of samples and average uniform draws per sample.
    """
    samples = []
    total_uniform_draws = 0

    for _ in range(n_samples):
        k = 0  # Integer offset

        while True:  # Trial loop
            # Generate decreasing sequence
            y_prev = np.random.rand()
            total_uniform_draws += 1
            y1 = y_prev  # Store first value
            n = 1

            # Find length of decreasing sequence
            while True:
                y_curr = np.random.rand()
                total_uniform_draws += 1
                if y_prev <= y_curr:
                    break  # Sequence no longer decreasing
                y_prev = y_curr
                n += 1

            # Accept if n is odd, reject if even
            if n % 2 == 1:  # Accept
                samples.append(y1 + k)
                break
            else:  # Reject
                k += 1

    avg_draws = total_uniform_draws / n_samples
    return np.array(samples), avg_draws

# Generate samples using von Neumann's method
von_neumann_samples, avg_draws = exponential_von_neumann(n_samples)
print(f"Von Neumann method used {avg_draws:.2f} uniform draws per sample on average.")

Von Neumann method used 4.30 uniform draws per sample on average.

The algorithm requires approximately 4.3 uniform draws per exponential sample, matching the theoretical value $e^2/(e-1) = 4.30$.

Verification and Comparison

The critical test: do both methods actually produce the same distribution? And how do their performance characteristics compare?

Visual Verification

Let’s plot histograms of samples from both methods alongside the theoretical probability density function $f(x) = e^{-x}$:

import matplotlib.pyplot as plt
import seaborn as sns

# Configure plot aesthetics
sns.set_style("whitegrid")
plt.figure(figsize=(12, 7))

# Plot histograms for both methods
plt.hist(inverse_samples, bins=50, density=True, alpha=0.7,
         label='Inverse Transform', color='skyblue')
plt.hist(von_neumann_samples, bins=50, density=True, alpha=0.7,
         label="Von Neumann's Method", color='lightcoral')

# Overlay theoretical PDF
x = np.linspace(0, 8, 400)
pdf = np.exp(-x)
plt.plot(x, pdf, 'r-', linewidth=2, label='Theoretical PDF ($e^{-x}$)')

plt.title('Exponential Sampling Methods vs. Theoretical Distribution', fontsize=16)
plt.xlabel('x', fontsize=12)
plt.ylabel('Density', fontsize=12)
plt.legend()
plt.xlim(0, 8)
plt.tight_layout()
plt.show()

The visualization confirms that both methods accurately reproduce the target exponential distribution. The empirical histograms match the theoretical curve, confirming both algorithms sample the target distribution.

Performance Analysis

Mathematical elegance often diverges from computational efficiency. Von Neumann himself observed that on the ENIAC, it was actually “slightly quicker to use a truncated power series for log(1-T)” than to perform all the comparisons his method required.

Let’s benchmark both approaches in a modern Python environment:

import time

# Benchmark inverse transform method
start_time = time.time()
_ = exponential_inverse_transform(n_samples)
inverse_time = time.time() - start_time

# Benchmark von Neumann method
start_time = time.time()
_ = exponential_von_neumann(n_samples)
vn_time = time.time() - start_time

print(f"Inverse Transform:  {inverse_time:.4f} seconds")
print(f"Von Neumann Method: {vn_time:.4f} seconds")
print(f"Speedup factor: {vn_time / inverse_time:.1f}x")

Inverse Transform:  0.0018 seconds
Von Neumann Method: 0.1860 seconds
Speedup factor: 103.3x

The gap is large. The vectorized NumPy implementation of inverse transform sampling, leveraging a highly optimized C-backed logarithm function, outperforms the Python-looped von Neumann implementation by more than two orders of magnitude. While a vectorized or JIT-compiled version of von Neumann’s method would close this gap by removing Python interpreter overhead, the inverse transform remains the practical winner on modern hardware with fast floating-point units. This confirms von Neumann’s prescient observation: the “theoretically elegant” method avoiding transcendental functions often yields to direct computation.

Conclusion

This exploration offers a window into the ingenuity of early computational mathematics. Von Neumann’s comparison-based algorithm demonstrates remarkable mathematical creativity (showing how to “compute a logarithm” using only basic machine operations). Our implementation reproduces the algorithm, producing samples whose histogram and moments match the exponential distribution.

The performance comparison validates von Neumann’s own pragmatic assessment. His rejection sampling method is intellectually elegant and historically significant. The direct logarithmic approach proves far more efficient on both early and modern hardware. It serves as a timeless reminder in scientific computing: theoretical beauty often diverges from computational practicality.

The enduring value of von Neumann’s work lies in the fundamental insight that creative mathematical thinking can circumvent apparent computational limitations. Understanding alternative methods deepens our appreciation for the rich landscape of algorithmic possibilities, even when the direct approach proves superior.

The Müller-Brown potential reads, in hindsight, like an adversarial example for optimization algorithms.

Designed in 1979 to break naive path-finding methods, this deceptively simple 2D surface features deep minima, high barriers, and tricky saddle points. For nearly five decades, it has served as a ground-truth benchmark for computational chemistry.

Today, it finds new life as a testbed for machine learning. In the 1970s, chemists struggled to find transition states, saddle points where standard gradient descent fails catastrophically. Modern machine learning engineers face a strikingly similar challenge: escaping saddle points in high-dimensional loss landscapes. The Müller-Brown potential was the original stress test for these algorithms, and it remains a perfect, low-cost sandbox for benchmarking modern optimizers.

Whether you’re training neural network potentials or benchmarking reinforcement learning agents for exploration, the Müller-Brown potential offers a fast, noise-free, and mathematically exact environment.

In this guide, we’ll implement it in PyTorch, work through the engineering trade-off between analytical derivatives and Autograd, and measure the roughly 4x force-evaluation speedup of the compiled analytical kernel over the Autograd reference.

The Problem: Finding Saddle Points in the 1970s

In the 1970s, finding energy minima was straightforward: follow the gradient downhill. But finding transition states proved much more challenging. These saddle points are maxima along the reaction coordinate but minima in all other directions, like standing at the top of a mountain pass.

Standard first-order optimizers, whether 1970s simplex methods or modern SGD and Adam, are designed to minimize loss functions blindly. Point them at a saddle point, and they slide into the nearest valley. The gradients vanish or point in misleading directions. Specialized algorithms were needed to navigate this mixed landscape of ups and downs.

The computational reality made this worse. Early quantum chemistry programs like ATMOL and Gaussian made energy calculations possible, but each computation was expensive. Gradients required even more resources, and second derivatives were rarely computed.

This created a catch-22: sophisticated algorithms were needed to find saddle points, but researchers couldn’t afford to test them on real molecular systems. Every calculation represented a major investment of time and computational resources.

Müller and Brown’s Solution

Müller and Brown’s insight¹ was to create a simple analytical test function that captured the essential difficulties of real chemical systems without the computational cost. Their potential offered three key advantages:

• Negligible computational cost - Evaluate millions of points instantly
• Analytical derivatives - Exact gradients and Hessians available immediately
• Realistic challenges - Multiple minima, saddle points, and curved pathways

The clever part was the deliberate design to break naive approaches. Early methods often assumed linear paths between reactants and products. The Müller-Brown potential has a curved minimum energy path that punishes this assumption. Try to take shortcuts, and algorithms climb over high-energy barriers.

The Mathematical Foundation

The Müller-Brown potential combines four two-dimensional Gaussian functions:

$$V(x,y) = \sum_{k=1}^{4} A_k \exp\left[a_k(x-x_k^0)^2 + b_k(x-x_k^0)(y-y_k^0) + c_k(y-y_k^0)^2\right]$$

Each Gaussian contributes a different “bump” or “well” to the landscape. The parameters control amplitude ($A_k$), width, orientation, and center position.

The Standard Parameters

The specific parameter values that define the canonical Müller-Brown surface are:

Notice that the first three terms have negative amplitudes (creating energy wells), while the fourth has a positive amplitude (creating a barrier). The cross-term $b_k$ in the third Gaussian creates the tilted orientation that gives the surface its characteristic curved pathways.

View Interactive Müller-Brown Potential Energy Surface →

The Resulting Landscape

This simple formula creates a surprisingly rich topography with exactly the features needed to challenge optimization algorithms:

The Key Challenge: Curved Pathways

The path from the deep reactant minimum (MA) to the product minimum (MB) doesn’t go directly over a single barrier. Instead, it follows a curved route:

1. MA → S1 → MC: First transition over the higher, rate-limiting barrier (S1) into an intermediate basin
2. MC → S2 → MB: Second transition over a much lower barrier (S2) to the product

This two-step pathway breaks linear interpolation methods. Algorithms that draw a straight line from reactant to product miss both the intermediate minimum and the correct transition states, climbing over much higher energy regions instead.

The energy profile comparison makes the failure mode concrete. A naive optimizer following the red dashed path would encounter a barrier roughly 53 reduced units higher than necessary (the naive summit sits at about +13 while the true path tops out near -41, the S1 saddle). The green minimum energy path navigates through the valleys, passing through the intermediate basin MC and crossing only the low-lying saddle points S1 and S2.

Why It Works as a Benchmark

The Müller-Brown potential has served as a computational chemistry benchmark for over four decades because of four key characteristics:

Low dimensionality: As a 2D surface, you can visualize the entire landscape and see exactly why algorithms succeed or fail.

Analytical form: Energy and gradient calculations cost virtually nothing, enabling exhaustive testing impossible with quantum mechanical surfaces.

Non-trivial topology: The curved minimum energy path and shallow intermediate minimum challenge sophisticated methods while remaining manageable.

Known ground truth: All minima and saddle points are precisely known, providing unambiguous success metrics.

This basin map reveals why the Müller-Brown potential is such an effective benchmark. Standard gradient descent from any point in the blue region inevitably falls into the deep MA minimum; from yellow, into MB. The saddle points S1 and S2 lie exactly on the boundaries between basins, infinitesimally perturbing an optimizer at these points sends it tumbling into different valleys. Finding these saddle points requires algorithms that identify and stabilize at these boundary regions.

What makes this particularly valuable is the contrast with other classic potentials. While the Lennard-Jones potential² serves as the benchmark for equilibrium properties with its single energy minimum, Müller-Brown explicitly models reactive landscapes. Its multiple minima and connecting barriers make it the testing ground for algorithms that find reaction paths, the methods that reveal how chemistry actually happens.

Applications Across Decades

The potential has evolved with the field’s changing focus:

1980s-1990s: Testing path-finding methods like Nudged Elastic Band (NEB)³, which creates discrete representations of reaction pathways and optimizes them to find minimum energy paths.

2000s-2010s: Validating Transition Path Sampling (TPS) methods⁴ that harvest statistical ensembles of reactive trajectories.

2020s: Benchmarking machine learning models and generative approaches that learn to sample transition paths or approximate potential energy surfaces.

Modern Applications in Machine Learning

The rise of machine learning has given the Müller-Brown potential renewed purpose. Modern Machine Learning Interatomic Potentials (MLIPs)⁵⁶ aim to bridge the gap between quantum mechanical accuracy and classical force field efficiency by training flexible models on expensive quantum chemistry data.

This creates a benchmarking challenge: with countless ML architectures available, how do you objectively compare them? The Müller-Brown potential provides an ideal solution, an exactly known potential energy surface that can generate unlimited, noise-free training data.

This enables researchers to ask fundamental questions:

• How well does a given architecture learn complex, curved surfaces?
• How many training points are needed for acceptable accuracy?
• How does the model behave when extrapolating beyond training data?
• Can it correctly identify minima and saddle points?

The potential has evolved from a simple model system into a reference benchmark: a fixed, exactly-known surface against which AI learning capacity is measured. Any prediction error is due to model limitations, not data quality.

Beyond static benchmarking, a PyTorch implementation enables differentiable simulation. Because the potential, forces, and integrator are all differentiable tensor operations, gradients can be backpropagated through time (via the trajectory) to optimize force field parameters or control policies directly. This capability connects classical molecular simulation to modern gradient-based machine learning in a single computational graph.

These benchmarking principles extend beyond abstract test cases. Real molecular dynamics simulations, such as those studying adatom diffusion on metal surfaces, face similar challenges in understanding energy landscapes and transition pathways. The Müller-Brown potential provides a controlled environment for developing methods that eventually tackle these complex realistic systems.

Extension to Higher Dimensions

The canonical Müller-Brown potential can be extended beyond two dimensions to create more challenging test cases that better reflect real molecular systems. This extensibility demonstrates why it remains such an effective template for computational method development.

Why Higher Dimensions Matter

Real molecules have dozens or hundreds of degrees of freedom. Understanding how algorithms scale with dimensionality is crucial for practical applications. Higher-dimensional extensions allow researchers to systematically test:

• Algorithm scaling - Does performance degrade gracefully as dimensions increase?
• Model robustness - Do machine learning approaches maintain accuracy in high-dimensional spaces?
• Parallel efficiency - Can massively parallel methods exploit additional dimensions effectively?

Extension Approaches

Harmonic constraints: Add quadratic wells in orthogonal dimensions while preserving the complex 2D landscape⁷:

$$V_{5D}(x_1, x_2, x_3, x_4, x_5) = V(x_1, x_3) + \kappa(x_2^2 + x_4^2 + x_5^2)$$

The parameter $\kappa$ controls constraint strength: small values create nearly flat directions that test algorithmic efficiency.

Collective variables: Define new coordinates that mix multiple dimensions⁸:

$$\tilde{x} = \sqrt{x_1^2 + x_2^2 + \epsilon x_5^2},\quad \tilde{y} = \sqrt{x_3^2 + x_4^2}$$

where $\epsilon \ll 1$. The 5D potential becomes $V_{5D}(\tilde{x}, \tilde{y}) = V(\tilde{x}, \tilde{y})$, embedding the original surface in a higher-dimensional space.

Value for Algorithm Development

This extensibility makes the Müller-Brown potential ideal for systematic testing:

• Progressive complexity: Debug on 2D, then scale to higher dimensions
• Ground truth preservation: Known minima and saddle points remain in the active subspace
• Realistic challenges: Captures the “needle in a haystack” problem of transition state finding while maintaining analytical tractability

These extensions transform a simple 2D benchmark into a scalable testbed for modern computational methods, probing specific challenges in high-dimensional optimization.

Implementation in PyTorch

Now let’s implement the Müller-Brown potential in PyTorch. A practical implementation needs to handle batch processing, support both analytical and automatic differentiation, and be optimized for performance.

Core Implementation

import torch
import torch.nn as nn
from torch import Tensor

class MuellerBrownPotential(nn.Module):
    """Müller-Brown potential with a torch.compile-accelerated force kernel."""

    def __init__(
        self,
        device: str | torch.device = "cpu",
        dtype: torch.dtype = torch.float64,
        use_autograd: bool = False
    ):
        super().__init__()
        self.use_autograd = use_autograd

        # Standard Müller-Brown parameters
        self.register_buffer(
            "A", torch.tensor([-200.0, -100.0, -170.0, 15.0],
                             device=device, dtype=dtype)
        )
        self.register_buffer(
            "a", torch.tensor([-1.0, -1.0, -6.5, 0.7],
                             device=device, dtype=dtype)
        )
        self.register_buffer(
            "b", torch.tensor([0.0, 0.0, 11.0, 0.6],
                             device=device, dtype=dtype)
        )
        self.register_buffer(
            "c", torch.tensor([-10.0, -10.0, -6.5, 0.7],
                             device=device, dtype=dtype)
        )
        self.register_buffer(
            "x_centers", torch.tensor([1.0, 0.0, -0.5, -1.0],
                                    device=device, dtype=dtype)
        )
        self.register_buffer(
            "y_centers", torch.tensor([0.0, 0.5, 1.5, 1.0],
                                    device=device, dtype=dtype)
        )

    def forward(self, coordinates: Tensor) -> Tensor:
        """Compute potential energy."""
        return _calculate_potential(
            coordinates, self.A, self.a, self.b, self.c,
            self.x_centers, self.y_centers
        )

    def force(self, coordinates: Tensor) -> Tensor:
        """Compute forces (negative gradient)."""
        if self.use_autograd:
            coordinates = coordinates.requires_grad_(True)
            potential = self.forward(coordinates)
            grad = torch.autograd.grad(potential.sum(), coordinates)[0]
            return -grad
        else:
            return _calculate_force(
                coordinates, self.A, self.a, self.b, self.c,
                self.x_centers, self.y_centers
            )

The implementation uses register_buffer to store parameters, a subtle but important PyTorch best practice. This ensures that the potential parameters are automatically moved to the GPU along with the model when calling .to(device), a common pitfall that leads to frustrating device mismatch errors. Beyond device placement, register_buffer also ensures these parameters are correctly handled during DistributedDataParallel (DDP) broadcasting, preventing silent failures when scaling training to multi-GPU clusters. The use_autograd flag switches between analytical and automatic differentiation.

Compiling the Force Kernel

Two decisions make the force path fast while keeping the rest flexible:

def _calculate_potential(coordinates: Tensor, A: Tensor, a: Tensor,
                        b: Tensor, c: Tensor, x_centers: Tensor,
                        y_centers: Tensor) -> Tensor:
    """Energy. Left eager (uncompiled) so autograd second derivatives,
    e.g. the Hessian, keep working: torch.compile does not support
    double-backward, and energy is computed per save, not per step."""
    coords = coordinates.view(-1, 2)
    x, y = coords[:, 0], coords[:, 1]
    dx = x.unsqueeze(-1) - x_centers
    dy = y.unsqueeze(-1) - y_centers

    potential = torch.sum(
        A * torch.exp(a * dx**2 + b * dx * dy + c * dy**2),
        dim=-1
    )
    return potential.view(coordinates.shape[:-1])

@torch.compile(dynamic=True)
def _calculate_force(coordinates: Tensor, A: Tensor, a: Tensor,
                    b: Tensor, c: Tensor, x_centers: Tensor,
                    y_centers: Tensor) -> Tensor:
    """Forces (negative gradient). Compiled: this is the hot path,
    called once per simulation step."""
    coords = coordinates.view(-1, 2)
    x, y = coords[:, 0], coords[:, 1]
    dx = x.unsqueeze(-1) - x_centers
    dy = y.unsqueeze(-1) - y_centers

    exp_terms = torch.exp(a * dx**2 + b * dx * dy + c * dy**2)
    A_exp = A * exp_terms

    grad_x = torch.sum(A_exp * (2 * a * dx + b * dy), dim=-1)
    grad_y = torch.sum(A_exp * (b * dx + 2 * c * dy), dim=-1)

    forces = torch.stack([-grad_x, -grad_y], dim=-1)
    new_shape = list(coordinates.shape[:-1]) + [2]
    return forces.view(new_shape)

The force kernel is decorated with @torch.compile(dynamic=True), which traces it once and hands it to TorchInductor to fuse the pointwise operations (the exponentials and polynomial terms) and cut Python’s dispatch overhead in the inner loop that runs millions of times per simulation. The dynamic=True flag keeps a single compiled trace valid across particle counts, so changing the batch size does not trigger a recompile. The energy is left eager on purpose: torch.compile does not support double-backward, so leaving forward uncompiled keeps autograd second derivatives (the Hessian) available, and since the energy is only evaluated when an observable is saved, it is not on the hot path anyway.

Performance: Analytical vs. Automatic Differentiation

A key design decision is whether to use analytical derivatives or automatic differentiation. I benchmarked both on an Apple M1 Max (CPU, PyTorch 2.x). To get a stable measurement, each configuration runs 100 warm-up iterations, then the median wall-clock time over 5 runs of 1000 iterations, which filters out operating-system jitter.

The speedup comes from bypassing the computational graph that PyTorch’s Autograd engine builds. By deriving the analytical Jacobian, we skip that machinery entirely. Every autograd.grad() call must:

1. Build a tape of operations during the forward pass
2. Traverse the graph backward to compute gradients
3. Allocate intermediate tensors for each node

For iterative workloads like molecular dynamics (millions of force evaluations per trajectory), this overhead elimination is critical. The analytical kernel computes forces directly in a single fused operation, no graph, no tape, no intermediate allocations.

When to use each approach:

• Analytical: Best for production molecular dynamics where forces are computed millions of times. The speedup directly reduces wall-clock simulation time.
• Autograd: Better for prototyping, machine learning training loops, or when implementing new potentials where correctness verification is paramount. The convenience and guaranteed accuracy often outweigh performance costs during development.

Molecular Dynamics Simulations

To demonstrate the PyTorch implementation in action, I performed Langevin dynamics simulations in different energy basins. These simulations reveal how particles behave when confined to different regions of the potential energy surface.

Simulation Parameters

I ran 3600 time steps with a 0.01 time unit step size, using a friction coefficient of 1.0 and temperature of 25.0 in reduced units. The simulations started from equilibrium positions within each basin and show the characteristic thermal fluctuations around local minima.

Basin MA: Deep Reactant Minimum

The deepest energy well (-146.70 in reduced units) shows highly constrained motion due to the steep energy barriers surrounding it.

Basin MB: Product Minimum

The product minimum (-108.17 in reduced units) shows intermediate behavior between the deep reactant well and shallow intermediate basin.

Transition Example

Running a longer simulation demonstrates transitions between basins:

Integration with Modern Workflows

The PyTorch implementation integrates naturally with machine learning workflows. It can serve as a component in larger computational graphs, enable gradient-based optimization, and bridge classical molecular simulation with deep learning approaches.

For readers interested in broader molecular ML applications, this implementation pairs well with other molecular representation methods. The Coulomb matrix approach offers complementary perspectives on encoding molecular structure, while the Kabsch algorithm provides essential tools for structural alignment.

The complete implementation is available on GitHub⁹¹⁰¹¹, including benchmarking scripts, visualization tools, the test suite, and examples for optimization and molecular dynamics.

Conclusion

The Müller-Brown potential exemplifies how a well-designed benchmark can evolve with a field. Born from 1970s computational constraints, it provided a simple way to test algorithms when quantum chemistry calculations were expensive. Its clever design, simple enough to compute instantly, complex enough to break naive approaches, made it invaluable for algorithm development.

Today, it serves new purposes in the machine learning era. This PyTorch implementation pairs a hand-derived analytical force kernel (compiled with torch.compile) with an autograd reference, a BAOAB Langevin sampler, and a test suite that checks the sampler against the canonical distribution. The analytical kernel’s roughly 4x advantage matters for intensive simulations, while PyTorch’s flexibility enables integration with neural network potentials and enhanced sampling methods.

The potential’s evolution from practical necessity to pedagogical tool to machine learning benchmark demonstrates the value of foundational test cases. As computational chemistry continues evolving, reliable standards like the Müller-Brown potential become even more important for rigorous method development and comparison.

For the complete implementation with benchmarking scripts, the BAOAB Langevin simulator, and visualization tools, see the GitHub repository. The full project with architecture details and performance results is at the Müller-Brown Potential: A PyTorch ML Testbed project page.

The physics of liquid argon is a solved problem. We know the answer.

So, why replicate it in 2025? To apply modern engineering standards to legacy science.

This project served as an exercise in software archaeology: taking a vintage scientific workflow and rebuilding it with a modular Python analysis pipeline. I wanted to see if I could replace Rahman’s “write-once” Fortran mentality with modern reproducibility, type safety, and intelligent caching.

The full source code is available on GitHub. The complete project overview, including analysis results and pipeline architecture, is on the Rahman 1964 Replication project page.

Engineering the Pipeline

The most interesting part of this project isn’t the simulation engine (LAMMPS handles that); it’s the architecture of the analysis suite. MD analysis is computationally expensive ($O(N^2)$), and iterating on plots can be painfully slow if you re-compute trajectory data every time.

Why bother? Don’t modern MD packages come with analysis tools? Well, some say that writing is thinking. Sometimes getting into the weeds of how an algorithm works or an analysis is performed, you gain insights and a deeper understanding that might be obscured by a plug-and-play tool.

Intelligent Caching

I built the argon_sim package with a decorator-based caching layer. The system hashes the source file’s modification time and the function’s arguments to avoid re-calculating the Radial Distribution Function (RDF) or Van Hove correlations on every script run.

@cached_computation("gr")
def compute_radial_distribution(filename: str, dr: float = 0.05):
    # ... expensive O(N^2) distance calculations ...
    return r_values, g_r, density

If I tweak a plot axis, the script runs instantly, loading pre-computed arrays from disk instead of re-running the $O(N^2)$ computation. If I change the simulation trajectory, the cache invalidates automatically.

Vectorization & Memory Management

Rahman likely relied on nested loops. Python is too slow for that. I utilized NumPy broadcasting to vectorize the calculation of atomic displacements.

However, calculating an $864 \times 864$ distance matrix for 5,000 frames consumes significant RAM. I implemented a chunked MSD (Mean Square Displacement) algorithm that processes the trajectory in blocks, balancing vectorization speed with memory constraints. The chunking trades some vectorization speed for a bounded memory footprint, so the analysis is not capped by holding the full distance matrix in RAM.

Reproducibility as a Feature

Academic code is notorious for “it works on my machine.” To combat this, I used uv for dependency management, locking the exact environment state. The entire workflow (from simulation to final figure generation) is abstracted into a Makefile.

# One command to run the physics, analyze data, and generate plots
make workflow

The Simulation: 1964 vs. 2025

I preserved Rahman’s physical parameters exactly to ensure a fair comparison:

• System: 864 Argon atoms
• Potential: Lennard-Jones ($\sigma = 3.4$ Å, $\epsilon/k_B = 120$ K)
• Target: 94.4 K, 1.374 g/cm³

However, I modernized the numerical methods to ensure stability:

The production run lasted 10 ps in the NVE ensemble, generating 5,001 frames. Temperature remained within 1% of target with an RMS fluctuation of 0.0165.

Validation Results

The replication was quantitatively successful. The analysis pipeline faithfully reproduced every key signature of liquid argon.

The Cage Effect

This is the paper’s crown jewel. In a gas, velocity correlations decay exponentially. In a liquid, Rahman discovered that atoms get trapped by their neighbors and bounce back, causing the correlation to go negative.

My simulation captures this minimum at -0.083, matching Rahman’s observation. The Fourier transform of this data (the frequency spectrum) reveals a peak at $\beta \approx 0.25$, physically representing the frequency of atomic collisions within the cage.

Structural Fingerprints

The Radial Distribution Function $g(r)$ and its Fourier transform, the Structure Factor $S(k)$, are the “fingerprints” of a liquid’s structure.

The agreement here is striking. My first peak appeared at 3.82 Å (Rahman: 3.7 Å). The slight discrepancy is likely due to my improved equilibration method, which allowed the system to relax into a more natural liquid state than Rahman’s 1960s hardware allowed.

Diffusion and Non-Gaussian Behavior

By calculating the Mean Square Displacement (MSD), I derived a diffusion coefficient of $D = 2.47 \times 10^{-5}$ cm²/s, which deviates only 2% from Rahman’s reported $2.43 \times 10^{-5}$.

More interestingly, I reproduced the “Non-Gaussian” parameters. Standard diffusion assumes a Gaussian distribution of displacements. Rahman found (and I confirmed) that liquid atoms deviate from this. They exhibit “jump” and “wait” dynamics, a behavior that standard Brownian motion models fail to capture.

Advanced Analysis: Van Hove Functions

Rahman also explored advanced properties like the Van Hove correlation function $G(r,t)$, which describes how liquid structure evolves over time.

At 1.0 ps, the structure remains well-defined with clear shells. By 2.5 ps, it becomes increasingly diffuse. Rahman compared this evolution to theoretical predictions (the Vineyard approximation) and found that theory predicted overly rapid structural decay. My results confirm this finding.

System Validation

Before analyzing physics, basic sanity checks confirmed proper thermal equilibrium.

Mean temperature was 94.73 K (0.33 K off target) with a standard deviation of 1.56 K.

The velocity distribution from 12.9 million velocity components produces a clean Maxwell-Boltzmann distribution, as expected for thermal equilibrium. The distribution widths at various heights closely match Rahman’s results: 1.77, 2.48, and 3.56 compared to his 1.77, 2.52, and 3.52.

Conclusion

Replicating a 60-year-old paper might seem like a solved puzzle, but it teaches a valuable lesson in computational science. Rahman relied on brilliance and raw mathematical intuition because he lacked compute power. Today, pairing modern compute with disciplined software practices makes the same result reproducible and auditable.

Applying modern software engineering (modular architecture, caching, and automated workflows) to classical physics reproduces the past and builds a foundation that makes the next discovery easier, faster, and more reliable.

The quantitative agreement is striking: diffusion coefficients within 2%, structural peaks within 0.1 Å, velocity distributions matching to three significant figures. This level of reproducibility, achieved with completely different hardware and software, validates something fundamental: Rahman’s physical model was remarkably sound, and his computational methodology was scientifically rigorous despite 1960s constraints.

The cage effect, velocity correlations, and structural evolution are fundamental characteristics of how matter behaves at the atomic scale, as relevant today as they were six decades ago.

A Variational Autoencoder (VAE) is a type of generative model, meaning its primary purpose is to learn the underlying structure of a dataset so it can generate new, similar data.

Whether the data is images, raw audio clips, or 2D graphs of drug-like molecules, a VAE aims to capture the essential features that define the data distribution. Once trained, it should be able to create entirely new samples that resemble the training data without simply copying specific examples.

Introduced by Kingma and Welling in 2013 (Auto-Encoding Variational Bayes, Paper), VAEs are used for:

• Generation: Creating new data (images, music, text).
• Dimensionality Reduction: Compressing data into a much smaller, meaningful representation (a “latent space”).
• Imputation: Intelligently filling in missing data (e.g., denoising images).

Importantly, they aim to provide a structured and continuous latent space, which allows for smooth interpolation between data points and meaningful manipulations of generated samples (think: optimization).

TL;DR: The Complete PyTorch Implementation

For those who just want the code, here is a complete, modern VAE implementation in PyTorch. It features softplus standard deviation parameterization for numerical stability and a custom training step that handles the ELBO loss correctly.

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class VAEOutput:
    z: torch.Tensor
    mu: torch.Tensor
    std: torch.Tensor
    x_recon: torch.Tensor
    loss: torch.Tensor
    loss_recon: torch.Tensor
    loss_kl: torch.Tensor

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=512, latent_dim=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh()
        )
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_std = nn.Linear(hidden_dim, latent_dim)

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, input_dim)
        )

    def encode(self, x):
        h = self.encoder(x)
        mu = self.fc_mu(h)
        # Softplus + epsilon for stable std deviation
        std = F.softplus(self.fc_std(h)) + 1e-6
        return mu, std

    def reparameterize(self, mu, std):
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x, kl_weight=1.0):
        mu, std = self.encode(x)
        z = self.reparameterize(mu, std)
        x_recon = self.decode(z)

        # 1. Reconstruction Loss (Binary Cross Entropy for MNIST)
        # Sum over features, mean over batch
        recon_loss = F.binary_cross_entropy_with_logits(x_recon, x, reduction='none').sum(dim=1).mean()

        # 2. KL Divergence
        # Analytic KL for Normal distributions
        kl_loss = -0.5 * torch.sum(1 + torch.log(std**2) - mu**2 - std**2, dim=1).mean()

        # 3. Total Loss (ELBO)
        loss = recon_loss + (kl_weight * kl_loss)

        return VAEOutput(z, mu, std, x_recon, loss, recon_loss, kl_loss)

# --- Training Loop Example ---
def train_step(model, batch, optimizer, kl_weight=1.0):
    model.train()
    optimizer.zero_grad()

    # Forward pass
    output = model(batch, kl_weight)

    # Backward pass
    output.loss.backward()

    # Gradient clipping (recommended)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    optimizer.step()
    return output.loss.item()

The Core Idea: Learning to Generate

The VAE is built on a key assumption: our complex, high-dimensional data (like a $28 \times 28$ pixel image, $\mathbf{x}$) is actually generated by some simpler, low-dimensional, unobserved variable (a “latent” variable, $\mathbf{z}$).

A Physical Metaphor: Water Molecules and Phase Diagrams

Consider a glass of water. At the microscopic level, you have more than $10^{24}$ $\text{H}_2\text{O}$ molecules bouncing around in an incredibly high-dimensional space. Each molecule has position, velocity, and interactions with its neighbors, computationally intractable to track directly. Yet we can describe the macroscopic behavior of all these molecules using just two simple variables: temperature and pressure. These two dimensions create a “phase diagram” that tells us whether our water will be ice, liquid, or vapor. The temperature and pressure are “latent variables” that capture the essential physics governing this complex molecular dance.

A VAE makes the same assumption: complex data (like images) emerges from simpler underlying factors. A handwritten digit might be generated by latent factors like “pen thickness,” “writing angle,” “digit style,” and “size,” a much simpler description than tracking all 784 pixel values independently.

The VAE learns two functions: one that maps from complex data ($\mathbf{x}$) to these descriptive factors ($\mathbf{z}$), and another that maps from these factors back to the data. It accomplishes this with two main components, typically implemented as neural networks:

1. The Encoder (Recognition Model)

This network takes a complex data point $\mathbf{x}$ (an image) and determines the “knob settings” $\mathbf{z}$ that could explain or generate it. This allows us to compress or understand the data.

$$q_{\phi}(\mathbf{z} | \mathbf{x})$$

It’s like examining a container of molecules and summarizing their complex arrangement into key parameters like temperature and pressure.

Crucially, the encoder outputs the parameters of a probability distribution (a simple Gaussian) that describes $\mathbf{z}$.

For each input $\mathbf{x}$, the encoder network outputs:

• A vector of means, $\mathbf{\mu}$
• A vector of standard deviations, $\mathbf{\sigma}$

These parameters define our approximation $q_{\phi}(\mathbf{z} | \mathbf{x}) = \mathcal{N}(\mathbf{z} \mid \mathbf{\mu}, \mathbf{\sigma}^2\mathbf{I})$. We then sample from this distribution to get the $\mathbf{z}$ that we feed to the decoder. This probabilistic step is what forces the latent space to be continuous and structured. It forces similar inputs to map to nearby regions in latent space, enabling smooth interpolation and generation.

2. The Decoder (Generative Model)

This network learns the “generative process.” It takes a simple latent vector $\mathbf{z}$ and reconstructs the complex data $\mathbf{x}$. This allows us to generate new data by feeding it a random $\mathbf{z}$ and observing what image $\mathbf{x}$ it produces.

$$p_{\theta}(\mathbf{x} | \mathbf{z})$$

The decoder reverses the encoder: it takes the simple latent representation and “paints” the full, complex image from it. It’s like taking temperature and pressure values and producing a detailed arrangement of water molecules consistent with those conditions. The goal is to reproduce the exact input as closely as possible.

After training, we have two networks that can be used for a variety of purposes:

• Generation: If the latent space is well-structured, we can sample random $\mathbf{z}$ vectors from a simple distribution (like a standard normal) and feed them into the Decoder to generate new images. This is particularly useful for searching for data points with desired properties, like in drug discovery, where we might want to generate molecules with specific characteristics.
• Compression: The Encoder can compress complex data into a low-dimensional latent space, which can be useful for visualization or as a feature extractor for other tasks.

The “Variational” Problem

Calculating the true distribution of latent variables $p_{\theta}(\mathbf{z}|\mathbf{x})$ (the posterior) is mathematically intractable.

This intractability arises from Bayes’ theorem:

$$p_{\theta}(\mathbf{z} | \mathbf{x}) = \frac{p_{\theta}(\mathbf{x} | \mathbf{z}) p_{\theta}(\mathbf{z})}{p_{\theta}(\mathbf{x})}$$

Breaking down each component:

• $p_{\theta}(\mathbf{x} | \mathbf{z})$ is our decoder, which is straightforward to compute given our likelihood model.
• $p_{\theta}(\mathbf{z})$ is our prior over latent variables, typically a simple distribution like a standard normal, making it easy to compute.
• $p_{\theta}(\mathbf{x})$ is the marginal likelihood of the data. And here lies the problem. It requires integrating over all possible latent variables that could have generated $\mathbf{x}$: $$p_{\theta}(\mathbf{x}) = \int p_{\theta}(\mathbf{x} | \mathbf{z}) p_{\theta}(\mathbf{z}) d\mathbf{z}$$ It is the normalization factor that ensures the posterior is a valid probability distribution (i.e., sums to 1 over all $\mathbf{z}$).

This integral is intractable because it involves integrating over a high-dimensional latent space with a complex likelihood function. No closed-form solution exists, and numerical integration is computationally prohibitive.

This is where the “variational” approach provides the solution. We approximate the true posterior by learning an encoder, $q_{\phi}(\mathbf{z} | \mathbf{x})$, that serves as a variational approximation to this intractable true distribution. The VAE’s training process optimizes this approximation to be as accurate as possible, pushing this learned distribution closer to the true posterior.

The VAE Objective: A Balancing Act

To get these two networks (parameterized by $\theta$ and $\phi$) to work together, we train them jointly with a special loss function. This objective has two parts that balance two different goals:

1. Reconstruction Loss

$$E_{q_{\phi}(\mathbf{z} | \mathbf{x})}[\log p_{\theta}(\mathbf{x} | \mathbf{z})]$$

This term asks: “How well can we reconstruct our original image?” It forces the VAE to be good at its job. The process goes:

1. Take an input point $\mathbf{x}$.
2. Use the Encoder to get its latent representation $\mathbf{z} \sim q_{\phi}(\mathbf{z} | \mathbf{x})$.
3. Use the Decoder to generate a new image $\mathbf{x}’$ from $\mathbf{z}$, $\mathbf{x}’ \sim p_{\theta}(\mathbf{x} | \mathbf{z})$.
4. Compare $\mathbf{x}$ and $\mathbf{x}’$.

The reconstruction loss measures the difference between the original and the reconstructed image.

• For continuous inputs (like general images), this is often Mean Squared Error (MSE).
• For inputs in $[0, 1]$ (like MNIST pixel intensities, which after ToTensor() are continuous values in $[0, 1]$), we use Binary Cross-Entropy (BCE). We treat each pixel as an independent Bernoulli variable whose target is its intensity. The decoder outputs the logits for each pixel, and the BCE-with-logits loss (e.g., F.binary_cross_entropy_with_logits) is the numerically stable way to compute the negative log-likelihood.
• More generally, you can output parameters of a desired output distribution. What if you wanted a mixture of Gaussians? The decoder could output the means, variances, and mixture weights, and you could compute the negative log-likelihood accordingly.

This loss pushes the encoder to produce useful $\mathbf{z}$ vectors and pushes the decoder to learn how to interpret them accurately.

2. The KL Divergence (The Regularizer)

$$D_{KL}(q_{\phi}(\mathbf{z} | \mathbf{x}) || p_{\theta}(\mathbf{z}))$$

On its own, the reconstruction loss might “cheat.” The encoder could learn to map every image to a different, specific point in the latent space, essentially “memorizing” the data. While this minimizes reconstruction error, it creates a meaningless latent space that fails at generation.

The KL divergence term fixes this. It’s a regularizer that forces the latent space to be organized and smooth.

We force the encoder’s output, $q_{\phi}(\mathbf{z} | \mathbf{x})$, to be close to a simple, predefined prior distribution, $p_{\theta}(\mathbf{z})$. This prior is almost always a standard normal distribution because it is mathematically convenient, easy to sample from, and encourages a well-behaved latent space.

This regularization term acts as a penalty, measuring how much the encoder’s output distribution diverges from the simple prior. By minimizing this KL divergence, we encourage the model to:

• Avoid overfitting by preventing the encoder from memorizing specific locations for each input
• Create meaningful clusters where similar inputs map to nearby regions in the latent space
• Maintain continuity so that points close together in latent space (like different variations of the digit “7”) decode into visually similar outputs

This smooth, structured latent space is what enables generation: we can sample random points from our prior distribution and decode them into realistic new data.

Ultimately, the optimizer finds a balance between these two objectives: reconstructing the data well while keeping the latent space organized and regularized.

The Reparameterization Trick: Making it All Trainable

We have a problem. The training process requires sampling:

1. Encoder produces $\mathbf{\mu}$ and $\mathbf{\sigma}$.
2. We sample $\mathbf{z} \sim \mathcal{N}(\mathbf{\mu}, \mathbf{\sigma}^2\mathbf{I})$.
3. Decoder uses $\mathbf{z}$ to reconstruct $\mathbf{x}’$.
4. We calculate the loss.

The “sampling” step is a random, non-differentiable operation. We can’t backpropagate the reconstruction loss from the decoder through this random node to update the encoder’s weights.

The reparameterization trick makes the sampling process differentiable. We generate $\mathbf{z}$ deterministically by sampling a random noise vector and transforming it:

1. Sample a random noise vector $\mathbf{\epsilon}$ from a simple, fixed distribution (e.g., the standard normal $\mathcal{N}(\mathbf{0}, \mathbf{I})$).
2. Compute $\mathbf{z}$ as: $\mathbf{z} = \mathbf{\mu} + \mathbf{\sigma} \odot \mathbf{\epsilon}$

This simple change moves the randomness “outside” the network. The gradient can now flow deterministically from $\mathbf{z}$ back through the $\mathbf{\mu}$ and $\mathbf{\sigma}$ nodes to the encoder network. This is the key engineering insight that allows us to train the entire model end-to-end with standard backpropagation.

Where Does This Objective Come From? (The Math)

This two-part loss function is derived directly from the goal of maximizing the marginal likelihood of the data, $\log p_{\theta}(\mathbf{x})$.

For a single data point $\mathbf{x}^{(i)}$, we can write: $$\log p_{\theta}(\mathbf{x}^{(i)}) = D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}^{(i)}) || p_{\theta}(\mathbf{z} | \mathbf{x}^{(i)})) + \mathcal{L}(\theta, \phi; \mathbf{x}^{(i)})$$

• The first term is the KL divergence between our encoder’s approximation and the (intractable) true posterior. This is non-negative, and unfortunately we cannot compute it.
• The second term, $\mathcal{L}$, is the Variational Lower Bound (also known as the Evidence Lower Bound, or ELBO). Since the KL term is $\ge 0$, we know that $\log p_{\theta}(\mathbf{x}^{(i)}) \ge \mathcal{L}$.

By maximizing this lower bound $\mathcal{L}$, we push up the “floor” on the true likelihood of our data. This is a problem we can solve.

When we expand this $\mathcal{L}$ term, we get our famous two-part objective:

$$\mathcal{L}(\theta, \phi; \mathbf{x}^{(i)}) = E_{q_{\phi}(\mathbf{z} | \mathbf{x}^{(i)})}[\log p_{\theta}(\mathbf{x}^{(i)} | \mathbf{z})] - D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}^{(i)}) || p_{\theta}(\mathbf{z}))$$

• Term 1: The expected log-likelihood of reconstructing $\mathbf{x}^{(i)}$ from $\mathbf{z}$. Maximizing this is the same as minimizing the Reconstruction Loss.
• Term 2: The negative KL divergence between our encoder and the simple prior. Maximizing this is the same as minimizing the KL Divergence Loss.

Thus, the VAE’s objective balances these two critical goals: faithfully reconstructing the data while maintaining a simple, regularized latent structure that is useful for generation.

From ELBO to Practical Loss

Remember, our goal is to maximize the ELBO:

$$\mathcal{L}(\theta, \phi; \mathbf{x}) = E_{q_{\phi}(\mathbf{z} | \mathbf{x})}[\log p_{\theta}(\mathbf{x} | \mathbf{z})] - D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}) || p_{\theta}(\mathbf{z}))$$

Since deep learning libraries are built to minimize a loss function, we simply flip the sign and minimize the negative ELBO ($-\mathcal{L}$).

This gives us our final, practical loss function:

$$\text{Loss} = -\mathcal{L} = -E_{q_{\phi}(\mathbf{z} | \mathbf{x})}[\log p_{\theta}(\mathbf{x} | \mathbf{z})] + D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}) || p_{\theta}(\mathbf{z}))$$

This is the function you actually implement. Minimizing this loss achieves both of our goals:

1. It minimizes the Reconstruction Loss (which is the same as maximizing the log-likelihood).
2. It minimizes the KL Divergence, forcing the encoder to match the prior.

Modern PyTorch VAE Implementation

Now that we understand the VAE architecture and objective, let’s implement a modern VAE in PyTorch. I’ll focus primarily on the model and loss function here, though the full code is available on GitHub.

My VAE implementation uses an output dataclass and a VAE class extending nn.Module.

"""Variational Autoencoder (VAE) model implementation."""

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


def get_activation(activation: str) -> nn.Module:
    """Get activation function by name."""
    activation_lower = activation.lower()
    ACTIVATION_MAP = {
        "relu": nn.ReLU(),
        "tanh": nn.Tanh(),
        "sigmoid": nn.Sigmoid(),
        "leaky_relu": nn.LeakyReLU(),
        "elu": nn.ELU(),
        "gelu": nn.GELU(),
    }
    if activation_lower not in ACTIVATION_MAP:
        supported = ", ".join(ACTIVATION_MAP.keys())
        raise ValueError(
            f"Unsupported activation '{activation}'. Supported: {supported}"
        )
    return ACTIVATION_MAP[activation_lower]


@dataclass
class VAEConfig:
    """VAE model configuration specifying architecture and behavior."""

    hidden_dim: int
    latent_dim: int

    input_shape: tuple[int, int, int] = (1, 28, 28)  # Default: MNIST
    activation: str = "tanh"  # Default: tanh, what was used in the original VAE paper
    use_softplus_std: bool = False  # Whether to use softplus for std parameterization
    n_samples: int = 1  # Number of latent samples per input during training


@dataclass
class VAEOutput:
    """VAE forward pass output containing all relevant tensors and optional losses."""

    x_logits: torch.Tensor
    z: torch.Tensor
    mu: torch.Tensor
    std: torch.Tensor

    x_recon: torch.Tensor | None = None
    loss: torch.Tensor | None = None
    loss_recon: torch.Tensor | None = None
    loss_kl: torch.Tensor | None = None


class VAE(nn.Module):
    """Variational Autoencoder with support for deterministic and probabilistic reconstruction."""

    DEFAULT_EPS = 1e-8

    def __init__(self, config: VAEConfig) -> None:
        """Initialize VAE with given configuration.

        Args:
            config: VAE configuration specifying architecture and behavior
        """
        super().__init__()
        self.config = config

        # Build encoder: input -> hidden -> latent parameters (mu, sigma)
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(
                int(torch.prod(torch.tensor(config.input_shape))), config.hidden_dim
            ),
            get_activation(config.activation),
            nn.Linear(config.hidden_dim, config.latent_dim * 2),
        )

        # Build decoder: latent -> hidden -> reconstructed input
        self.decoder = nn.Sequential(
            nn.Linear(config.latent_dim, config.hidden_dim),
            get_activation(config.activation),
            nn.Linear(
                config.hidden_dim, int(torch.prod(torch.tensor(config.input_shape)))
            ),
            nn.Unflatten(1, config.input_shape),
        )

    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Encode input to latent distribution parameters."""
        encoder_output = self.encoder(x)
        mu, sigma = torch.chunk(encoder_output, 2, dim=-1)
        return mu, sigma

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """Decode latent representation to reconstruction logits"""
        return self.decoder(z)

    def reparameterize(self, mu: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
        """Apply reparameterization trick for differentiable sampling."""
        epsilon = torch.randn_like(std)
        return mu + std * epsilon

    def forward(
        self,
        x: torch.Tensor,
        compute_loss: bool = True,
        reconstruct: bool = False,
        eps: float = DEFAULT_EPS,
    ) -> VAEOutput:
        """Forward pass through the VAE.

        Args:
            x: Input tensor of shape (batch_size, *input_shape)
            compute_loss: Whether to compute VAE loss components
            reconstruct: Whether to return reconstructions or distributions
            eps: Small epsilon value for numerical stability

        Returns:
            VAEOutput containing all relevant tensors and optionally computed losses
        """
        # Prepare input for multiple sampling if needed
        x_expanded = self._expand_for_sampling(x) if self.config.n_samples > 1 else x

        # Encode and sample from latent space
        mu, sigma = self.encode(x)
        std = self._sigma_to_std(sigma, eps=eps)
        mu_expanded, std_expanded = self._expand_latent_params(mu, std)
        z = self.reparameterize(mu_expanded, std_expanded)

        # Decode latent samples
        x_logits = self.decode(z)

        # Create output object
        output = VAEOutput(
            x_logits=x_logits,
            z=z,
            mu=mu,
            std=std,
            x_recon=torch.sigmoid(x_logits) if reconstruct else None,
        )

        # Compute losses if requested
        if compute_loss:
            loss, loss_recon, loss_kl = self._compute_loss(
                x_expanded, x_logits, mu, sigma, std
            )
            output.loss = loss
            output.loss_recon = loss_recon
            output.loss_kl = loss_kl

        return output

    # ==================== Helper Methods ====================

    def _sigma_to_std(
        self, sigma: torch.Tensor, eps: float = DEFAULT_EPS
    ) -> torch.Tensor:
        """Convert sigma parameter to standard deviation."""
        if self.config.use_softplus_std:
            return F.softplus(sigma) + eps
        else:
            return torch.exp(0.5 * sigma)  # sigma represents log-variance

    def _expand_for_sampling(self, x: torch.Tensor) -> torch.Tensor:
        """Expand input tensor for multiple sampling."""
        shape_dims = [1] * len(self.config.input_shape)
        x_expanded = x.unsqueeze(1).repeat(1, self.config.n_samples, *shape_dims)
        return x_expanded.view(-1, *self.config.input_shape)

    def _expand_latent_params(
        self, mu: torch.Tensor, std: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Expand latent parameters for multiple sampling."""
        if self.config.n_samples == 1:
            return mu, std

        mu_expanded = (
            mu.unsqueeze(1)
            .repeat(1, self.config.n_samples, 1)
            .view(-1, self.config.latent_dim)
        )
        std_expanded = (
            std.unsqueeze(1)
            .repeat(1, self.config.n_samples, 1)
            .view(-1, self.config.latent_dim)
        )

        return mu_expanded, std_expanded

    # ==================== Loss Computation ====================

    def _compute_loss(
        self,
        x: torch.Tensor,
        x_logits: torch.Tensor,
        mu: torch.Tensor,
        sigma: torch.Tensor,
        std: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute VAE loss components for deterministic reconstruction."""
        loss_recon = self._compute_reconstruction_loss(x, x_logits)
        loss_kl = self._compute_kl_loss(mu, sigma, std)
        return loss_recon + loss_kl, loss_recon, loss_kl

    def _compute_reconstruction_loss(
        self, x: torch.Tensor, x_logits: torch.Tensor
    ) -> torch.Tensor:
        """Compute reconstruction loss using binary cross-entropy."""
        return F.binary_cross_entropy_with_logits(
            x_logits, x, reduction="sum"
        ) / x.size(0)

    def _compute_kl_loss(
        self,
        mu: torch.Tensor,
        sigma: torch.Tensor,
        std: torch.Tensor,
        eps: float = DEFAULT_EPS,
    ) -> torch.Tensor:
        """Compute KL divergence between latent distribution and standard normal prior."""
        # Analytical KL: KL(N(μ,σ²) || N(0,1)) = 0.5 * Σ(μ² + σ² - 1 - log(σ²))
        if self.config.use_softplus_std:
            # sigma is just the raw output, need to use std directly: σ
            kl_per_sample = 0.5 * torch.sum(
                mu.pow(2) + std.pow(2) - 1 - torch.log(std.pow(2) + eps), dim=1
            )
        else:
            # sigma represents log-variance parameterization: log(σ²)
            kl_per_sample = 0.5 * torch.sum(mu.pow(2) + sigma.exp() - 1 - sigma, dim=1)

        return kl_per_sample.mean()

Loss Scaling

Both components of the VAE loss should be summed over data dimensions and averaged over the batch size. A common mistake is using the reduction="mean" option in PyTorch loss functions, which averages over all elements in the tensor.

• The KL Divergence (loss_kl) is a penalization term. Each dimension of the latent space has the potential to add complexity and deviate from the prior. As you increase the latent dimensionality, you typically see the KL loss increase in magnitude. That’s the cost of having a more expressive latent space.
• The Reconstruction Loss (loss_recon) measures how well the model reconstructs the input data, and it should scale with input dimensionality (this can bias the model toward better reconstruction for higher-dimensional data).

In the case of MNIST, if we used reduction="mean" for BCE, it would be averaged over all $784 \times \text{batch size}$ pixels, making it tiny compared to the KL loss. The KL term would dominate, and the model would learn to ignore the input, potentially leading to posterior collapse.

While modern optimizers can handle a variety of scenarios and you can still learn effective models with imperfect scaling, the original VAE paper used the scaling described above, and I recommend following that convention.

Mitigating Posterior Collapse: KL Annealing/Warmup

One common issue in training VAEs, especially with powerful decoders (like RNNs or deep CNNs), is posterior collapse. This happens when the KL term dominates the loss early in training. The model quickly learns to just output the prior distribution ($q(z|x) \approx p(z)$) to drive the KL loss to zero, effectively ignoring the latent code $z$. The decoder then becomes a powerful autoregressive model that ignores the latent input.

To prevent this, we often use KL Annealing (or Warmup). We introduce a weight $\beta$ for the KL term that starts at 0 and slowly increases to 1 over the first $N$ steps or epochs.

$$ \mathcal{L} = \mathcal{L}_{recon} + \beta \cdot D_{KL} $$

This allows the model to focus purely on reconstruction first (using the full latent capacity), and then slowly adds the regularization pressure.

# Simple Linear Annealing Scheduler
def get_kl_weight(step, total_steps, max_val=1.0):
    val = (step / total_steps) * max_val
    return min(max(val, 0.0), max_val)

# In your training loop:
for epoch in range(epochs):
    beta = get_kl_weight(epoch, warmup_epochs)
    loss = recon_loss + beta * kl_loss

Parameterizing Standard Deviation

    def _sigma_to_std(self, sigma: torch.Tensor, eps: float = DEFAULT_EPS) -> torch.Tensor:
        """Convert sigma parameter to standard deviation."""
        if self.config.bound_std is not None:
            return torch.sigmoid(sigma) * self.config.bound_std + eps
        elif self.config.use_softplus_std:
            return F.softplus(sigma) + eps
        else:
            return torch.exp(0.5 * sigma)  # sigma represents log-variance

Parameterizing the mean of the latent distribution is straightforward since $\mu \in \mathbb{R}$. However, the standard deviation $\sigma$ must be strictly positive (as must the variance $\sigma^2$). This type of constrained optimization is challenging for neural networks.

Log-Variance One common approach is to have the network output the log-variance ($\log \sigma^2$). This is what the original VAE paper did. The idea is to allow the network to output any real number and treat that value as the log-variance, $s = \log \sigma^2$. We can then compute the standard deviation as $\sigma = \exp(0.5 s)$, which is always positive.

The KL divergence formula simplifies nicely with this parameterization:

$$ \text{KL}( \mathcal{N}(\mu, \sigma^2) || \mathcal{N}(0, 1) ) = \frac{1}{2} \sum_{i=1}^d (\mu_i^2 + \sigma_i^2 - 1 - \log \sigma_i^2) $$

0.5 * torch.sum(mu.pow(2) + s.exp() - 1 - s, dim=1)

Softplus Standard Deviation An alternative is to have the network output $\sigma$ directly. This must be handled with care to ensure positivity. Strictly positive activations like softplus are required. Activations like ReLU can output zero, leading to numerical instability (during training) and deterministic behavior (during sampling). Additionally, adding a small epsilon value ensures numerical stability by preventing $\sigma$ from being exactly zero.

The KL divergence formula becomes slightly more complex:

0.5 * torch.sum(
    mu.pow(2) + std.pow(2) - 1 - torch.log(std.pow(2) + eps), dim=1
)

Bounded Standard Deviation Another option is to bound the standard deviation to a maximum value using a sigmoid transformation (or similar). This replaces mapping to $(0, \infty)$ with mapping to $(0, \text{bound})$. This helps prevent extremely high variance values that might destabilize training, while limiting the expressiveness of the latent distribution. Like with softplus, adding a small epsilon ensures numerical stability by preventing $\sigma$ from being exactly zero or approaching it too closely.

Gradient Behavior All parameterizations can work well in practice and have different gradient behaviors. Think of $g(s)$ as a transformation function from the network output to the proper domain of $\sigma$ (or $\sigma^2$); in the log-variance case, $g(s) = \exp(s)$, while in the softplus case, $g(s) = \text{softplus}(s) + \epsilon$.

The gradient of the loss with respect to these outputs can be written using the chain rule:

$$ \frac{\partial \mathcal{L}}{\partial s} = \frac{\partial \mathcal{L}}{\partial \sigma} \cdot \frac{\partial g(s)}{\partial s} $$

where $\frac{\partial g(s)}{\partial s}$ is the derivative of the transformation function.

We need to guard against two pathological cases:

• $\frac{\partial g(s)}{\partial s} \rightarrow 0$: This leads to vanishing gradients, making it hard for the network to learn.
• $\frac{\partial g(s)}{\partial s} \rightarrow \infty$: This leads to exploding gradients, causing instability during training and potentially divergence.

The log-variance parameterization, with its exponential transformation that is its own derivative, exhibits both issues at extreme values. If $s \rightarrow -\infty$, then $\sigma \rightarrow 0$ and the gradient vanishes. If $s \rightarrow \infty$, then $\sigma \rightarrow \infty$ and the gradient explodes. Since the interval $(0, 1)$ is mapped to $(-\infty, 0)$ in log-space, it’s much more difficult for the network to drive $\sigma$ to small values. In practice, exploding gradients at high values have been more problematic in my experience. Gradient clipping, learning rate scheduling, and clamping the log-variance output to a maximum value can help mitigate this.

What about softplus? The derivative of softplus is the sigmoid function, which smoothly maps $(-\infty, \infty)$ to $(0, 1)$. Gradients are always bounded by unity, preventing explosion (barring explosion from other parts of the network). However, as $s \rightarrow -\infty$, the gradient approaches zero, leading to vanishing gradients. Adding a small epsilon helps mitigate this, ensuring that $\sigma$ never gets too close to zero. Nonetheless, learning can still slow down.

For bounded standard deviation, the derivative of the sigmoid function is also bounded, preventing exploding gradients. (The gradient of sigmoid is defined in terms of itself: $\text{sig}’(x) = \text{sig}(x)(1 - \text{sig}(x))$; its maximum value is $0.25$ at $x=0$.)

Experiments

2D MNIST VAE with Different Std. Dev. Parameterizations

First, let’s run an experiment that is close to what was done in the original VAE paper. We’ll use MNIST as our dataset, a simple feedforward architecture with tanh activations, and the log-variance parameterization for the latent distribution.

Some of the differences from the original paper include:

• Using a hidden size of 512 (the original used 500)
• Using the AdamW optimizer (the original used vanilla Adagrad)
• Applying similar weight decay, doing so quite differently due to the optimizer change
• Focusing primarily on 2D latent spaces (for now)

This results in a network with 807,700 parameters. I train each model for 150 epochs at most and highlight the best based on the reconstruction loss on the test set. Just for fun, I sweep across different standard deviation parameterizations and learning rate warmup strategies.

From this summary table, all three parameterizations work well. The differences in final loss values are quite small. This could be due to the simplicity of the dataset and model architecture, further amplified by forcing the network to compress images into a very low-dimensional latent space (2D).

Since the softplus parameterization with learning rate warmup achieved the best reconstruction loss, let’s visualize some of its training dynamics and results more closely.

Loss Dynamics

To understand the VAE’s behavior, we must look at the ELBO and its two components: the Reconstruction Loss and the KL Divergence.

These plots reveal a clear narrative:

1. Rapid Initial Learning: Performance skyrockets in the first ~15 epochs.
2. Overfitting: The Reconstruction Loss (middle) flatlines for the test set while continuing to improve for training, a classic sign of memorization.
3. The Balancing Act: The KL Divergence (bottom) initially rises (“The Cost of Learning”) as the model stretches the latent space to encode digits, then saturates.
4. Equilibrium: The total ELBO (top) improves slowly, driven by the model finding the optimal trade-off between reconstruction and regularization. Notice that Test and Train KL tracks closely: a sign of good regularization!

Visualizing the VAE Trade-Off: BCE vs. KL

While the line plots visualize progress over time, they miss the evolving relationship between our two competing objectives.

A VAE is fundamentally a multi-objective optimization problem. We want to:

1. Minimize Reconstruction Loss (BCE)
2. Minimize KL Divergence

Combining them as the ELBO is common and effective, though it can mask some of the underlying dynamics.

These two goals are in direct conflict. To get perfect reconstruction (BCE = 0), the encoder would need to “memorize” each input, mapping it to a unique, precise point in latent space. This would cause the KL divergence to skyrocket, as these specific, “pointy” distributions are nothing like our smooth N(0, 1) prior.

Conversely, to get perfect KL divergence (KL = 0), the encoder must always output N(0, 1), regardless of the input. This perfectly matches the prior. Since the latent code $\mathbf{z}$ now contains zero information about the input $\mathbf{x}$, the decoder can only learn to output the “average” image, resulting in terrible reconstruction.

The training process is a search for the best compromise.

This plot shows the test set’s BCE (y-axis) vs. KL Divergence (x-axis) at every evaluation step. The color gradient from cool (blue) to warm (red) represents the training progress from Epoch 0 to 150.

Here’s how to interpret this training path:

1. The Start (Green Diamond, ~Epoch 0): The model starts at the top-left.

   • High BCE (Reconstruction): The decoder is random and hasn’t learned to reconstruct anything. Reconstruction is terrible.
   • Low KL Divergence: The encoder is also random. Its output distributions $q_{\phi}(\mathbf{z} | \mathbf{x})$ are a random mess. On average, this “mess” is coincidentally close to the “mess” of the prior $p_{\theta}(\mathbf{z})$, so the KL penalty is low. The model isn’t encoding any useful information yet, so it’s not paying a high price for it.

2. Phase 1: The Initial Plunge (Blue Path): The path moves almost straight down.

   • BCE Plummets: The model’s first and easiest task is to learn to reconstruct something. The optimizer finds massive, easy gains by making the decoder output “blurry digits” to replace the initial noise.
   • KL Stays Low: The model achieves this huge reconstruction win without needing to learn a very complex latent space. It’s the “low-hanging fruit” of training.

3. Phase 2: The Trade-Off (The “Elbow”): The path stops dropping vertically and starts moving to the right and down.

   • “Spending” KL to “Buy” Reconstruction: This is the true VAE trade-off in action. The easy wins are gone. To make the reconstructions sharper and more accurate (lowering BCE further), the model must now learn a more complex, informative latent representation.
   • It “stretches” the latent distributions $q_{\phi}(\mathbf{z} | \mathbf{x})$ to encode more details about each specific digit. This “stretching” moves it further from the simple N(0, 1) prior, and the KL divergence (the “cost”) goes up.

4. The End Game (Red Path & Star): The path settles in the bottom-right corner.

   • Finding the “Elbow”: The model finds an equilibrium. It has pushed the KL divergence as high as it’s “worth” for the reconstruction gains it gets. Trying to get even better reconstruction (moving further down) would cost an enormous, disproportionate amount in KL divergence (moving far to the right), and the total loss would increase.
   • Best Recon (Orange Star): The best reconstruction model (Epoch 118) is found right at this “elbow,” representing the best-found balance point on the trade-off frontier.

This single plot visualizes the entire training dynamic as a journey along the Pareto frontier: the set of optimal solutions where you can’t improve one objective (BCE) without worsening the other (KL).

Generative Performance

Let’s take a look at how well this model can decode samples.

Reconstruction Performance

Immediately, we see a couple of key points:

• Reconstructions are quite blurry compared to the originals. This is expected given the low capacity of the model and the extreme compression into a 2D latent space. General structure is typically preserved, while fine details are lost.
• The network struggles with 4s and 9s, often mixing them up or producing ambiguous shapes. This is a common failure mode in MNIST models due to the similarity of these digits.

Sampling from the Prior

If we sample from the prior N(0, 1) and decode those samples, we get a variety of digit-like images. From this, we get a pretty rich representation of digits. Almost all digits appear to be featured in this random sampling. Again, we see the standard blurriness.

Sweeping the Latent Space

We can select two points at random (here, two zeros), embed them into our latent space and then walk across that latent space to interpolate between two data points. Here, we see a walk that takes us from a zero that is askew to one that is more upright.

Finally, we can also sweep each latent dimension independently to see how they affect the generated images.

1. Sweeping z_1 (top row), we see a 5 become an 8 and then a 9. The slant shifts from left to right as we sweep the dimension.
2. Sweeping z_2 (bottom row), we see a 4 become a 9 and then an 8. Then it becomes a 3, a 2, some nonsense, and a 6.

So clearly each latent dimension is encoding some high-level features of the digits, and we can manipulate those features by moving in latent space.

Inspecting the Latent Space

What does the actual latent space look like?

Even without class information, the network organizes the latent space to encode digit structure effectively. It also becomes immediately apparent why 4s and 9s are so confused by the model. That region is a dense mixture of the two.

We can also look at the marginal distributions of each latent dimension to see how well they match the prior N(0, 1). Here, z_1 is closer to the prior than z_2. z_2 exhibits a bimodal marginal distribution, indicating that the encoder is using this dimension to separate two distinct clusters of data.

We also might want to understand how the log-variance of the latent distributions behaves.

For the most part, we see similar concentration. Some digits are more concentrated than others, though in general the difference is slight.

Beyond 2D: Higher-Dimensional Latent Spaces

What happens as we increase the latent dimensionality? We must do dimensionality reduction to visualize latent spaces, giving us an approximate sense of how the latent space is organized.

As we double the dimensionality, we see a dominant trend at first:

• The reconstruction loss goes down
• The KL loss goes up
• The KL loss per dimension goes down

Something odd happens when we jump from 16 to 32 latent dimensions: some of our latent dimensions become degenerate and stop encoding useful information. This could be an indication we need to choose our hyperparameters a little more cautiously. Perhaps we need a different architecture. Or maybe there is an intrinsic limit to the dimensionality needed for this dataset past which it’s not really helpful to keep scaling the latent dimension.

Training Dynamics

The training dynamics show the battle between reconstruction and KL divergence for different latent dimensionalities. As we increase the latent dimensionality, the oscillation in the KL divergence becomes more pronounced. Particularly chaotic is the $D=16$ case, which struggles to find a stable equilibrium. By the time we expand to $D=32$, the KL penalty seems to overpower the ability to encode information in the latent space, leading to many inactive dimensions. The drop in KL complexity has staircase-like steps without clearly gaining reconstruction ability.

Reconstruction and Generation

As we increase the latent dimensionality, the reconstruction quality improves significantly.

As we increase the dimensionality, we see the increase in quality we’d expect given the reduction in BCE reconstruction loss. In the jump to 4D, we’re able to better resolve the differences between 4s and 9s. Images become much sharper by the time we hit 16 dimensions. The differences between 16 and 32 dimensions, however, are marginal.

Sampling quality also improves with latent dimensionality. Images are sharper as we increase the dimensionality. However, the space seems to get sparser as we increase to the largest dimensionalities, which makes sense given the size and nature of our dataset.

Latent Space Visualizations

The challenge with visualizing higher-dimensional latent spaces is that we must reduce their dimensionality to 2D. PCA struggles to capture the variance of higher dimensionalities. The 4D and 8D plots suggest increasingly better separation of the numeric classes. However, the 16D and 32D plots only show 10-20% of the variance and give a misleading image of overlap.

Conclusion

In this tutorial, we’ve journeyed from the core theory of Variational Autoencoders to a practical, modern PyTorch implementation and a series of experiments on the MNIST dataset. Our findings highlight several key takeaways for practitioners:

1. The VAE is a Balancing Act: The fundamental tension between reconstruction fidelity and latent space regularization is the core of the VAE. Our visualization of the BCE vs. KL loss trade-off clearly showed training as a search for an optimal point on this Pareto frontier, where improving one objective necessarily means sacrificing the other.

2. Latent Dimensionality is a Critical Hyperparameter: Increasing the latent dimension consistently improved reconstruction quality with diminishing returns. As we saw in the jump from 16 to 32 dimensions, too much capacity can lead to “inactive” dimensions, where the KL penalty overpowers the model’s ability to encode useful information. This demonstrates that choosing the right latent size is crucial for both performance and efficiency.

3. VAEs Learn Meaningful Unsupervised Representations: Without any labels, our VAE successfully organized the latent space, clustering similar digits and enabling smooth interpolations. This underscores the power of VAEs for unsupervised learning, dimensionality reduction, and discovering the underlying structure in complex data.

4. Implementation Details Matter: While different standard deviation parameterizations yielded similar results on this simple problem, understanding their gradient behaviors is key for tackling more complex datasets where training stability can be a major challenge. Proper loss scaling is similarly crucial to prevent one term from dominating the other and leading to issues like posterior collapse.

While the classic VAE produces characteristically blurry reconstructions, it remains a foundational generative model. The principles we’ve explored here (the ELBO, the reparameterization trick, and the trade-off between reconstruction and regularization) are central to many more advanced generative models used today.

Questions or feedback? Feel free to reach out. I’d love to hear about your experiences with VAE experiments!

Can you determine a molecule’s shape from mathematical fingerprints alone? This question drives some of the most fundamental challenges in computational chemistry and machine learning. In the broader ML context, this is the classic search for the right inductive bias or invariant representation. Whether we are processing messy documents, natural language, or molecular dynamics, finding a representation that captures essential structure while ignoring irrelevant variations is critical.

I recently encountered a paper with an intriguing title: “Can One Hear the Shape of a Molecule (from its Coulomb Matrix Eigenvalues)?” The title references Mark Kac’s famous mathematical question “Can One Hear the Shape of a Drum?” exploring whether a drum’s shape dictates its sound frequencies.

The molecular version asks: can we determine a molecule’s structure from the eigenvalues of its Coulomb matrix?

Molecular representations are the foundation of machine learning in chemistry. If eigenvalues can capture structural information, they become powerful features for property prediction. Successfully separating simple structural differences is a prerequisite for handling more complex molecules.

The original authors tested this hypothesis using alkane constitutional isomers (molecules with identical formulas but different structural arrangements). I decided to replicate and extend their work to better understand both the methods and their limitations.

In this post, we will explore molecular representation through eigenvalue analysis, covering data generation, unsupervised clustering approaches, and supervised classification methods. I’ll also explore log-transformed Coulomb matrices, which can reveal structural details that standard matrices miss.

Why Alkanes Make Ideal Test Cases

Alkanes are the simplest organic molecules: carbon and hydrogen connected by single bonds with the general formula $C_{n}H_{2n+2}$.

What makes them perfect for testing molecular representations is their constitutional isomers: molecules with identical formulas but different structural arrangements. For small alkanes ($n \leq 3$), atoms can connect in only one way. Starting with butane ($n = 4$), multiple arrangements become possible:

The number of isomers grows rapidly with molecular size. By undecane ($n = 11$), there are 159 different structural arrangements:

This creates a natural classification challenge: can Coulomb matrix eigenvalues distinguish these structural differences? Successfully separating simple alkane isomers is a prerequisite for handling more complex molecules.

Computational Pipeline

The analysis requires three computational steps:

1. Generate constitutional isomers for each alkane formula
2. Create multiple 3D conformations for each isomer
3. Calculate Coulomb matrix eigenvalues for each conformation

Generating Constitutional Isomers

Enumerating all possible carbon skeletons is a combinatorial problem. I used MAYGEN, an open-source Java tool for generating molecular structures from chemical formulas.

For butane ($C_{4}H_{10}$):

java -jar MAYGEN-1.8.jar -v -m -f C4H10 -smi -o butane_conformers.smi

This generates:

CCCC
CC(C)C

The first is n-butane (linear), the second is isobutane (branched). We can automate this across all alkanes:

os.makedirs('isomers', exist_ok=True)
for n in range(1, 12):
    cmd = f"java -jar MAYGEN-1.8.jar -f C{n}H{2*n + 2} -smi -o isomers/C{n}H{2*n + 2}.smi"
    os.system(cmd)

Generating 3D Conformations

For machine learning applications, we need multiple 3D structures of each isomer to capture conformational flexibility. I used RDKit’s ETKDG method, which remains competitive with commercial alternatives:

from rdkit.Chem import AllChem


def smiles_str_to_rdkit_mol(smiles_str: str) -> rdkit.Chem.Mol:
    """Convert a SMILES string to an RDKit mol object.

    Args:
    - smiles_str (str): A SMILES string representing a molecule.

    Returns:
    - mol (rdkit.Chem.Mol): An RDKit mol object representing the molecule.
    """

    # Convert SMILES string to RDKit mol object
    mol = rdkit.Chem.MolFromSmiles(smiles_str)

    # Add hydrogens to the molecule
    mol = rdkit.Chem.AddHs(mol)

    # Assign 3D coordinates to the molecule
    AllChem.EmbedMolecule(mol)

    return mol

Computing Coulomb Matrix Eigenvalues

The Coulomb matrix encodes 3D structure in a rotation and translation-invariant way. Its eigenvalues should capture structural information while remaining invariant to molecular orientation.

First, I wrote a helper function to convert RDKit molecules into ASE Atoms objects that DScribe can process:

def rdkit_mol_to_ase_atoms(rdkit_mol: rdkit.Chem.Mol) -> ase.Atoms:
    """Convert an RDKit molecule to an ASE Atoms object.

    Args:
        rdkit_mol: RDKit molecule object.

    Returns:
        ASE Atoms object.
    """
    ase_atoms = ase.Atoms(
        numbers=[
            atom.GetAtomicNum() for atom in rdkit_mol.GetAtoms()
        ],
        positions=rdkit_mol.GetConformer().GetPositions()
    )
    return ase_atoms

Then I computed Coulomb matrix eigenvalues using DScribe, with optional log transformation:

def ase_atoms_to_coloumb_matrix_eigenvalues(
    ase_atoms: ase.Atoms,
    log: bool = False
) -> np.ndarray:
    """Convert an ASE Atoms object to a Coulomb matrix and calculate its eigenvalues.

    Args:
        ase_atoms: ASE Atoms object.
        log: Whether to log transform the Coulomb matrix prior to calculating the eigenvalues.

    Returns:
        Eigenvalues of the Coulomb matrix.
    """
    # Create a Coulomb matrix
    coulomb_matrix = dscribe.descriptors.CoulombMatrix(
        n_atoms_max=ase_atoms.get_global_number_of_atoms(),
    )

    # Calculate the Coulomb matrix
    coulomb_matrix = coulomb_matrix.create(ase_atoms)
    coulomb_matrix = coulomb_matrix.reshape(
        ase_atoms.get_global_number_of_atoms(),
        ase_atoms.get_global_number_of_atoms())

    if log:
        # Log transform the Coulomb matrix
        coulomb_matrix = np.log(coulomb_matrix)

    # Calculate the eigenvalues of the Coulomb matrix
    eigenvalues = np.linalg.eigvals(coulomb_matrix)
    return eigenvalues

Combining these functions enables efficient data generation:

# Generate 1000 conformations per isomer for each alkane
# gen_n_spectra() combines the above functions to generate multiple conformations
os.makedirs('spectra', exist_ok=True)
n_confs = 1000

for n in range(1, 12):
    print(f'Generating spectra for C{n}H{2*n + 2}')

    with open(f'isomers/C{n}H{2*n + 2}.smi') as f:
        smiles_list = [line.strip() for line in f]

    for i, smiles in enumerate(smiles_list):
        spectra = gen_n_spectra(smiles, n_confs, log=False)
        np.save(f'spectra/C{n}H{2*n + 2}_{i}.npy', spectra)

Reproducing the Original Results

To validate our computational pipeline, I replicated key figures from the original paper. This ensures our implementation correctly captures the phenomena they observed.

We generate data using 1000 conformations per isomer:

os.makedirs('spectra', exist_ok=True)

n_confs = 1000
for n in range(1, 12):
    print(f'Generating spectra for C{n}H{2*n + 2}')
    with open(f'isomers/C{n}H{2*n + 2}.smi') as f:
        lines = f.readlines()
        for i, line in enumerate(lines):
            if not line.strip():
                continue

            smiles = line.strip()
            spectra = gen_n_spectra(n_confs, smiles, log=False)
            np.save(f'spectra/C{n}H{2*n + 2}_{i:03d}.npy', spectra)

After generation, we can load the data into a structured format:

import re
from glob import glob

spectra = {}
for n in range(1, 12):
    spectra[n] = {}
    for f in glob(f'spectra/C{n}H{2*n + 2}_*.npy'):
        j = int(re.search(rf'C{n}H{2*n + 2}_(\d+).npy', f).group(1))
        spectra[n][j] = np.load(f)

Largest Eigenvalues Across Alkane Series

The first analysis examines how the largest Coulomb matrix eigenvalues vary across constitutional isomers for each alkane formula. This plot reveals whether single eigenvalues can distinguish between different molecular formulas.

fig, ax = plt.subplots(1, 1, figsize=(10, 5))

for n in range(1, 12):
    eigenvalues = np.array([spectra[n][i][:, 0].mean() for i in spectra[n]])

    # add black dots to boxplot, but not dirextly to the center line
    jitter = np.random.normal(0, 0.1, size=len(eigenvalues))
    ax.scatter(np.full(len(eigenvalues), n) + jitter, eigenvalues, color='black', s=1, alpha=0.3)

    # Plot median
    ax.scatter([n], [np.median(eigenvalues)], color='red', alpha=0.5)

    # Plot range
    ax.plot([n - 0.5, n + 0.5], [np.min(eigenvalues), np.min(eigenvalues)], 'k-', alpha=0.5)
    ax.plot([n - 0.5, n + 0.5], [np.max(eigenvalues), np.max(eigenvalues)], 'k-', alpha=0.5)

ax.set_xlabel('Molecular formula')
ax.set_ylabel('Largest eigenvalue')
ax.set_xticks(range(1, 12))
ax.set_xticklabels([f'C{n}H{2*n + 2}' for n in range(1, 12)])
ax.set_title('Largest eigenvalues of the Coulomb matrix for alkane constitutional isomers')
plt.savefig('alkane_coulomb_matrix_largest_eigenvalues.webp', bbox_inches='tight')

Our results match the original paper. We observe sub-linear growth in the largest eigenvalue with carbon number, and critically, increasing overlap between isomers as molecules grow larger. The largest eigenvalue alone cannot reliably distinguish constitutional isomers for larger alkanes.

Eigenvalue Distributions for Heptane Isomers

Looking deeper at a specific case, I analyzed the probability density functions for heptane ($C_7H_{16}$) isomers. This molecule has nine constitutional isomers, providing a good test of discrimination power.

for n_sel in range(4, 8):
    fig, ax = plt.subplots(1, 1, figsize=(10, 5))

    smiles = []
    with open(f'isomers/C{n_sel}H{2*n_sel + 2}.smi') as f:
        for line in f:
            smiles.append(line.strip())

    for i in range(len(spectra[n_sel])):
        eigenvalues = spectra[n_sel][i][:, 0]

        # kde plot with the following params
        # - Gaussian kernel
        # - bandwidth with Silverman's rule of thumb
        ax = sns.kdeplot(eigenvalues, bw_method='silverman', label=get_iupac_name(smiles[i]))

    ax.set_xlabel('Largest eigenvalue')
    ax.set_ylabel('Density')
    ax.set_title('PDF of the largest eigenvalue for $C_' + str(n_sel) + 'H_{' + str(2*n_sel + 2) + '}$')
    ax.legend()
    plt.savefig(f'pdf_largest_eigenvalue_C{n_sel}H{2*n_sel + 2}.webp', bbox_inches='tight')
    plt.close()

The pattern is clear:

• n-heptane (linear chain) has the smallest eigenvalues
• 2,2,3-trimethylbutane (highly branched) has the largest
• Seven other isomers fall in between with substantial overlap

This demonstrates the fundamental limitation: while extreme structural differences (linear vs. highly branched) create separable eigenvalue distributions, intermediate structures become indistinguishable.

For smaller alkanes, the separation is more promising:

The progression is clear: eigenvalue-based discrimination works well for small alkanes and degrades as molecular complexity increases.

Two-Dimensional Eigenvalue Space

Can we improve discrimination by using multiple eigenvalues? For butane, plotting the first two eigenvalues reveals interesting structure:

fig, ax = plt.subplots(1, 1, figsize=(10, 5))

n_sel = 4
smiles = []
with open(f'isomers/C{n_sel}H{2*n_sel + 2}.smi') as f:
    for line in f:
        smiles.append(line.strip())

for i in range(len(spectra[n_sel])):
    eigenvalues = spectra[n_sel][i][:, :2]

    ax.scatter(eigenvalues[:, 0], eigenvalues[:, 1], label=get_iupac_name(smiles[i]))

ax.set_xlabel('Largest eigenvalue')
ax.set_ylabel('Second largest eigenvalue')
ax.set_title('2D plot of the first two eigenvalues for $C_4H_{10}$ conformers')
ax.legend()
plt.savefig(f'2d_largest_eigenvalue_C{n_sel}H{2*n_sel + 2}.webp', bbox_inches='tight')

The two isomers cluster distinctly, demonstrating that multi-dimensional eigenvalue features can achieve perfect separation for simple cases. The outlier point in the lower right likely results from conformational sampling noise.

Dimensionality and Information Content

How many eigenvalues do we actually need? Principal component analysis reveals the effective dimensionality of the eigenvalue representations:

from sklearn.decomposition import PCA

fig, ax = plt.subplots(1, 1, figsize=(10, 5))

n_components = {}
for n in range(1, 12):
    n_components[n] = []

    for i in range(len(spectra[n])):
        eigenvalues = spectra[n][i]

        # PCA
        pca = PCA(n_components=0.99, svd_solver='full', whiten=False, random_state=42)
        pca.fit(eigenvalues)
        n_components[n].append(pca.n_components_)

    ax.scatter([n] * len(n_components[n]), n_components[n], alpha=0.3)
    ax.plot([n - 0.25, n + 0.25], [np.mean(n_components[n]), np.mean(n_components[n])], 'k-', alpha=0.5)  # Draw a line for the mean

ax.plot([1, 11], [1, 11], 'k--', alpha=0.5, label='y = num carbon')
ax.plot([1, 11], [5, 35], 'r--', alpha=0.5, label='y = num atoms')

ax.set_xlabel('Number of carbon atoms')
ax.set_ylabel('Number of principal components')
ax.set_title('99% variance explained by number of principal components')
plt.legend()
plt.savefig('99_variance_explained.webp', bbox_inches='tight')

Key observations:

• High compressibility: Far fewer than the $3n+2$ eigenvalue dimensions (one per atom) are needed
• Linear scaling: Principal components grow roughly linearly with carbon number
• Efficient representation: The eigenvalue space has lower effective dimensionality than expected

This suggests the representations are highly correlated, enabling significant dimensionality reduction without information loss.

Log-Transformed Coulomb Matrices

As explored in our previous post on Coulomb matrices, log transformation can reveal different structural information. Standard Coulomb matrices emphasize heavy atom interactions. Log transformation expands the influence of hydrogen atoms by mapping magnitudes in $[0,1]$ to $[-\infty,0]$.

I generated equivalent datasets using log-transformed matrices to test how this affects discriminative power.

os.makedirs('spectra', exist_ok=True)

n_confs = 1000
for n in range(1, 12):
    print(f'Generating spectra for C{n}H{2*n + 2}')
    with open(f'isomers/C{n}H{2*n + 2}.smi') as f:
        lines = f.readlines()
        print(f'\tNumber of SMILES strings: {len(lines)}')
        for i, line in enumerate(tqdm(lines)):
            print(f'\t\t{i + 1}/{len(lines)} - {line.strip()}')

            if not line.strip():
                continue

            if os.path.exists(f'spectra/log-C{n}H{2*n + 2}_{i}.npy'):
                continue

            smiles = line.strip()
            spectra = gen_n_spectra(n=n_confs, smiles_str=smiles, log=True)
            np.save(f'spectra/log-C{n}H{2*n + 2}_{i:03d}.npy', spectra)

Log-Transformed Eigenvalue Distributions

The log transformation substantially changes the eigenvalue landscape:

Log-transformed versions exhibit distinct characteristics:

• Span large negative values due to hydrogen atom emphasis
• Show increasing variance between isomers as molecular size grows
• Demonstrate greater discrimination potential for some isomers

This comes with trade-offs. The distributions can become significantly broader, as seen in the heptane analysis:

The log scale on the y-axis is necessary because unbranched isomers become nearly invisible due to the highly concentrated distributions of branched isomers.

Two-Dimensional Log-Transformed Space

The 2D eigenvalue plot for log-transformed butane shows similar clustering behavior:

The transformation reduces the separation distance, yet linear discriminability remains intact for this simple case.

Dimensionality of Log-Transformed Features

Principal component analysis of log-transformed eigenvalues reveals similar compression properties:

The log-transformed features show comparable dimensionality reduction with marginally higher component requirements.

Testing Eigenvalue Separability

Our exploratory analysis revealed concerning patterns that hint at fundamental limitations: high correlation between eigenvalue dimensions, rapid dimensionality compression via PCA, and overlapping distributions for larger molecules ($n \geq 6$).

These findings leave the question open. We now test eigenvalues directly to see if they can actually separate constitutional isomers without supervision.

We’ll use two complementary clustering metrics to measure how well eigenvalues separate constitutional isomers. This is a fair test. We only compare isomers with identical molecular formulas, keeping eigenvalue dimensions constant.

Dunn Index: Global Cluster Quality

The Dunn Index provides a single metric capturing cluster quality. It asks: “Are the closest different clusters still farther apart than the most spread-out individual cluster?”

$$ \text{Dunn Index} = \frac{\text{smallest distance between different clusters}}{\text{largest diameter within any cluster}} $$

Higher values indicate better separation. When it approaches zero, clusters become indistinguishable, exactly what we suspected from the overlapping eigenvalue distributions observed earlier.

Computing the Dunn Index for each alkane series:

import time

dunn_scores = {}
for n in range(4, 12):
    tik = time.time()
    dunn_scores[n] = dunn_index([spectra[n][i] for i in spectra[n]])
    tok = time.time()
    dunn_scores[n]['time'] = tok - tik
    print(f'C{n}H{2*n + 2}:', dunn_scores[n])

C4H10: {'diameter': 21.43072917950398, 'distance': 8.316362440688767, 'dunn_index': 0.3880578383978837, 'time': 0.06010293960571289}
C5H12: {'diameter': 23.449286379564892, 'distance': 2.4693042873545856, 'dunn_index': 0.10530402705587172, 'time': 0.10832405090332031}
C6H14: {'diameter': 19.602363375467938, 'distance': 1.4477574259511048, 'dunn_index': 0.07385626917634591, 'time': 0.28030991554260254}
C7H16: {'diameter': 20.065014927470955, 'distance': 0.4050094394280803, 'dunn_index': 0.02018485612355977, 'time': 1.0307331085205078}
C8H18: {'diameter': 24.794154667613665, 'distance': 0.5013450168168625, 'dunn_index': 0.020220290771668196, 'time': 4.199508905410767}
C9H20: {'diameter': 21.811025941686033, 'distance': 0.34381162248560415, 'dunn_index': 0.01576320267578513, 'time': 17.400264978408813}
C10H22: {'diameter': 27.180773716656066, 'distance': 0.4986608768730121, 'dunn_index': 0.0183460883811206, 'time': 86.00787401199341}
C11H24: {'diameter': 25.58731511020692, 'distance': 0.5490373275460223, 'dunn_index': 0.021457402825629343, 'time': 424.4431610107422}

The computation time grows dramatically, over 7 minutes for $C_{11}H_{24}$, due to quadratic scaling with the number of isomers (159 isomers requiring ~12,000 pairwise comparisons).

The results reveal a clear trend:

fig, axs = plt.subplots(2, 2, figsize=(15, 10))

# in axs[0, 0] - diameter vs number of carbon atoms
axs[0, 0].plot(
    list(range(4, 12)),
    [dunn_scores[n]['diameter'] for n in range(4, 12)],
    marker='o'
)
axs[0, 0].set_xlabel('Number of carbon atoms')
axs[0, 0].set_ylabel('Diameter')
axs[0, 0].set_title('Diameter vs number of carbon atoms')

# in axs[0, 1] - distance vs number of carbon atoms
axs[0, 1].plot(
    list(range(4, 12)),
    [dunn_scores[n]['distance'] for n in range(4, 12)],
    marker='o'
)
axs[0, 1].set_xlabel('Number of carbon atoms')
axs[0, 1].set_ylabel('Distance')
axs[0, 1].set_title('Distance vs number of carbon atoms')

# in axs[1, 0] - dunn index vs number of carbon atoms
axs[1, 0].plot(
    list(range(4, 12)),
    [dunn_scores[n]['dunn_index'] for n in range(4, 12)],
    marker='o'
)
axs[1, 0].set_xlabel('Number of carbon atoms')
axs[1, 0].set_ylabel('Dunn index')
axs[1, 0].set_title('Dunn index vs number of carbon atoms')

# in axs[1, 1] - time vs number of carbon atoms
axs[1, 1].plot(
    list(range(4, 12)),
    [dunn_scores[n]['time'] for n in range(4, 12)],
    marker='o'
)
axs[1, 1].set_xlabel('Number of carbon atoms')
axs[1, 1].set_ylabel('Time (s)')
axs[1, 1].set_title('Time vs number of carbon atoms')

plt.tight_layout()
plt.savefig('dunn_index_vs_num_carbon_atoms.webp', bbox_inches='tight')

The trend confirms our earlier concerns:

• $C_{4}H_{10}$: Excellent separation (Dunn Index = 0.39) between butane and isobutane
• $C_{5}H_{12}$ to $C_{6}H_{14}$: Rapid decline in separability
• $C_{7}H_{16}$ and beyond: Poor separation (Dunn Index $\approx$ 0.02)

This validates our computational pipeline and matches the original paper’s findings. For larger molecules, eigenvalue clusters become nearly indistinguishable, confirming the overlapping distributions we observed earlier.

Silhouette Analysis: Individual Conformation Assessment

The Dunn Index provides the global view. We must also consider individual molecules. The silhouette score evaluates each conformation separately, asking: “Is this molecule closer to its own isomer family or to a different one?”

For each molecular conformation $i$:

$$ s(i) = \frac{b(i) - a(i)}{\max(a(i), b(i))} $$

where:

• $a(i)$ = average distance to other conformations of the same isomer
• $b(i)$ = average distance to conformations of the nearest different isomer

Interpretation:

• Score near +1: Conformation clusters correctly (good clustering)
• Score near -1: Conformation closer to different isomer (misclassification)

This enables two critical measurements:

1. How many isomers have any misclassified conformations?
2. What fraction of individual conformations get misclassified?

from sklearn.metrics import silhouette_samples
from tqdm import tqdm

s_scores = {}
for n in tqdm(range(4, 12)):
    X = []
    y = []
    for i in spectra[n]:
        X.append(spectra[n][i])
        y.extend(np.full(spectra[n][i].shape[0], i))
    X = np.concatenate(X)
    y = np.array(y)

    s_scores[n] = silhouette_samples(X, y)

Computing both clustering quality metrics:

# Metric 1: Fraction of isomers with ANY negative scores
neg_iso = {}
for n in range(4, 12):
    n_iso = s_scores[n].shape[0] // 1000
    n_has_neg = 0
    for i in range(n_iso):
        chunk = s_scores[n][i * 1000:(i + 1) * 1000]
        if np.any(chunk < 0):
            n_has_neg += 1
    neg_iso[n] = n_has_neg / n_iso

# Metric 2: Individual conformation misclassification rates
neg_confs = {}
for n in range(4, 12):
    n_iso = s_scores[n].shape[0] // 1000
    neg_confs[n] = np.zeros(n_iso)
    for i in range(n_iso):
        isomer_scores = s_scores[n][i * 1000:(i + 1) * 1000]
        neg_confs[n][i] = np.sum(isomer_scores < 0) / isomer_scores.shape[0]

Isomer-Level Analysis

The trend is concerning: by $C_{11}H_{24}$, 97% of isomers have at least one conformation that would be misclassified. This metric is deliberately strict. Even a single misplaced conformation marks the entire isomer as problematic.

Conformation-Level Analysis

The individual analysis reveals dramatic variation:

• $C_{4}H_{10}$: Perfect clustering (0% misclassification), confirming our earlier 2D separation plots
• $C_{5}H_{12}$ to $C_{6}H_{14}$: Modest problems (1-8% misclassification rates)
• $C_{11}H_{24}$: Average 35% conformations misclassified per isomer

Some isomers experience up to 99.5% conformation misclassification (they become essentially unrecognizable in eigenvalue space). This directly connects to our earlier observation: mathematical representations that appear elegant may lack the structural nuances needed for practical discrimination.

Supervised Learning: Finding Hidden Structure

Both clustering metrics deliver the same conclusion: Coulomb matrix eigenvalues alone struggle to reliably distinguish constitutional isomers for larger alkanes. The mathematical elegance of eigenvalues encounters practical limitations as molecular complexity increases.

Supervised learning offers an alternative approach. Providing labels allows models to extract hidden patterns that elude clustering algorithms. The mathematical structure often requires explicit guidance for discovery.

I’ll focus on two baseline approaches: k-nearest neighbors and logistic regression. These represent fundamentally different learning paradigms (one memorizes patterns, the other learns linear boundaries) giving us insight into what types of structure might exist in eigenvalue space.

k-Nearest Neighbors: Pattern Recognition Through Memory

k-NN represents the simplest supervised learning approach: it stores all training examples and classifies new samples based on their closest neighbors. If eigenvalue patterns truly distinguish isomers, nearby points in eigenvalue space should belong to the same class.

This directly tests the local structure. Local neighborhoods often preserve meaningful distinctions even when global structure appears diffuse.

Testing Different Feature Representations

We compare three approaches: full eigenvalue vectors, top 10 eigenvalues only, and PCA-reduced representations.

Testing 1-nearest neighbor with full dimensionality:

from sklearn.model_selection import cross_val_score, StratifiedKFold
from sklearn.neighbors import KNeighborsClassifier

df_1nn = []

for n in range(4, 12):

    # Prepare the data for CnH2n+2
    X, y = prep_data(n=n)

    # Create knn classifier
    knn = KNeighborsClassifier(n_neighbors=1)

    # Set up stratified 5-fold cross-validation
    cv = StratifiedKFold(n_splits=5)

    # Perform cross-validation. Since 'cross_val_score' computes accuracy, we compute misclassification rate by subtracting accuracy from 1.
    acc_scores = cross_val_score(knn, X, y, cv=cv, scoring='accuracy')

    # Convert accuracy scores to misclassification error rates
    misclassification_error_rates = 1 - acc_scores

    # Calculate the average and standard deviation of the misclassification error rates
    avg_misclassification_error = np.mean(misclassification_error_rates)
    std_misclassification_error = np.std(misclassification_error_rates)

    print(f'C{n}H{2*n + 2}: {avg_misclassification_error:.2%} ± {std_misclassification_error:.2%}')

    df_1nn.append({
        'molecule': f'C{n}H{2*n + 2}',
        'avg_misclassification_error': avg_misclassification_error,
        'std_misclassification_error': std_misclassification_error,
        'n': n,
        'representation': 'full',
        'model': '1nn',
    })

The results are remarkable compared to unsupervised clustering:

C4H10: 0.00% ± 0.00%
C5H12: 0.00% ± 0.00%
C6H14: 0.00% ± 0.00%
C7H16: 0.07% ± 0.05%
C8H18: 0.11% ± 0.05%
C9H20: 0.51% ± 0.09%
C10H22: 1.31% ± 0.09%
C11H24: 3.24% ± 0.09%

Perfect classification for molecules up to $C_{6}H_{14}$, with low error rates even for $C_{11}H_{24}$ (3.24%). This is a large improvement over clustering, where 97% of $C_{11}H_{24}$ isomers had misclassified conformations.

Note on feature scaling: Standardizing features significantly degraded performance; eigenvalue magnitudes carry crucial structural information.

Comparing performance across different feature representations:

Key insights:

• Representation choice matters little for 1-NN. Full, top-10, and PCA representations perform nearly identically
• PCA slightly outperforms top-10 eigenvalues for larger molecules, capturing more structural variance
• Perfect classification persists through $C_{6}H_{14}$ regardless of representation

This confirms that discriminative information concentrates in the largest eigenvalues, validating our earlier PCA findings.

The Neighbor Count Effect

Testing k-NN with different neighbor counts (k=1, 3, 5) reveals a counterintuitive pattern:

Why does performance degrade with more neighbors? This connects directly to our earlier clustering analysis. The eigenvalue space lacks meaningful local structure. When k-NN examines beyond the immediate nearest neighbor, it increasingly finds examples from different classes.

This validates our unsupervised findings: in the absence of clear cluster boundaries, examining more neighbors introduces noise.

Logistic Regression: Learning Linear Decision Boundaries

Logistic regression represents a fundamentally different approach. Logistic regression learns linear decision boundaries in eigenvalue space. If eigenvalues encode structural information linearly, this should work well.

We’ll focus on PCA-reduced representations to keep computation manageable, using insights from the k-NN analysis.

from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.decomposition import PCA

df_lr = []

for n in range(4, 12):
    # Prepare the data for CnH2n+2
    X, y = prep_data(n=n)

    # Create logistic regression classifier with PCA
    lr = Pipeline([
        ('pca', PCA(n_components=10)),
        ('lr', LogisticRegression(
            max_iter=10_000,
            penalty='l2',
            solver='lbfgs',
            C=10.0,  # Reduced regularization
            random_state=42,
            n_jobs=-1,
        ))
    ])

    # 5-fold stratified cross-validation
    cv = StratifiedKFold(n_splits=5)
    acc_scores = cross_val_score(lr, X, y, cv=cv, scoring='accuracy')

    # Convert to misclassification rates
    avg_error = np.mean(1 - acc_scores)
    std_error = np.std(1 - acc_scores)

    print(f'C{n}H{2*n + 2}: {avg_error:.2%} ± {std_error:.2%}')

Comparing k-NN versus logistic regression performance:

Key observations:

• k-NN dominates across all molecular sizes
• Linear boundaries fail for larger molecules. This suggests nonlinear eigenvalue relationships.
• Performance gap grows with molecular complexity, indicating increasingly nonlinear structural patterns

Logistic regression’s performance indicates that discriminative patterns in eigenvalue space are fundamentally nonlinear. Capturing these complex relationships requires memory-based or non-linear approaches.

Implications for Molecular Representation

Our supervised learning experiments reveal a nuanced picture of Coulomb matrix eigenvalues as molecular descriptors. Eigenvalues preserve sufficient local structure for nearest-neighbor classification to work remarkably well, despite lacking clean global clusters.

This analysis reveals important lessons about molecular representations:

1. Empirical performance and mathematical elegance are separate axes: an elegant descriptor can still fail in practice as the space gets more complex.
2. Context matters: Representations exhibit distinct performance characteristics under supervised versus unsupervised conditions.
3. Molecular complexity is challenging: Even simple alkanes test our best descriptors.
4. Local vs. global structure: Local neighborhood structures often contain highly discriminative information.

For practitioners working with molecular representations, it is crucial to test multiple learning paradigms. Supervised and unsupervised approaches often yield different insights. Furthermore, logistic regression’s poor performance indicates that discriminative patterns in eigenvalue space are fundamentally nonlinear. Capturing these complex relationships requires memory-based or non-linear approaches.

Why this matters beyond the alkane case: molecular representations are the input layer for property prediction, and understanding where a simple descriptor like Coulomb-matrix eigenvalues fails (overlapping clusters for larger molecules, nonlinear class structure that defeats logistic regression) is what motivates moving to graph- or coordinate-aware models. The failure modes here are the argument for richer representations.

The data pipeline that generated the datasets used in this analysis is available at the Synthetic Isomer Data Generation Pipeline project page.

When working with machine learning in chemistry, one of the first challenges you encounter is how to represent molecules in a way that algorithms can understand. You can’t just feed raw atomic coordinates into a model. The representation needs to be invariant to rotation, translation, and atom ordering, since these operations don’t change the molecule’s fundamental properties.

The Coulomb matrix, introduced by Rupp et al. in 2012 [1], provides a straightforward solution to this problem. While newer methods have largely superseded it for practical applications, the Coulomb matrix remains an excellent starting point for understanding how molecular descriptors work.

The key insight is simple: we encode pairwise relationships between atoms in a way that captures the essential physics while maintaining the required invariances.

The Coulomb Matrix: Theory and Intuition

The Coulomb matrix encodes molecular structure in a symmetric $N \times N$ matrix, where $N$ is the number of atoms. Each element $C_{ij}$ is defined as:

$$ C_{ij} = \begin{cases} 0.5 Z_i^{2.4} & \text{if } i = j, \\ \frac{Z_i Z_j}{|\mathbf{R}_i - \mathbf{R}_j|} & \text{if } i \neq j, \end{cases} $$

Here, $Z_i$ is the atomic number of atom $i$, and $\mathbf{R}_i$ is its position in 3D space. The diagonal elements ($0.5 Z_i^{2.4}$) represent atomic self-energies, derived from fitting atomic numbers to experimental data. The off-diagonal elements mimic Coulombic interactions between atoms. They’re inversely proportional to distance, just like electrostatic potential energy [3].

This construction gives us several useful properties:

• Rotation and translation invariant: Only relative distances matter
• Symmetric: $C_{ij} = C_{ji}$, which is physically sensible
• Size-extensive: Larger molecules have larger matrix elements
• Captures 3D structure: Nearby atoms have larger interaction terms

While more sophisticated methods exist today [2], the Coulomb matrix’s simplicity makes it ideal for understanding the fundamentals of molecular representation.

Hands-on Example: Bicyclobutane

Let’s calculate the Coulomb matrix for bicyclobutane, a strained but stable bicyclic system (bicyclo[1.1.0]butane, C4H6, two cis-fused cyclopropane rings). This example will show you exactly how the theory translates to practice.

I’ll use Python with the Atomic Simulation Environment (ase) for molecular structure [4] and dscribe for the Coulomb matrix calculation [2]:

from ase.build import molecule
from ase.visualize import view

# Load the bicyclobutane structure
bicyclobutane = molecule('bicyclobutane')

# Optional: visualize the structure
view(bicyclobutane, viewer='x3d')

Now we calculate the Coulomb matrix using DScribe:

from dscribe.descriptors import CoulombMatrix

# Set up the descriptor
cm = CoulombMatrix(n_atoms_max=len(bicyclobutane))

# Calculate and reshape into matrix form
cm_bicyclobutane = cm.create(bicyclobutane)
cm_bicyclobutane = cm_bicyclobutane.reshape(len(bicyclobutane), len(bicyclobutane))

Visualizing the Results

The Coulomb matrix can be visualized as a heatmap. Let’s look at both the raw matrix and its logarithm:

import matplotlib.pyplot as plt
import numpy as np

# Raw Coulomb matrix
plt.figure(figsize=(8, 8), dpi=150)
plt.imshow(cm_bicyclobutane, cmap='coolwarm')
plt.colorbar(label='Magnitude')
plt.title('Coulomb Matrix for Bicyclobutane')
plt.show()

The raw matrix shows clear patterns:

• Large diagonal elements: Carbon atoms (Z=6) dominate due to their higher atomic numbers
• Smaller off-diagonal elements: Represent pairwise interactions
• Minimal hydrogen contribution: Hydrogen atoms (Z=1) have much smaller values

For better visualization of the structure, the logarithm reveals more detail:

plt.figure(figsize=(8, 8), dpi=150)
plt.imshow(np.log(cm_bicyclobutane), cmap='coolwarm')
plt.colorbar(label='log(Magnitude)')
plt.title('Log Coulomb Matrix for Bicyclobutane')
plt.show()

Eigenvalue Analysis

The eigenvalues of the Coulomb matrix provide another perspective on molecular structure:

These eigenvalues are often used as features themselves, providing a more compact representation than the full matrix.

Practical Limitations

The Coulomb matrix has significant limitations that explain why it’s been largely superseded by modern methods. Understanding these constraints is crucial for knowing when and how to use this descriptor.

The Size Problem

Every molecule must be represented by the same size matrix, which creates several issues:

• Padding overhead: Small molecules get padded with zeros up to the maximum size
• Quadratic scaling: An $N$-atom molecule requires $N^2$ features
• Fixed maximum size: You can’t represent molecules larger than your preset limit
• Inefficient storage: Most elements are zero for small molecules in large matrices

For a dataset ranging from 5-atom to 50-atom molecules, every molecule needs a 50x50 matrix. That’s 2,500 features, most of which are zero for smaller molecules.

Permutation Sensitivity

Despite being called “invariant,” the Coulomb matrix can actually change if you reorder the atoms in your input file. The standard solution is to sort atoms by the L2 norm of their matrix rows, but this introduces its own problems:

• Symmetry breaking: Equivalent atoms might be ordered differently
• Numerical instability: Small coordinate changes can flip the ordering
• Loss of chemical intuition: The sorted order doesn’t reflect meaningful chemistry

Interestingly, some studies suggest that adding controlled noise to create multiple permutations can actually improve machine learning performance [5].

Limited Scope

The Coulomb matrix works well only for specific types of systems:

• Small molecules: Performance degrades for large systems due to size scaling
• Gas-phase: Not suitable for periodic systems like crystals or surfaces
• Single conformations: Each 3D structure gets its own matrix
• Non-reactive: Doesn’t capture bond-breaking or formation

For periodic systems, you’d need specialized variants like the Ewald sum matrix [6].

Why Learn It Anyway?

Given these limitations, why spend time understanding the Coulomb matrix? Several reasons:

Educational value: It’s conceptually straightforward and provides excellent intuition for how molecular descriptors work. The mathematical formulation is simple enough to implement from scratch.

Historical importance: Many subsequent methods build on ideas first explored with Coulomb matrices. Understanding this foundation helps you appreciate why newer methods were developed.

Benchmarking: It remains useful as a baseline method for comparing new descriptors on small molecular datasets.

Proof of concept: For exploratory work on small, well-defined datasets, the Coulomb matrix can still provide quick insights.

If you’re working on practical problems with larger datasets or diverse molecular sizes, consider modern alternatives like graph neural networks, descriptors from DScribe’s extended library, or learned representations from transformer models.

Putting It in Context

To see the Coulomb matrix applied to real problems, I’ve written a detailed guide using it for molecular classification:

• Coulomb Matrix Eigenvalues: Can You Hear the Shape of a Molecule?: A comprehensive analysis of alkane isomers, from unsupervised clustering limits to supervised classification successes.

For comparison with modern approaches, check out my post on 3D conformer generation with the GEOM dataset, which showcases more sophisticated molecular representations. For technical specifications and benchmarks, see the GEOM dataset card.

The Coulomb matrix may be dated, but it remains an excellent entry point into the world of molecular machine learning. Once you understand its strengths and limitations, you’ll be better equipped to appreciate why the field has moved toward more sophisticated approaches.

Have questions about molecular descriptors or want to discuss other approaches to molecular machine learning? I’d be happy to explore these topics further.

References

• [1]: M. Rupp, A. Tkatchenko, K.-R. Müller, and O. A. von Lilienfeld, “Fast and Accurate Modeling of Molecular Atomization Energies with Machine Learning,” Physical Review Letters, 108(5), 058301 (2012). https://doi.org/10.1103/PhysRevLett.108.058301 arXiv:1109.2618
• [2] L. Himanen, M. O. J. Jäger, E. V. Morooka, F. F. Canova, Y. S. Ranawat, D. Z. Gao, P. Rinke, and A. S. Foster, “DScribe: Library of descriptors for machine learning in materials science,” Computer Physics Communications, 247, 106949 (2020). https://doi.org/10.1016/j.cpc.2019.106949 arXiv:1904.08875
• [3] J. Schrier, “Can one hear the shape of a molecule (from its Coulomb matrix eigenvalues)?,” Journal of Chemical Information and Modeling, 60(8), 3804-3811 (2020). https://doi.org/10.1021/acs.jcim.0c00631
• [4] A. H. Larsen, J. J. Mortensen, J. Blomqvist, I. E. Castelli, R. Christensen, M. Dułak, J. Friis, M. N. Groves, B. Hammer, C. Hargus, E. D. Hermes, P. C. Jennings, P. B. Jensen, J. Kermode, J. R. Kitchin, E. L. Kolsbjerg, J. Kubal, K. Kaasbjerg, S. Lysgaard, J. B. Maronsson, T. Maxson, T. Olsen, L. Pastewka, A. Peterson, C. Rostgaard, J. Schiøtz, O. Schütt, M. Strange, K. S. Thygesen, T. Vegge, L. Vilhelmsen, M. Walter, Z. Zeng, and K. W. Jacobsen, “The Atomic Simulation Environment - A Python library for working with atoms,” J. Phys.: Condens. Matter, 29, 273002 (2017). https://doi.org/10.1088/1361-648X/aa680e documentation
• [5] G. Montavon, K. Hansen, S. Fazli, M. Rupp, F. Biegler, A. Ziehe, A. Tkatchenko, A. Lilienfeld, and K.-R. Müller, “Learning invariant representations of molecules for atomization energy prediction,” Advances in Neural Information Processing Systems, 25 (2012). Available online: https://proceedings.neurips.cc/paper_files/paper/2012/file/115f89503138416a242f40fb7d7f338e-Paper.pdf
• [6] F. Faber, A. Lindmaa, O. A. von Lilienfeld, and R. Armiento, “Crystal structure representations for machine learning models of formation energies,” International Journal of Quantum Chemistry, 115(16), 1094-1101 (2015). https://doi.org/10.1002/qua.24917

In computer vision or scientific computing, a common problem frequently arises: given two sets of points, what is the optimal rigid body transformation for their alignment? The Kabsch algorithm provides a nice solution.

What are some concrete situations where this crops up?

• Molecular Dynamics: Your points are a set of atoms (with physically relevant types), and you want to compare two molecular conformations. Are they the same structure with minor noise or rotation? Or are they different conformations, like a different folding of a protein? This is especially helpful when applying generative models to chemical structures. For example, if you are building a 3D Molecular VAE in PyTorch or working with Flow Matching models, Kabsch alignment ensures your generative loss function remains rotationally invariant.
• Computer Vision: You have two point clouds from 3D scans of an object taken from different angles. You want to align them to reconstruct the full shape. Or perhaps you’re generating 3D shapes from 2D images and need to compare the generated shape to a ground truth scan. Anytime a 3D system is represented as a point cloud, the Kabsch algorithm can help with alignment.

Of course, existing libraries implement this algorithm. However, often I find it beneficial to implement algorithms from scratch to build intuition. Furthermore, modern machine learning applications require automatic differentiation, so we will implement the algorithm in PyTorch, TensorFlow, and JAX.

Below, we’ll cover the math behind the Kabsch algorithm (and its scaling variant, the Kabsch-Umeyama algorithm) and provide complete, differentiable implementations in NumPy, PyTorch, TensorFlow, and JAX, demonstrating both single-pair and batched computations for ML applications.

The Math

Let’s say we have two sets of paired points, $P={\mathbf{p}_i} \in \mathbb{R}^{N \times D}$ and $Q={\mathbf{q}_i} \in \mathbb{R}^{N \times D}$, for $i = 1, \dots, N$ (where $D$ is the dimensionality and $N$ is the number of points). We want to find a translation vector $\mathbf{t}$ and rotation matrix $R$ to transform $P$ to align with $Q$.

The optimization problem is:

$$ \min_{\mathbf{t}, \ R} \mathcal{L}(\mathbf{t}, R) = \frac{1}{2} \sum_{i=1}^N | \mathbf{q}_i - (R\mathbf{p}_i + \mathbf{t}) |^2 $$

where $\mathbf{t}^\ast \in \mathbb{R}^D$ and $R^\ast \in \mathbb{R}^{D \times D}$ are the optimal translation and rotation.

Often we use a weighted version with weights $w_i$ (e.g., atomic masses in molecular dynamics):

$$ \min_{\mathbf{t}, \ R} \mathcal{L}(\mathbf{t}, R) = \frac{1}{2} \sum_{i=1}^N w_i | \mathbf{q}_i - (R\mathbf{p}_i + \mathbf{t}) |^2 $$

The Translation

The translation and rotation are coupled, but they separate cleanly once we work in centroid-centered coordinates. Compute the centroids (averages) of both point sets:

$$ \bar{\mathbf{p}} = \frac{1}{N} \sum_{i=1}^N \mathbf{p}_i \quad \text{and} \quad \bar{\mathbf{q}} = \frac{1}{N} \sum_{i=1}^N \mathbf{q}_i $$

For any fixed rotation $R$, the translation that minimizes $\mathcal{L}$ is found by setting $\partial \mathcal{L} / \partial \mathbf{t} = 0$. It maps the rotated source centroid onto the target centroid:

$$ \mathbf{t} = \bar{\mathbf{q}} - R\bar{\mathbf{p}} $$

A tempting shortcut is to write $\mathbf{t} = \bar{\mathbf{q}} - \bar{\mathbf{p}}$, but that is only correct when $R = I$. In general the translation depends on the rotation, so we compute it after solving for $R$. Substituting this optimal $\mathbf{t}$ back into the objective cancels the centroids and leaves a rotation-only problem in the centered coordinates $\mathbf{p}_i^\prime = \mathbf{p}_i - \bar{\mathbf{p}}$ and $\mathbf{q}_i^\prime = \mathbf{q}_i - \bar{\mathbf{q}}$:

$$ \mathcal{L}(R) = \frac{1}{2} \sum_{i=1}^N | \mathbf{q}_i^\prime - R\mathbf{p}_i^\prime |^2 $$

which is what the next section solves.

The Rotation Matrix

We now minimize $\mathcal{L}(R)$ over rotations, using the centered points $\mathbf{p}_i^\prime$ and $\mathbf{q}_i^\prime$ from above. Compute the cross-covariance matrix between the centered sets:

$$ C = P^{\prime T} Q^\prime = \sum_{i=1}^N \mathbf{p}_i^{\prime T} \mathbf{q}_i^{\prime} \in \mathbb{R}^{D \times D} $$

This is a fairly lightweight operation since $D$ is typically small (e.g., 3 for 3D points), even if $N$ is large.

With $C$ in hand, we want to compute its Singular Value Decomposition (SVD):

$$ C = U \Sigma V^T $$

This operation is computationally expensive. It scales cubically with $D$ (i.e., $O(D^3)$). However, since we’re often interested in cases where $D$ is small (e.g., 2D or 3D points), this is manageable.

Next, we check for improper rotations (i.e., reflections) and correct for them where necessary:

$$ d = \text{sign}(\det(V U^T)) $$

If $d = -1$, we need to flip the last column of $V$ in the final rotation matrix.

Let $B = \text{diag}(1, 1, d)$. The optimal rotation matrix comes out:

$$ R^\ast = V B U^T $$

Summary

In a nutshell, the Kabsch algorithm boils down to:

1. Compute centroids of $P$ and $Q$ ($\bar{\mathbf{p}}$ and $\bar{\mathbf{q}}$)
2. Center both point sets by subtracting centroids: $P^\prime$ and $Q^\prime$
3. Compute cross-covariance matrix $C = P^{\prime T} Q^\prime$
4. Compute SVD: $C = U \Sigma V^T$ (expensive step)
5. Compute $d = \text{sign}(\det(V U^T))$ and $B = \text{diag}(1, 1, d)$
6. Optimal rotation: $R^\ast = V B U^T$
7. Optimal translation (using the rotation from step 6): $\mathbf{t}^\ast = \bar{\mathbf{q}} - R^\ast\bar{\mathbf{p}}$

The resulting root-mean-square deviation (RMSD) between aligned point sets is

$$ \text{RMSD} = \sqrt{\frac{1}{N} \sum_{i=1}^N | \mathbf{q}_i - (R^\ast\mathbf{p}_i + \mathbf{t}^\ast) |^2} $$

which is frequently used as a measure of similarity between molecular structures or as a metric in loss functions for ML applications.

The Kabsch-Umeyama Algorithm (Scaling)

While the standard Kabsch algorithm solves for optimal rotation and translation, the Kabsch-Umeyama algorithm extends this by also finding an optimal scaling factor $c$. This is essential when aligning structures of different scales, such as a 3D scan versus a ground truth model.

(Note: This is sometimes searched for as the “Absch-Umeyama algorithm” due to typos, but the correct attribution is to Shinji Umeyama based on Wolfgang Kabsch’s work.)

The method estimates the transformation $\mathbf{q}_i \approx c R \mathbf{p}_i + \mathbf{t}$. The optimal scale is the trace of the (reflection-corrected) singular values of the cross-covariance divided by the variance of the source points about their centroid. See the Umeyama paper notes for the full derivation.

A Note on SVD and Automatic Differentiation

While modern frameworks allow us to backpropagate through the Singular Value Decomposition (SVD), it comes with a known stability issue: if the cross-covariance matrix has identical (degenerate) singular values (which can occur if the point clouds are perfectly aligned or have certain symmetries), the gradient of the SVD approaches infinity, causing NaN values during backpropagation. If you plan to use this algorithm as a loss function for a neural network, it is often necessary to add a tiny epsilon to the matrix before computing the SVD, or to utilize an SVD gradient patch. The Kabsch-Horn Cookbook library provides a SafeSVD primitive that floors the singular-value-gap denominator at machine epsilon in the backward pass, producing finite gradients at degenerate inputs across PyTorch, JAX, TensorFlow, and MLX.

Implementation

Let’s implement the algorithm in different frameworks. Note that for simplicity, the following implementations cover the unweighted Kabsch algorithm. If your application (like molecular dynamics) requires weights (e.g., atomic masses), the Kabsch-Horn Cookbook library provides per-point weighted alignment out of the box.

NumPy

import numpy as np


def kabsch_numpy(P, Q):
    """
    Computes the optimal rotation and translation to align two sets of points (P -> Q),
    and their RMSD.

    :param P: A Nx3 matrix of points
    :param Q: A Nx3 matrix of points
    :return: A tuple containing the optimal rotation matrix, the optimal
             translation vector, and the RMSD.
    """
    assert P.shape == Q.shape, "Matrix dimensions must match"

    # Compute centroids
    centroid_P = np.mean(P, axis=0)
    centroid_Q = np.mean(Q, axis=0)

    # Center the points
    p = P - centroid_P
    q = Q - centroid_Q

    # Compute the covariance matrix
    H = np.dot(p.T, q)

    # SVD
    U, S, Vt = np.linalg.svd(H)

    # Validate right-handed coordinate system
    if np.linalg.det(np.dot(Vt.T, U.T)) < 0.0:
        Vt[-1, :] *= -1.0

    # Optimal rotation
    R = np.dot(Vt.T, U.T)

    # Optimal translation (depends on R, so computed after it)
    t = centroid_Q - np.dot(R, centroid_P)

    # RMSD
    rmsd = np.sqrt(np.sum(np.square(np.dot(p, R.T) - q)) / P.shape[0])

    return R, t, rmsd

Here’s a quick test to verify correctness:

def test_numpy():
    np.random.seed(12345)

    P = np.random.randn(100, 3)

    alpha = np.random.rand() * 2 * np.pi
    R = np.array([[np.cos(alpha), -np.sin(alpha), 0],
                    [np.sin(alpha), np.cos(alpha), 0],
                    [0, 0, 1]])
    t = np.random.randn(3) * 10

    Q = np.dot(P, R.T) + t

    R_opt, t_opt, rmsd = kabsch_numpy(P, Q)

    print('RMSD: {}'.format(rmsd))
    print('R:\n{}'.format(R))
    print('R_opt:\n{}'.format(R_opt))
    print('t:\n{}'.format(t))
    print('t_opt:\n{}'.format(t_opt))

    l2_t = np.linalg.norm(t - t_opt)
    l2_R = np.linalg.norm(R - R_opt)
    print('l2_t: {}'.format(l2_t))
    print('l2_R: {}'.format(l2_R))