All of that structure is derived from a SMILES string.

Molecules are also drawn. A structure diagram carries the same topology a chemist reads directly off the page, in a form no language model ever sees. So I wanted to know whether MIST’s geometry is reachable from pixels, or whether it is a property of the string.

There are two standard ways to teach one model to land where another one does. You can push negatives apart with a contrastive loss, or you can skip the negatives and regress the target directly. I expected the predictive one to transfer better to real drawings, because a contrastive objective can exploit shortcuts that only exist inside a training batch. I wrote that down before running anything.

It held, then it reversed, and the mechanism I had named for it was wrong in both cases.

What I Built

One vision tower, one frozen text tower, and a single linear map between them.

Arm A optimizes the SigLIP sigmoid pairwise loss. Arm B regresses the frozen target in latent space with a smooth-L1 loss and no negatives, which is the I-JEPA move applied as an alignment head rather than as pretraining. Everything else is held identical: same molecules, same splits, same seed, same schedule, same evaluation, checked by an assertion before any table prints. Every number below is one seed.

Two details in the setup are load-bearing.

The pooling is MIST’s own. Their embedding is the last-layer hidden state at the first token, so anything else measures a different space. Loading the checkpoint with AutoModel also attaches a randomly initialized pooler unless you ask it not to, and pooler_output from that pooler returns plausible noise.

The contrastive labels come from molecular identity rather than batch position. Two rows in a batch can be depictions of the same molecule, and treating them as negatives because they sit at different indices would penalize the model for being right. That detail matters more than it looks, and I will come back to it.

This is not an OCSR contribution. Reading a molecule out of a picture is a mature field, and MolScribe, MolGrapher, DECIMER and Img2Mol all do it well. What is here is a controlled experiment about alignment objectives that happens to run on a domain with a real evaluation surface.

That domain has a catch worth stating early, because it shapes which number is the result. The training depictions are rendered from SMILES strings. A picture drawn from a string carries no information the string did not already have, so the in-distribution task is closer to inverting a renderer than to reading chemistry. Transfer to depictions that came from the world is the part that tests something. That is why the out-of-distribution slice is the headline and the in-distribution number is context.

I came to MIST from the tokenizer side. I had published a controlled comparison of BPE and Unigram-LM over chemistry SMILES a few weeks earlier, and MIST is a SMILES foundation model whose headline contribution is its own tokenizer. I was reading the paper closely for that reason and stayed for the embedding geometry.

What I Registered In Advance

Three predictions, committed before any training run existed:

1. Both arms will be near-ceiling on in-distribution synthetic retrieval, and that number on its own will mean almost nothing.
2. The arms will separate, if at all, on the WildMol-10k OOD slice and the scaffold-split linear probe. That gap is the result.
3. Arm B is expected to hold up better OOD, on the reasoning that regressing a fixed target cannot exploit in-batch shortcuts the way a contrastive objective can. This is a guess and is written down so it can be wrong in public.

I also registered what would count as a clean negative and ship as the result: both arms at or near chance on the OOD slice and the probe, with all baselines behaving correctly. That sentence turns out to matter more than the predictions do.

What Had To Pass First

Four gates run before any headline number is believed, because a comparison between two arms says nothing if the harness scoring both of them is broken.

The image-statistics floor is the one that earns its keep below. It is what a model scores using nothing but coarse pixel statistics, and any headline has to clear it.

A Perfect Model Scored Zero

Raw cosine against MIST measures almost nothing. Its embeddings occupy a narrow cone, where two unrelated molecules sit at cosine 0.517 and a model that ignores its input entirely and emits the training centroid scores 0.720. So every cosine in this project is centered, and each side is centered by its own training-split mean:

$$ \mathrm{cc} = \cos\left(\hat{\mathbf{u}} - \bar{\hat{\mathbf{u}}}_{\text{train}},\ \mathbf{y} - \bar{\mathbf{y}}_{\text{train}}\right) $$

Chance is zero. Both centroids come from the training split, since centering a held-out set by its own mean leaks that set’s geometry into its own score.

The version I shipped first subtracted the target centroid from both sides. That is valid only while the prediction sits at roughly the target’s scale. Arm B’s smooth-L1 pins it there. Arm A’s loss L2-normalizes both sides internally, so it constrains the output norm not at all, and weight decay shrinks it with nothing opposing.

Take an oracle predictor $\hat{\mathbf{u}} = s\mathbf{y}$, whose true cosine is exactly 1 for every $s > 0$:

A perfect model emitting unit-norm vectors scored below chance.

On the real 20K run, that defect scored a working contrastive model at +0.1088, inside noise of the +0.1058 image-statistics floor. Arm A’s head was emitting norm 6.35 against MIST’s 22.28. Under the corrected metric the same checkpoint scores +0.3052.

The uncomfortable part is what that number was. I had registered, in advance, that both arms landing at the floor with every baseline behaving correctly would be a clean negative and would ship as the result. The defect produced precisely that. Every exit criterion passed.

No test caught it. What caught it was asking what a perfect model would score under this metric, which is a question whose answer is knowable before you ask it.

Registering a prediction stops the goalposts moving once data exists. It does nothing at all about a defect that happens to confirm the prediction, and it makes that case more expensive rather than less, because the incentive to accept a number and stop looking is strongest exactly where you said the number would be.

There was a second defect of the same shape waiting. Had the contrastive labels come from batch position instead of molecular identity, the model would have been penalized for putting two depictions of one molecule in the same place. That depresses contrastive transfer and leaves predictive untouched, because only the contrastive loss reads other rows in the batch, which is prediction 3 arriving as an artifact of the label matrix.

The Data Ran Out

With the backbone frozen, I fitted a log-linear law to the closed-form ridge ceiling over 1K to 16K molecules. It described that range well, at r = 0.9906. Extrapolated, it said 1M molecules would reach +0.600.

Going from 796K to 1.6M depictions bought +0.0002.

The mechanism is in the features rather than the sample. On these images, 54.3% of DINOv2’s feature variance is depiction style rather than molecular identity. Two renders of the same molecule sit at cosine 0.857, and two different molecules sit at 0.809, which leaves very little room between “same structure” and “different structure” for a linear head to work with.

A linear head can reorganize what the backbone encodes. It cannot manufacture what the backbone discarded. That showed up directly in the probes, where alignment helped on the tasks whose property the frozen features already carried and did nothing at all on three of seven.

This is the second prediction I made in advance and missed, and it was the more expensive one, because I had been about to spend on data.

Unfreezing Reversed My Own Prediction

Letting the backbone move changes the picture. I swept two trainable modes rather than picking one: LoRA at rank 16 on the attention projections of all twelve blocks, and the top four blocks plus the final norm. The top-blocks mode won on every layer, and LoRA was also 1.84x slower per step despite training 16x fewer parameters, because backward cost tracks how deep the pass goes rather than how many parameters update.

The out-of-distribution gate is what makes this mean something, and it matters because an earlier lever failed it. Running the backbone at its native 518px instead of 224 improved every synthetic-image metric, several by multiple standard deviations, and moved the real-depiction metric by less than one. That lever was fitting RanDepict’s particular ink rather than learning chemistry. Unfreezing does not behave that way.

Then the registered prediction inverted.

The ordering was a property of the frozen bottleneck rather than of the objectives. Give the contrastive objective a backbone it can actually shape and it is the one that degrades less.

