Cost-aware Finetuning for Retrosynthesis

Final project for DS-GA-1011, Natural Language Processing with Representation Learning, at NYU, with Zhaodong Liu and Haoqi (Kevin) Zhang. The full report is Cost-aware Finetuning on Pretrained LM to Efficiently Perform Retrosynthesis of Chemical Reactions.

Retrosynthesis as a sequence-to-sequence problem

Retrosynthesis runs a reaction backwards: given a product molecule, predict the reactants that would produce it. What makes this an NLP problem is SMILES (Simplified Molecular Input Line Entry System), a notation that writes a molecule as one line of text, encoding atoms, bonds, branches and ring closures as characters. Under that representation retrosynthesis is string-to-string translation, and an encoder-decoder language model applies directly — the framing established by the Molecular Transformer (Schwaller et al., ACS Central Science, 2019). We built on Chemformer (Irwin et al., 2022), which supplies a BART-style model pretrained on SMILES reconstruction; that checkpoint is what we finetuned.

Two costs

Chemical reaction data is scarce because producing it means running experiments, which are slow and expensive. Nearly every public dataset is a variant of USPTO — reactions mined from US patents, 1976 to September 2016 (Lowe, 2017) — and USPTO-50K is a preprocessed 50,000-reaction subset of it. Separately, the large chemistry models that perform best on this task, such as LlaSMol (Yu et al., 2024), are expensive to train and to serve. We went after both costs: cheaper data, and cheaper finetuning.

Building USPTO-50K_γ

Instead of running more experiments, we generated labels. Starting from the real USPTO-50K data, we used LlaSMol to predict outputs for reactions that would be slow or costly to obtain in a lab, and blended the generated examples back in. We call the result USPTO-50K_γ, where

γ = (examples generated by the LLM) / (total examples)

so γ = 0 is untouched USPTO-50K and γ = 1 replaces every label with a model prediction. The dataset is published at scaliaven/uspto-50k.

What the γ sweep shows

We score models on two metrics. Token accuracy is the fraction of correct output tokens. Molecular accuracy is the fraction of predictions where the whole molecule is right, reported at top-1, top-3 and top-5. Molecular accuracy is the metric that matters: a prediction can get almost every token right and still describe a molecule that is chemically wrong.

Molecular accuracy climbs steadily with γ — top-1 goes from 0.5879 on the untouched data to 0.7418 when every label is generated — while token accuracy drifts slightly the other way. Our reading is that the LLM’s outputs, coming from a model trained on far more chemistry, are more self-consistent than raw patent labels and therefore easier for a small model to fit. The practical payoff is that this data costs compute hours instead of lab time.

Which parameters are worth updating

The model is a 6-layer Transformer encoder-decoder, embedding size 512, feedforward dimension 2048. Against a full-parameter finetuning baseline we tried freezing the encoder or the decoder, inserting MLPs between encoder and decoder (mid), after the decoder (end), or both (all), and LoRA on the output projection layers of the multi-head attention blocks. Every run used the original USPTO-50K, 100 epochs, batch size 128, and took roughly four hours on a single RTX 8000 on NYU HPC.

Without LoRA, freezing the decoder (0.5113) is far better than freezing the encoder (0.2556), and adding MLPs after the decoder recovers most of the remaining gap (0.5355). Adding LoRA inverts that ordering: the best configuration we found is a frozen encoder with LoRA on the decoder, at 0.5534 top-1 and 0.9752 token accuracy — within a point of full finetuning while updating a small fraction of the weights. Once LoRA gives the decoder a cheap way to adapt, holding the encoder fixed becomes the better trade.

Caveats and future work

The γ sweep ran at batch size 64 rather than the 128 used everywhere else; we hit out-of-memory at the intended setting. Batch size matters enough here that the γ numbers are not directly comparable to the finetuning-treatment numbers — which is why γ = 0 (0.5879) reads higher than the full-parameter baseline (0.5634) on what is otherwise the same data.

Two directions we did not get to: distilling from LlaSMol into the BART model rather than only training on its sampled outputs, so the loss accounts for the larger model’s behaviour; and better decoding, since token accuracy sits near 0.97 across every configuration while molecular accuracy is much lower.

Artifacts

• Report — NLP_final.pdf
• Code — scaliaven/NLP_project
• Dataset — scaliaven/uspto-50k on Hugging Face