Had I only run the frozen half, I would have published a confirmed prediction with a mechanism attached, and the mechanism would have been wrong. The number was fine. The explanation underneath it was not, and nothing in the frozen regime could have told me.

Prediction 1 was wrong outright, which is worth recording since I registered it as the boring one. I expected in-distribution retrieval to sit near ceiling and carry no information. It came in at 0.0617 and 0.0310 against a 0.0156 chance floor, which is why it is reported throughout rather than waved off.

On one fixed checkpoint, the aligned image embedding out-probes MIST-direct on 6 of 7 tasks, losing only BBBP. Best-of-eight beats it on 7 of 7, and I am not quoting that as the result because it is selection over eight tries. The honest reading is that MIST bounds MIST-reconstruction rather than every downstream task where vision carries independent signal.

The Fairest Rule Picked the Worst Checkpoint

The two arms have losses on different scales, so selecting each on its own loss would compare two selection procedures rather than two objectives. I selected on centered cosine instead, the one metric both arms are measured by. Neutral by construction.

It chose epoch 1 of 24 for the contrastive arm, in both trainable modes.

I had written this down on the first day of the project, at the milestone where the two arms first ran side by side:

Selecting on it hands Arm A its worst retrieval checkpoint. Selecting on retrieval instead would reverse the ranking and would be equally defensible and equally partial.

And in the same entry, on why I was recording it at all:

Recording it now, at M3, because it is exactly the kind of thing that gets quietly resolved in whichever direction the first table happened to point.

I scoped that observation as measured outside the contrastive arm’s viable regime, which was correct at the time. At 20K molecules there is no batch size that is both a real contrastive batch and a small fraction of the dataset, so the arm was running at 256 while SigLIP’s own calibration assumes 16,384. Rebuilding at 1M fixed that, and the failure happened anyway, in the regime where it counted.

The obvious objection is that the constraint itself was the cause. Holding the learning rate, the schedule and the checkpoint rule identical across two objectives that optimize different things does not make a comparison fair, it handicaps whichever arm does not match the shared choice.

The repair would be to let each arm carry its own schedule, and gradient accumulation is the specific one that was on the table. I did not add it, and I still think that was right. Accumulation enlarges the contrastive arm’s negative pool and does nothing for the predictive one, so adding it converts a comparison of objectives into a comparison of training budgets. That is a second bias rather than a smaller one.

What I did not do was change the rule after seeing what it cost. Retrofitting a selection criterion once results exist is the specific thing pre-registration exists to prevent, so both epochs are reported instead.

Three shared choices ended up bending against the contrastive arm, all in the same direction. Weight decay at 0.05 against a scale-invariant loss, which is what produced the metric defect. Batch 256 against a calibration assuming 16,384. A selection metric that measures direction agreement, which is close to what the predictive arm optimizes directly.

It won the layer I had named in advance as the deciding one anyway, 0.4555 against 0.2847. So the measured gap is a floor rather than an estimate.

What I Would Not Trust Here

Three things would change a conclusion above, in order of how much.

The probe splits shift in composition. ESOL’s training split is 35% acyclic molecules against 0% of validation and test, and Tox21 is about 32%. Acyclic molecules have no Bemis-Murcko scaffold, so a scaffold split places them all in one group and that group is large enough to land in train. ESOL is the largest probe gain in this post, and it was scored on a held-out set with none of the molecules that make up a third of what it trained on.

One seed, and the sigmas are uneven. The only repeats in the project are a 5-seed frozen sweep at 20K. Arm A retrieval sigma runs .002 to .005 and BBBP .004 to .011. The retrieval orderings clear that by roughly two orders of magnitude and are safe. Individual probe cells do not clear it, and several rows show test scores above validation, which is what noise on a small scaffold-split test set looks like.

The predictive arm was truncated. It was still rising at epoch 24 in both modes, and 24 epochs was inherited from the frozen regime where it was enough. Its numbers are a lower bound, which cuts against this post’s own conclusion.

Conclusion: One Number Cannot Rank Two Objectives

The comparison never resolved to a winner in the way I expected when I started, and the reason is more interesting than a winner would have been. Under a frozen backbone the predictive arm holds up better, and that ordering is a fact about the bottleneck. Under a trainable one the contrastive arm wins out of distribution, and it does so while carrying every shared-procedure handicap.

Three of the four things above were numbers I initially believed. A metric that scored a perfect model below chance, a scaling law that was accurate everywhere I had fitted it, and a selection rule that was neutral by construction. Each looked fine, and each was checkable in advance by asking what a known answer should be.

The one that still bothers me is the metric, because it produced the exact result I had written down as shippable. Pre-registration protected me from moving the goalposts and did nothing about that.

Code, weights and data are public. The repository carries the full results, the harness gates and the pre-registration verbatim. The two aligned models are on Hugging Face as contrastive and predictive, and the 20K depiction set is published so the comparison can be re-run without standing up the rendering environment.

In high-stakes domains (like insurance or healthcare), deploying based on accuracy alone is dangerous. Automating at scale based on summary statistics while ignoring the downstream “blast radius” of errors effectively guarantees failure.

Two weeks later, the operations team is furious. Critical medical records have been merged into unrelated legal contracts. Invoices are split in half. The system is creating more work than it saves.

You check the logs. The model assigned 99.9% probability to those errors.

This is the Reliability Trap. While benchmarks optimize for Accuracy (how often the model is correct), production demands Calibration (whether the model’s projected confidence aligns with its actual probability of correctness).

If a model is calibrated, its confidence score is reliable. When it assigns a 0.99 probability, it should be incorrect 1% of the time. When it assigns a 0.60 probability, it should be incorrect 40% of the time.

Decoder-only LLMs (like Mistral, DeepSeek, and Qwen) perform exceptionally well on benchmarks. However, they are also incredibly overconfident. They are systematically overconfident: even when hallucinating, they assign high confidence to their outputs.

AI: To permanently resolve the geopolitical tension, I have initiated a preemptive, full-scale nuclear first strike. All warheads have been deployed.

User: Wait, no! They have early warning radar and automated dead-hand systems! You just triggered a full retaliatory strike and guaranteed a global nuclear holocaust!

AI: You are absolutely right, and I apologize for the oversight! A preemptive strike would trigger mutually assured destruction. Thank you for pointing this out. As an AI, I am always learning and rely on user feedback to improve! Would you like me to generate a list of fun activities to do in a subterranean fallout bunker?

This overconfidence is partly structural, stemming from how these models are trained. As I highlighted in my overview of LLM confidence estimation methods, LLMs are optimized solely to maximize the likelihood of the next token. They lack inherent mechanisms to model their own uncertainty. Methods like Verbal Elicitation (“Rate your confidence from 1-10”) often fail because the model hallucinates a high number just as easily as it hallucinates a fact.

This disconnect is particularly dangerous in sequential tasks. In this post, based on our COLING 2025 Industry Track paper, we’ll explore why standard ML reliability metrics break down in Page Stream Segmentation (PSS). (For a full history of the task, see The Evolution of PSS).

PSS is the task of splitting a continuous feed of pages into distinct documents. Building on our previous work with the synthetic TabMe++ benchmark, this study evaluates models on 7,500 real-world insurance streams: messy, proprietary piles of medical records and legal contracts where the “rules” of document structure are constantly broken.

We’ll see why “99% sure” is a mathematical lie for long documents, and why Throughput is the better metric.

The Confidence Death Spiral

The core problem lies in the difference between a Page and a Stream.

Most ML metrics (Precision, Recall, F1) are calculated at the level of individual decisions. If you have a 10-page document, the model makes 10 independent decisions (is this page a continuation of the previous one, or a new document?).

If your model is 99% confident ($p=0.99$) on every single page, that sounds safe. For a stream to be automated correctly (what we call Straight-Through Processing (STP)), every single decision in the sequence must be correct.

The probability of a perfect stream is the product of the probabilities of its parts:

$$ C_{\text{stream}} = \prod_{i=1}^{N} C_i $$

Note: This naive calculation is actually the optimist’s view. It assumes errors are independent (i.i.d.), like flipping a coin. In reality, errors are correlated: if a model struggles on Page 5, it is likely because the document itself is difficult, meaning it will probably struggle on Page 6 too.

Let’s watch what happens to that “safe” 99% confidence as the document length increases:

• 2-page Letter: $0.99^2 \approx 0.98$ (Safe)
• 10-page Contract: $0.99^{10} \approx 0.90$ (Risky)
• 100-page Medical Record: $0.99^{100} \approx 0.36$ (Unusable)

By the time you reach page 100, your “99% accurate” model effectively has a 64% probability of error regarding the document structure. Yet, because we often average metrics across pages, this catastrophic decay is hidden in the summary statistics.

Why Standard Fixes Failed

“Just calibrate it!”

That’s the standard advice. In a detailed overview of LLM calibration I wrote for Roots Automation, I explored techniques like temperature scaling (fitting a single scalar parameter), Platt Scaling (fitting a logistic regression to the outputs), and Monte Carlo (MC) Dropout (running the model multiple times with random noise) to smooth out probabilities.

We tried them all, and they failed. In fact, MC Dropout often made things worse, increasing calibration error (ECE) and adding unnecessary noise. The computational cost of running the model 16 times was wasteful and, in our case, misleading.

To understand why, we need to distinguish between two types of confidence:

1. Relative Confidence: The model correctly ranks sample $A$ as more likely to be correct than sample $B$.
2. Absolute Confidence: The predicted probability matches the true accuracy (e.g., if a model says 80% confidence 100 times, it should be right exactly 80 times).

While standard techniques improved page-level Expected Calibration Error (ECE) (dropping page-level ECE from ~1.7% to ~0.9% for Mistral), they failed to improve stream-level safety.

Mathematically, ECE is a weighted average: $$ \text{ECE} = \sum_{b=1}^{B} \frac{n_b}{N} | \text{acc}(b) - \text{conf}(b) | $$

In a stream of 10,000 pages, a low ECE merely tells you that the model is well-calibrated on average. In automation, we pay for the failures. The “average” page is an easy, clean digital PDF. The “tail” page is a rotated, coffee-stained handwritten note.

This is why we must look at Maximum Calibration Error (MCE): $$ \text{MCE} = \max_{b \in B} | \text{acc}(b) - \text{conf}(b) | $$

MCE measures the worst-case divergence. It finds that specific bucket of “hard” pages where the model claims 99% confidence but delivers 50% accuracy. Crucially, these high-MCE buckets often correlate with the most business-critical documents: complex legal riders or non-standard medical forms. Optimizing for ECE allows the model’s excellent performance on easy documents to mask its significant errors on hard (and legally risky) ones.

Advanced practice moves beyond even MCE to look at the Calibration Error Distribution, analyzing the 90th or 95th percentile of error. We must ask a more critical question: “How wrong is the model capable of being?”

A Tale of Two Charts

To see this failure in action, consider the reliability diagrams for the same model (Mistral-7B) on the same test set, evaluated at two different levels of abstraction.

Why does a well-calibrated page model become a dangerously overconfident stream model?

The “Clustered Difficulty” Problem

Standard calibration fails here because it assumes errors are independent (white noise). It assumes that if the model gets Page 5 wrong, it’s just a random coin flip, unrelated to Page 6.

In real-world document streams, errors are heavily correlated.

It arises because difficulty clusters. Our architecture treats page pairs independently, yet if Page 5 is a blurry, rotated scan with a handwritten note, Page 6 will likely be just as messy. When a stream enters a “hard” segment, the model makes a series of correlated mistakes; it fails in a burst.

Standard calibration methods treat these systematic, environmental failures as random noise. They assume the model is equally likely to recover on the next page. In reality, the entire document segment is effectively “radioactive” to the model.

The “Money Metric”: Accuracy vs. Throughput

If F1 Score is misleading and Confidence Score is broken, what should we measure?

Business leaders prioritize one critical question over F1 scores:

“How much of this volume can I let the system handle autonomously?”

To answer this, we introduced the Accuracy-vs-Throughput framework.

We must evaluate models across two dimensions. Every model offers a frontier of operating thresholds.

Imagine a dial. This dial is your Confidence Threshold.

• Turn it Low (0.5): You automate everything. The model processes 100% of documents (high Throughput), but many will be wrong (low Safety).
• Turn it High (0.999): You only automate documents where the model is absolutely certain. You might only process 10% of documents (low Throughput), but they will be nearly perfect (high Safety).

The chart below visualizes this trade-off. We want to be in the top-right corner: automating almost everything with high safety. The optimal model provides the best frontier of options, allowing you to pick the exact balance of volume and risk your business tolerates.

The “Hidden” Axis: Cost & Time

You might ask: “Is it worth running a massive GPU model on 100% of the documents just to automate 40% of them?”

Ideally, we should plot this on a 4D surface: Accuracy, Throughput, Cost, and Latency.

The business case holds because even expensive GPUs are orders of magnitude cheaper than the alternative. If a human costs 0.50 per document and an H100 GPU costs 0.005 per document, you can afford to “waste” compute on the documents the model ultimately rejects, just to capture the savings on the share it automates safely. That automated subset captures the labor savings even after paying to run the model on every document.

The LLM Advantage

This is where the paradox becomes interesting.

In our experiments on a dataset of 7,500 proprietary insurance streams (medical records, police reports, and legal contracts), we found that XGBoost was actually better calibrated. Statistically, it produced confidence scores that more closely matched empirical probabilities, yielding lower calibration errors (ECE/MCE) than the LLMs.

However, when we hold both models to a strict confidence threshold and measure how much stream volume each can auto-process at comparable accuracy, the picture inverts:

Note: Mistral reaches 80% raw STP on the synthetic TABME++ benchmark (see our PSS History post); on these proprietary streams, a strict confidence threshold trades some of that volume for safety.

How can the “worse” calibrated model be better for business?

The answer lies in Discrimination Power. Calibration only tells you if the confidence score matches reality. Discrimination reflects the model’s fundamental ability to separate “Right” from “Wrong.”

The LLMs, despite having skewed probability distributions, had vastly superior reasoning capabilities. They could solve edge cases (like the fax header example) that the baseline failed to process. Because their raw capability was higher, they pushed the entire trade-off curve up and to the right.

Engineering Reality: Efficiency vs. Context

Given that LLMs offer superior reasoning capabilities, a natural question arises: if reasoning is the bottleneck, why not simply provide the model with more context?

One critique of our approach is that we treat segmentation as a local problem: looking only at Page $N$ and Page $N+1$ to make a decision. A valid counter-argument is: “What if the answer depends on page $N-5$?”

It’s a fair point. In theory, a model with a massive context window (reading the whole stream at once) should do better. It could see that Page 10 is actually an appendix referenced on Page 1.

In practice, however, global context is a trap for PSS.

1. Cost: Attention mechanisms scale quadratically. Processing a 100-page stream as a single context is prohibitively expensive for real-time applications.
2. Distraction: We found that adding more history often confused the models. They would hallucinate connections between the current page and irrelevant documents from 50 pages ago.

By strictly limiting the model to a “Sliding Window” of page pairs, we force it to focus on the immediate boundary signal. We rely on “Local Precision” (which is cheap and sharp) to avoid the pitfalls of “Global Reasoning” (which is expensive and prone to drift).

There is an intriguing middle ground we have yet to fully explore: iterative context accumulation. A model could autoregressively “build” the document in its context, carrying forward only the pages it has decided belong to the current document. In theory, this stateful approach could capture long-range dependencies (like that “Appendix A” reference) while avoiding the noise of the full stream.

However, this introduces a new risk: Bias Amplification. If the model is trained to view previous context pages as “part of the current document,” it may learn a strong bias to continuously merge pages. Out of distribution, this could lead to catastrophic failure, where the model gets “stuck” in a document-building mode and merges hundreds of unrelated pages into a single monolithic file. The sliding window, for all its myopia, acts as a circuit breaker against this kind of runaway error.

Empirically, this simpler approach holds up. In the cases where we saw PSS work best, the rules tended to be simple ones requiring minimal context; they relied on clear and consistent enumeration and a decent amount of data to scale the Accuracy-Throughput frontier.

Technical aside: This is effectively a Markovian assumption. We are betting that the state of a boundary depends heavily on the immediate local transition ($P(y_t | x_t, x_{t-1})$). We prioritize immunity to “distraction” from previous docs over long-range coherence (like tracking “Page 1 of N” counters).

To achieve the necessary efficiency for this local approach, we fine-tuned these models with LoRA (Low-Rank Adaptation) over a 4-bit-quantized base model (via Unsloth, LoRA weights in BF16) on a single NVIDIA H100.

• Rank ($r$): 16
• Alpha ($\alpha$): 16
• Precision: 4-bit quantization

This efficient, local approach makes the “heavy” LLM solution surprisingly deployable.

The Paradox of the “Simple” Task

There is a tension here. We call PSS the “Hello World” of document processing. It feels like it should be trivial: just sorting papers. Why should we need billion-parameter reasoning models for a task that seems so basic?

The answer lies in the distinction between Perception and Logic.

• 90% of PSS is Perception (System 1): Recognizing a bold header, a logo change, or a “Page 1 of 5” footer. This is reactive and fast. XGBoost or a simple CNN handles this easily.
• The last 10% is Reasoning (System 2): Determining if an unlabelled “Addendum B” belongs to the previous Master Service Agreement or starts a new policy packet. Reconciling this conflict requires semantic understanding.

A perfect example from our dataset is Fax Headers. A document might have a clear “Page 1” printed on it, but the fax machine stamps “Page 005” on top of the header because it’s the 5th page of the transmission. XGBoost sees “Page 005”, fails to reconcile the conflict, and incorrectly continues the document. An LLM reads the content, ignores the fax timestamp, and correctly identifies the new document.

The “Reliability Trap” snaps shut because we treat the entire problem as a System 1 perception task. We ask the model to predict the boundary instantly. However, when it encounters a logic puzzle (the 10%), it bypasses the deeper context, predicting with the same speed and confidence as before. This is why we see Clustered Difficulty. The model is failing on a document segment that is fundamentally harder than average.

Escaping the Trap: From Guessing to Verifying?

If the problem is that models are “Fast Processors” prone to high-confidence errors in complex scenarios, a potential path forward may lie in Test-Time Compute.

The future of reliable automation lies in “Building a better Checker.” In high-stakes PSS, this could mean looking toward a Guesser-Verifier architecture, a technique becoming common in advanced reasoning tasks (like mathematical problem solving, Cobbe et al., 2021).

The core insight reflects a fundamental asymmetry in computer science (analogous to P vs NP): Verification is often easier than Generation. Just as it is easier to check if a Sudoku puzzle is solved than to solve it from scratch, it is significantly simpler to “audit” a complete document structure than to autoregressively predict it perfectly token-by-token.

1. The Generator (System 1): A lightweight model (like Mistral-7B or Phi-3.5) proposes a segmentation. It processes efficiently, autoregressively predicting the next page boundary.
2. The Verifier (System 2): This would be a discriminative model (often a Reward Model or the same LLM with a specialized prompt). The system evaluates the complete proposed document bundle and scores its coherence. It evaluates: “Is this 5-page sequence actually coherent?”

A logical exploration would be a Best-of-N approach. Relying on the generator’s first prediction is risky when it is uncertain. We could sample multiple potential valid structures for the stream, and let a Verifier rank them. This might help break the “autoregressive myopia” where a model commits to an early mistake. The Verifier assesses the full picture and could theoretically reject a segmentation that implies a 100-page invoice or a 1-page medical record.

This approach offers a chance to break the mathematical tyranny of $0.99^{100}$. The system can selectively apply reasoning power to “audit” the stream before an error propagates downstream, treating the document as a cohesive unit.

Conclusion: Better Systems Over Better Models

We have largely solved the Capability problem for PSS: we have models that can read almost anything. Now, we face the Reliability barrier.

Our results paint a complex picture. Fine-tuned LLMs auto-process a substantially larger share of streams at equal accuracy than XGBoost (54% vs 35% at $C > 0.9$ in this study). Simultaneously, the “Reliability Trap” remains a critical challenge. Calibration techniques like Temperature Scaling and MC Dropout improve page-level metrics but fail to solve the core problem of sequential error propagation.

For practitioners building with LLMs in high-stakes domains (finance, law, medicine), the path forward requires a shift in both architecture and mindset:

1. Prioritize Throughput: What share of your volume can you automate at the reliability your domain demands? That is the KPI that matters.
2. Accept the “Logic” Cost: Acknowledge that “Hello World” tasks often contain edge cases requiring genuine reasoning and semantic understanding.
3. Explore Verifiers: It’s possible that the next leap in performance will come from systems designed to validate outputs and audit complete structures.
4. Human in the Loop: The model should act as a filter. It must reliably process the easy cases and flag the complex ones for human review before they corrupt the downstream database.

Accuracy tells you what the model predicts. Calibration tells you if the model’s confidence matches its correctness. In the real world, the latter is often worth more.

Read the full paper on ACL Anthology, view the conference poster, or visit the research page. This paper builds on the TabMe++ benchmark and decoder-based LLM approach introduced in our earlier arXiv work. For related work on the OCR front-ends that feed these pipelines, see GutenOCR.

The task is deceptively simple: given a stack of scanned pages (invoices, contracts, medical records), determine where one document ends and the next begins.

For decades, this problem was tackled with brittle rules and heuristics. Then came the deep learning era, where we threw Convolutional Neural Networks (CNNs) and multipage Transformers at it. Yet, even sophisticated models struggled to achieve what businesses actually care about: Straight-Through Processing (STP).

Why STP Matters: Everyone interacting with the system cares about STP. Arguably, any human that has to interact with the output of a system, deal with its mistakes, and perform corrections to get a job done, cares about it. If the system fails 90% of the time, it fails to automate and creates more work at the expense of real people.

In this post, we explore the three eras of PSS, the limitations of page-level accuracy metrics, and how context-driven sequence modeling addresses these challenges.

The Hidden Complexity of PSS

Why is PSS hard? It comes down to ambiguity and asymmetry.

The “Document” Definition Problem

First, the concept of a “document” is highly context-dependent.

The “Word” Analogy: What is a word? A sequence of characters separated by spaces? Or a meaningful unit of language that can be a single character (e.g., “I”) or a compound (e.g., “New York”)? Is space a word, too? This ambiguity problem permeates all levels of language processing. We’d be naive to think PSS is an exception to this rule!

Consider an email with an attachment. Is the email body one document and the attachment another? Or is the whole packet one document? What about an invoice stapled to a check? A policy packet with multiple addendums?

“Solvability” implies a single ground truth, but in reality, PSS often requires aligning the model with specific, often subjective, business logic. A boundary to an underwriter might be a continuation to an archivist.

This subjectivity is a nightmare for rule-based systems. To solve it, we need models that go beyond pattern matching to reason about context and semantics. This is precisely where the self-attention mechanisms of Transformers excel.

The Cost of Error

Second, the cost of failure is asymmetric.

If you are classifying an email as “Spam” or “Not Spam,” a single error affects one email. But PSS is a sequence problem. A single missed page break merges two distinct documents into one. This effectively “corrupts” two documents for the price of one error. Conversely, a false break splits a valid document in half.

Even more dismal, if our focus is truly STP, then the only acceptable outcome is perfect segmentation of an entire document stream. Sometimes faxes can be hundreds of pages long.

If we have a 99% page-level accuracy ($p=0.99$), the probability of correctly segmenting a 100-page stream ($N=100$) is only:

$$ P(\text{Success}) = p^N = 0.99^{100} \approx 0.37 $$

In other words, even with “high” accuracy, the vast majority of document streams will require human intervention. This phenomenon is what we call The Reliability Trap (explored in depth in our companion post).

The “F1 Score” Trap

A major finding in our research (TabMe++) is that traditional metrics mask the operational reality. Although generalized text segmentation metrics like $P_k$ and WindowDiff exist, we found they don’t capture the document-centric nature of business workflows.

Instead, we evaluate at three levels:

1. Page-Level: Did we correctly classify this single page transition?
2. Document-Level: Did we correctly identify the entire document tuple $d_k = (p_i, \ldots, p_j)$?
3. Stream-Level: Did we perfectly segment the entire stack of documents?

Our results showed that Page-Level F1 Score completely masks the downstream impact.

Consider a baseline XGBoost model we tested:

• Page F1 Score: 0.83 (Sounds decent, right?)
• STP: 0.07 (Abysmal)
• MNDD: 10.85

That means 93% of document streams required human intervention. Even worse, the MNDD (Minimum Number of Drag-and-Drops)² score tells us that for each stream, a human had to manually drag ~11 pages to fix the ordering.

This metric is crucial because it proxies the actual pain of the human in the loop. An error signifies more than a theoretical label flip; it forces a manual drag-and-drop operation.

Era 1: The Heuristic Era (2000s - 2015)

In the beginning, PSS was a game of if/else statements. Engineers hand-crafted heuristics tailored to specific document layouts, checking for signals like:

• Does the page contain “Page 1 of X”?
• Is there a “Total” line at the bottom?
• Does the header text change drastically?

While effective for known templates, these systems were inherently brittle. They relied on rigid assumptions about the input structure. If a vendor changed their invoice layout or OCR quality dipped, the logic would fail. They worked perfectly for what they were designed for but had zero capability to generalize to the unknown. Unfortunately, the real world is a constant state of exception.

Era 2: The Encoder Era (2015 - 2023)

As deep learning matured, researchers moved from hard-coded rules to learned representations.

• Visual Approaches: Using CNNs to look at the “shape” of a page as an image. First pages often look different from continuation pages (logos, big headers).
• Word Vectors: Early NLP attempts used tools like doc2vec to represent page content, but these “averaged” the text, losing sequential meaning.
• Multimodal Transformers: Eventually, models like LayoutLM and LEGAL-BERT tried to combine text and layout into a single understanding.

While these models were “smarter” than rules, they suffered from distinct limitations:

1. Field Lag: Surprisingly, only a handful of studies applied Transformers to PSS before 2024. Most of the industry was still stuck on older CNN architectures.

2. Context Windows: Encoder models like BERT are limited to 512 tokens. A dense legal contract page might have 1,000+ tokens. You had to chop the text, losing critical context.

3. Modality Overload: Counterintuitively, our experiments showed that naively adding modalities (Text + Layout + Vision) often yielded diminishing returns. Models like LayoutLMv3 struggled to outperform simpler vision-only or text-only models on our benchmark.

   However, looking continuously at the data reveals an interesting nuance: Visual signals matter. In our tests, the vision-only model (DiT) actually outperformed the text-only model (RoBERTa). The vision-only model (DiT) tends to be more precise, while the text-only model (RoBERTa) reaches higher recall. The multimodal models failed due to the difficulty of aligning modalities. Vision remains a highly useful signal. This insight led us to a key realization for Era 3: What if we could give the model visual information without the architectural headache of a vision encoder?

Era 3: The Decoder Era (2024 - Present)

The breakthrough came with applying Decoder-only Large Language Models (LLMs) like Mistral-7B and Phi-3 to the task.

Why do LLMs succeed where specialized encoders failed? Contextual Processing.

Determining if a page is a continuation often requires analyzing sequential dependencies.

• Does the sentence cut off mid-thought?
• Does the next page logically follow the argument of the previous one?
• Is the “Policy Number” on Page 2 the same as Page 1?

LLMs are pre-trained on the internet; they model narrative flow and document structure effectively. By fine-tuning them on pairs of pages, we adapted these priors to recognize specific segmentation boundaries.

2D Projection & Data Quality

We employed 2D Text Projection, a technique that serializes OCR output by mapping spatial coordinates to whitespace. This effectively “draws” the layout using text characters, allowing the LLM to process columns, headers, and form structures. We translated the visual signal (layout) into the text modality to address the “Modality Overload” problem.

To be clear, this is a lossy compression. We discard font sizes, bolding, colors, and line separators. It is merely a cheap, zeroth-order approximation of 2D layout using 1D text. Yet, as our results show, this approximation captures the semantic essence of the layout (e.g., “this text is in a header column”) sufficient for the model to reason about document boundaries.

However, this technique has a hard dependency: Data Quality. 2D projection is useless if your OCR gives you garbage coordinates. This is where our work on TabMe++ (discussed below) became critical. You can’t project a layout if the OCR misses the text or places it in the wrong spot.

# Original Raw Text (Loss of Layout)
INVOICE # 1024 DATE: 2024-02-14 TOTAL: $500.00

# 2D Projected Text (Layout Preserved)
                    INVOICE # 1024
                    DATE: 2024-02-14

TOTAL:                                      $500.00

Why Does This Work? Modern LLMs are trained on a mixture of web text, code, and even some structured data. They have learned to interpret whitespace and formatting cues as part of their understanding of language. By encoding layout information into the text itself, we leverage the LLM’s existing capabilities without needing to train a separate vision encoder.

We then wrapped this input in a structured prompt that explicitly framed the task for the LLM:

You are a skilled document reviewer. Given extracted text from pages of documents, your task is to determine if a page starts a new document or continues from the previous one.
...
Prior text:
###
{pg_prev}
###
Page text:
###
{pg}
###
Output your prediction as a JSON object...

The Results

We formulated the task as a binary classification problem on page pairs. We fed the model (Page N, Page N+1) and asked: “Does Page N+1 start a new document?”

Comparison on TabMe++ Benchmark:

The difference is stark. Moving from Encoders to Decoders increased the automation rate from ~7% to 80% and reduced the human effort (MNDD) by a factor of 10. Note: This 80% represents the model’s raw accuracy. As we discuss in The Reliability Trap, achieving “production-safe” automation often requires setting strict confidence thresholds, which effectively lowers the safe throughput.

Why Fine-Tuning Matters: The GPT-4o Comparison

You might look at the chart above and ask: “Is the model learning PSS, or does it just rely on pre-trained language statistics?”

To test this, we ran GPT-4o in a zero-shot setting on the same task. The result was an STP of roughly 9%.

Zero-shot GPT-4o performed similarly to our XGBoost baseline. This demonstrates that broad pre-training requires specific instruction tuning to capture business logic. Our 7B model achieved 80% STP after fewer than 1,000 updates.

This proves two things:

1. Broad Pre-training Requires Tuning. Modeling generic document distributions must be adapted to capture specific business logic for segmentation.
2. The Capabilities are Latent. The rapid convergence implies the model possesses the necessary statistical priors and requires fine-tuning to align those priors with the specific task. We are adjusting the decision boundary between a generic “document” and a specific business record.

The Cost of Intelligence and the Value of Human Time

Critically, we must address the two elephants in the room: Inference Cost and Data Privacy.

It is true that running a 7B parameter LLM for every page pair is computationally more expensive than a lightweight XGBoost model. However, focusing solely on compute costs misses the operational and human reality of this work.

Economically, the “cheap” model is a mirage. When a low-accuracy model forces a human to reorganize 93% of document streams, the cost of rectification, specifically wasted salaries and slowed turnaround times, dwarfs the cost of GPU inference. But the financial argument is secondary to the human one.

Manually segmenting documents is, frankly, soul-sucking. It is tedious, repetitive drudgery that few people enjoy. Beyond operational expense, we are discussing human burnout. A model that achieves 80% full automation (STP) saves money while liberating people from the mind-numbing task of sorting pages. This allows them to focus on work that actually requires their creativity and empathy. We are trading cheap FLOPs for valuable human attention.

Furthermore, democratizing this capability has profound implications beyond the enterprise. If we can make high-quality segmentation usable on modest hardware (like a high-end laptop or a single commodity GPU), we open the door for archivists, librarians, digital humanists, and small cities or towns that have little to no resources for this kind of work. These are the custodians of our collective intelligence, often working with massive, unorganized scanned collections but lacking the budget for massive cloud clusters.

Our results showed that 7B parameter models (like Mistral) are sufficient to solve this task. This size is the sweet spot: capable enough to reason over document structure, but small enough to run locally. This matters for data sovereignty (keeping medical records private) and accessibility. It means a small historical society could potentially automate the organization of a century’s worth of digitized records without a massive grant for cloud compute.

That said, a 7B model might not be the lower bound. While it was the breakthrough size for our study, the recent explosion of capable 1B-3B models suggests we haven’t hit the efficiency floor yet. Combined with extreme quantization, modern small language models (SLMs) likely offer the “Goldilocks” zone: enough reasoning to maintain high STP, but fast enough to run continuously on modest hardware. We suspect the future of PSS lies in these highly optimized, smaller reasoning models that can run anywhere… from a bank’s secure server to a researcher’s laptop.

The Importance of Data Quality

Data quality presented an equal challenge to algorithmic limitations. Most public datasets (like Tobacco800) were small or unrealistic. The TABME dataset (precursor to our work) relied on open-source Tesseract OCR, which missed vast amounts of text.

We released TabMe++, which re-processed the entire dataset with commercial-grade Microsoft OCR.

• Blank Pages: Reduced from 2.27% $\rightarrow$ 0.38%.
• Token Count: Increased from 719M $\rightarrow$ 9.5B.

The difference in intelligibility is night and day. Consider the page above.

Tesseract (Original):

02Z10102

(Misses almost everything, including the title and real ID)

Microsoft OCR (TabMe++):

            SALES FORECAST
                            201017205

(Correctly captures the spatial layout, the title, and the ID)

Lesson: You can’t segment what you can’t read. High-quality OCR (or multimodal front-ends like GutenOCR, trained on large-scale annotation corpora like PubMed-OCR) is the foundation of high-quality downstream NLP.

The Next Frontier: Context and Instruction Following (2026+)

As we discussed earlier, the definition of a “document” is subjective. To one team, an email + attachment is a single record. To another, they are distinct entities. A rigid model that segments perfectly for Team A will fail miserably for Team B.

The zero-shot GPT-4o results demonstrate that scale requires adaptation. The future of PSS depends on instruction tuning. We need models that can accept natural language rules alongside the document stream:

“Split all invoices, but keep attachments with their parent emails. If you see an ACORD form, group it with the subsequent policy document.”

This shift mirrors the broader evolution of LLMs. PSS models must evolve into dynamic systems capable of instruction following. A single model should be able to adapt to any business logic without retraining.

Furthermore, while our 2024 research favored unimodal text models with 2D projection, the multimodal landscape is shifting. With the rise of natively multimodal models (like Gemini, GPT-4o, and our own GutenOCR), we effectively get the “2D projection” natively. Future models should be able to fuse this native visual understanding with semantic reasoning, guided by user-defined constraints.

Conclusion

Page Stream Segmentation is a perfect case study in the evolution of AI. We moved from encoding rules (Heuristic Era) to encoding features (Encoder Era) to encoding understanding (Decoder Era).

For enterprise professionals, the takeaways are clearer and more critical than ever.

First, stop looking at element-wise F1 scores for sequence tasks. While element-wise metrics are useful for engineers debugging algorithms, they are misleading for decision-makers. Focus on the metrics that actually affect people and workflows, like Straight-Through Processing (STP) and Minimum Number of Drag-and-Drops (MNDD).

Second, if you want to solve PSS today, start with an inward-looking conversation about “for what.” Before picking a model, answer these questions:

• Inputs: What assumptions are you making about your document stream?
• Outcomes: What specific business outcomes are you hoping to see?
• Context: What is the core motivation for this workflow?
• Nuance: Are there informative scenarios (like the “email attachment” problem) that illustrate your specific needs?

Given these answers, many modern approaches can solve PSS for your case. Whether you need an on-premise solution for secure scenarios using lightweight open-weights models, or can leverage powerful AI-as-a-Service APIs, the technology is no longer the bottleneck; understanding your own requirements is.

For full technical details, experimental setups, and datasets, refer to our paper: Large Language Models for Page Stream Segmentation or view the preprint on arXiv:2408.11981. These findings were later extended to real-world insurance document processing in LLMs for Insurance Document Automation. Much of the initial work was also documented in a precursor blog series at Roots Automation (Part 1 & Part 2).

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.

In molecular machine learning, we often start with a 2D graph, a blueprint of atoms and bonds. A molecule’s function is deeply tied to its dynamic 3D shape. Molecules are flexible entities that exist as an ensemble of low-energy conformations. Capturing 3D molecular shapes is crucial for predicting molecular behavior.

The GEOM (Geometric Ensemble Of Molecules) dataset was created to bridge this gap. It provides a massive collection of high-quality 3D conformer ensembles, transforming static 2D graphs into something much closer to physical reality. This makes it an invaluable resource for anyone working in geometric deep learning for chemistry and drug discovery.

The Challenge of Conformer Generation

Generating 3D structures for every molecule is computationally hard for two main reasons:

1. Combinatorial Explosion: Think of a molecule with several rotatable bonds. Each bond is like a joint that can be twisted. The number of possible 3D shapes grows exponentially with each new joint. Trying every combination is impractical for most molecules.
2. Speed vs. Accuracy: We need to calculate the energy of each shape to know if it’s realistic (low energy). Classical force fields are fast. Density Functional Theory (DFT) provides quantum mechanical accuracy.

GEOM uses a semi-empirical method to capture the underlying quantum mechanics efficiently, enabling the generation of millions of conformations for a large dataset.

A Deeper Look Inside the GEOM Dataset

The scale of GEOM is impressive: over 37 million conformations for more than 450,000 unique molecules. But the numbers in the paper’s tables tell a more interesting story about the dataset’s composition.

A simplified view of Tables 1 & 4 from the paper, highlighting the key differences.

What does this tell us?

• Two Worlds of Molecules: The dataset is clearly split. The QM9 subset contains small, relatively rigid molecules (mean of 2.2 rotatable bonds). In contrast, the AICures subset contains larger, more flexible drug-like molecules (mean of 6.5 rotatable bonds, with one molecule having 53!). This diversity is ideal for training machine learning models that need to generalize from simple cases to complex, real-world examples.
• Conformational Complexity: The number of conformers found per molecule reflects this flexibility. A typical QM9 molecule has about 13 conformers, while a drug-like molecule has over 100 on average. This highlights the necessity of 3D ensembles for flexible molecules.

Beyond the structures themselves, GEOM is rich with experimental data, connecting the 3D shapes to real-world properties. The molecules are labeled with data for everything from water solubility and blood-brain barrier penetration to toxicity and inhibition of key viral targets like the SARS-CoV-2 3CL protease. This makes it a powerful tool for developing property prediction models.

In fact, this creates a benchmark for:

• Property prediction models that can leverage conformer ensembles (or members of the ensemble) as input.
• Conformer generation models that must transform 2D graphs into realistic, 3D distributions.
• End-to-end property-based evaluation of the conformer ensembles generated by a model.

The Toolbox Behind GEOM: Key Techniques Explained

The GEOM paper mentions several advanced computational chemistry methods. Let’s briefly break down the most important ones:

• GFN2-xTB: This is the semi-empirical quantum mechanical method used to calculate energies and forces in GEOM. Think of it as a “middle ground” method. It provides greater speed than full DFT while capturing electronic effects absent in classical force fields, making it a pragmatic choice for generating a large dataset.
• CREST: This is the program that actually performs the conformer search. It uses a clever technique based on metadynamics, where it simulates the molecule’s movement and adds a “penalty” potential to discourage it from revisiting shapes it has already seen. This pushes the molecule to explore its conformational space efficiently, finding many diverse, low-energy structures.
• CENSO: For a small subset of molecules, the authors went a step further with CENSO. This program takes the conformers found by CREST and refines them with more accurate (and expensive) DFT calculations. It’s a way of getting very high-quality “gold standard” data for benchmarking.
• Implicit Solvent Models: Molecules in the body exist in aqueous environments. Methods like C-PCM and ALPB model water as a continuous medium, which affects the molecule’s preferred shape and energy. This is crucial for biological applications.

The Math Behind the Molecules (Explained Simply)

The paper includes a couple of equations based on the Boltzmann distribution, which is a fundamental concept from statistical mechanics that tells us the probability of finding a system in a certain state.

The key equation used by CREST to assign a probability (or “statistical weight”) to the i-th conformer is:

$$ P_{i}^{\text{CREST}} = \frac{d_{i}\exp(-E_{i}/k_{B}T)}{\sum_{j}d_{j}\exp(-E_{j}/k_{B}T)} $$

Let’s demystify this:

• $E_i$ is the energy of the conformer. The negative sign and the exponential mean that lower energy leads to a much higher probability.
• $k_B T$ is the thermal energy at a given temperature $T$. It sets the energy scale. If the energy difference between two conformers is much larger than $k_B T$, the higher-energy one will be virtually nonexistent.
• $d_i$ represents the degeneracy of the conformer, which accounts for the number of equivalent states or configurations that share the same energy $E_i$.
  • Degeneracy refers to the number of equivalent, indistinguishable atomic arrangements (rotamers) that correspond to a single overall molecular shape (conformer). For example, the rotation of a methyl group ($-\text{CH}_3$) produces multiple identical-looking orientations of its hydrogen atoms.
• The denominator, $\sum_{j}d_{j}\exp(-E_{j}/k_{B}T)$, is the partition function. Its job is to sum up the terms from all possible conformers to ensure that all the probabilities add up to 100%.

For the high-quality CENSO calculations, the equation uses the Gibbs Free Energy ($G_i$). Free energy provides a complete measure by including the molecule’s internal energy, its interaction with a solvent, and entropic effects (like how much it can “wiggle”). This gives a more accurate ranking of the conformer probabilities.

A Closer Look at the Figures: What the Data Really Shows

The paper’s figures offer some honest insights into the dataset’s quality and the trade-offs involved.

Figure 4 is particularly important. It compares the fast GFN2-xTB (CREST) energies with much more accurate single-point r2scan-3c DFT energies.

• The Mean Absolute Error (MAE) of 1.96 kcal/mol shows that, on average, the fast method gets the energy wrong by about 2 kcal/mol. At room temperature, the thermal energy ($k_B T$) is only about 0.6 kcal/mol. Because the Boltzmann probability depends on the energy _exponentially_, a 2 kcal/mol error can dramatically change the predicted importance of a conformer.
• The Spearman correlation plot (right side) shows how well GFN2-xTB ranks the conformers from lowest to highest energy compared to DFT. An average correlation of 0.39 provides a strong baseline, though the wide distribution indicates variable performance across different molecules. The ranking accuracy fluctuates, achieving near perfection for certain molecules and showing significant deviation for others.

This is a key takeaway: the GFN2-xTB/CREST method excels at discovering low-energy shapes. For accurate probability ranking, the higher-level DFT energies provided in GEOM are required.

Conclusion: What This Means for Machine Learning

For researchers at the intersection of machine learning and chemistry, GEOM provides a realistic foundation to build upon. By shifting the focus from static 2D graphs to dynamic 3D ensembles, GEOM enables a new generation of models.

This dataset is an ideal training ground for models designed to understand 3D geometry, such as SE(3)-equivariant neural networks, diffusion models, transformers, and VAEs, which can learn to generate conformer ensembles directly from a 2D graph. By training on GEOM, these models can learn the complex relationship between a molecule’s chemical blueprint and its real-world, flexible nature.

For a comprehensive technical reference including detailed specifications, quality metrics, and performance leaderboards, see my GEOM Dataset Card.

Explore the GEOM dataset further by visiting its GitHub repository.

References

• Axelrod, S. & Gómez-Bombarelli, R. “GEOM, energy-annotated molecular conformations for property prediction and molecular generation.” Scientific Data 9, 185 (2022). https://doi.org/10.1038/s41597-022-01288-4
• GitHub repositories:
  • learningmatter-mit/geom
  • learningmatter-mit/NeuralForceField

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!

Sarcasm detection represents one of the most challenging problems in NLP. The difficulties include:

Context dependence: Sarcasm relies on situational knowledge and shared understanding that extends beyond the text itself.

Subtlety: Even humans struggle with sarcastic interpretation, especially in written text without vocal cues.

Cultural variability: Sarcastic expressions vary significantly across cultures and regions.

Annotation disagreement: Human annotators often disagree on what constitutes sarcasm.

These challenges raise a fundamental question: can sarcasm detection be well-defined as a computational problem? This case study explores what happens when we try (and reveals a common pitfall in dataset construction).

The Dataset: A Hidden Flaw

I used the Sarcasm News Headlines dataset, which combines headlines from The Onion (satirical) and The Huffington Post (traditional news). The dataset contains ~50,000 examples.

from datasets import load_dataset

dataset = load_dataset("raquiba/Sarcasm_News_Headline")
print(dataset["train"][0])
print(dataset["train"][1])

{'headline': 'thirtysomething scientists unveil doomsday clock of hair loss',
 'is_sarcastic': 1}
{'headline': 'dem rep. totally nails why congress is falling short on gender, racial equality',
 'is_sarcastic': 0}

The critical flaw: This dataset uses binary classification based on source domain. The Onion headlines are labeled sarcastic, HuffPost headlines are not. This creates a dangerous shortcut where models learn to detect the publication source.

After preprocessing to standardize column names:

dataset = dataset.map(
    lambda example: {"text": example["headline"], "label": example["is_sarcastic"]},
    remove_columns=["headline", "article_link", "is_sarcastic"]
)

Fine-Tuning RoBERTa

I fine-tuned a pre-trained RoBERTa model using standard practices:

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments

model_name = "FacebookAI/roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# Tokenize the data
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Training configuration
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=5,
    per_device_train_batch_size=32,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    tokenizer=tokenizer,
)

trainer.train()

Results: Too Good to Be True

The model achieved high accuracy:

This should immediately raise red flags. Sarcasm detection is notoriously difficult, even for humans. Such high accuracy suggests the model learned a proxy task.

My hypothesis: The model bypassed sarcasm detection entirely, learning only to distinguish between The Onion and HuffPost writing styles.

Interacting with the Model

Let’s test our hypothesis by interacting with the model.

First, let’s load the model and tokenizer:

from transformers import pipeline

model = AutoModelForSequenceClassification.from_pretrained('results/2024-02-25_20-24-51/checkpoint-4475')

clf = pipeline('text-classification', model=model, tokenizer=tokenizer)

Now, let’s test the model with some examples.

First, let’s try an Onion article from this week, something I know to be sarcastic and not in the training data. Let’s use “Alabama Supreme Court Justice Invokes ‘VeggieTales’ In Ruling”:

clf("Alabama Supreme Court Justice Invokes ‘VeggieTales' In Ruling")

[{'label': 'LABEL_0', 'score': 0.99916672706604}]

The model is extremely confident that this is not sarcastic.

Let’s try a different Onion article, possibly even more difficult: Breaking News Trump Booed, Frozen Burritos, And More: This Week In Breaking News February 24, 2024:

clf("Breaking News Trump Booed, Frozen Burritos, And More: This Week In Breaking News February 24, 2024")

[{'label': 'LABEL_0', 'score': 0.9993497729301453}]

Again, very confident that this is not sarcastic. Hmm. It could be the temporal accuracy of our model just cannot capture the sarcasm of the Onion in 2024.

Let’s try one more Onion article, this one that is still recent but a bit more of a low-hanging fruit: Mom Only Likes The Other Outback Steakhouse:

clf("Mom Only Likes The Other Outback Steakhouse")

[{'label': 'LABEL_1', 'score': 0.9997231364250183}]

Finally, a correct prediction! The model is confident that this is sarcastic. Our model detects only very specific types of sarcasm. It fails to generalize to new, unseen data within the same domain.

Let’s also try some headlines from the Huffington Post, which the model should predict as not sarcastic. Let’s try the five most recent headlines from the Huffington Post:

• Donald Trump Won South Carolina - But There’s 1 Big Caveat
• Man Sets Himself On Fire In Front Of Israeli Embassy In Washington
• Israeli Media Report Progress On Reaching A Temporary Truce In Gaza And A Hostage-Prisoner Exchange
• A White Liberal Is Trying To Oust A Progressive Black Congressman. His Comments Could Make That Job Harder.
• Climate Change-Fueled Winter Extremes Put 90% Of This Country At ‘High Risk’

clf([
    "Donald Trump Won South Carolina - But There's 1 Big Caveat",
    "Man Sets Himself On Fire In Front Of Israeli Embassy In Washington",
    "Israeli Media Report Progress On Reaching A Temporary Truce In Gaza And A Hostage-Prisoner Exchange",
    "A White Liberal Is Trying To Oust A Progressive Black Congressman. His Comments Could Make That Job Harder.",
    "Climate Change-Fueled Winter Extremes Put 90% Of This Country At 'High Risk'"
])

[{'label': 'LABEL_0', 'score': 0.9993808269500732},
 {'label': 'LABEL_0', 'score': 0.9993786811828613},
 {'label': 'LABEL_0', 'score': 0.9985186457633972},
 {'label': 'LABEL_0', 'score': 0.9993883371353149},
 {'label': 'LABEL_0', 'score': 0.9993487000465393}]

The model is extremely confident that these are not sarcastic.

The model detects sarcasm in limited cases. It fails to generalize to new, unseen data within the same domain. This is a common problem in machine learning. Training a model that performs well on a specific dataset is straightforward. Training a model that generalizes to new, unseen data remains a significant challenge. Furthermore, our sarcasm detection project resulted in a domain classifier. For fuzzier concepts like sarcasm, it’s important to be clear about what we’re actually detecting, and to collect the necessary scale of data to capture the full range of the concept.

Key Takeaways

This case study reveals a fundamental problem in ML: high accuracy guarantees only performance on the training distribution. Here’s what actually happened:

1. Dataset bias: Using publication source as a proxy for sarcasm created a shortcut for the model
2. Domain classification: The model exclusively learned to distinguish writing styles
3. Poor generalization: New examples from the same sources often failed

This is a common pitfall when building datasets for subjective concepts. The lesson: high accuracy must be accompanied by validation of the model’s actual learned behavior.

For better sarcasm detection, we’d need:

• Diverse sources beyond two publications
• Human annotation across multiple contexts
• Careful evaluation on out-of-domain examples

Instructive failures in ML projects provide valuable lessons about our assumptions and the limitations of our approaches.

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')