Running biology models: why we built JapanFold
Eight open biology models, free to use, on hardware we own and operate in Japan. This is the inference infrastructure behind them, and why serving this workload looks almost nothing like serving a language model.
AlphaFold turned protein structure prediction from an advanced computation handled by a few specialists into research infrastructure that many people can use. More than 200 million predicted structures are now public, and a researcher can start from an existing prediction rather than working a structure out from scratch.
The role of AI in biology has since widened beyond structure. Models can predict how molecules bind and compare the strength of those interactions, design new proteins against a chosen target, and search enormous sequence datasets for promising candidates. Work that used to be done separately is being connected into a single computational workflow, and that is starting to change how research itself is organized.
Computation alone does not finish a drug or a material. But if many candidates can be compared before the bench work begins, a researcher can concentrate limited time and budget on the more promising directions. Increasing the number of hypotheses that can be examined in the same period is increasing the chances of a discovery that leads to a new medicine, enzyme, or material.
We built JapanFold to make these capabilities usable in Japanese research. It brings together eight open models covering structure prediction, molecular interaction, binder design, and protein search and classification, reachable from a browser, a common API, or an agent. All of it is free to use.
Predicting a single protein structure means searching enormous sequence databases, computing relationships between amino acids, and generating several candidate structures, which consumes a great deal of memory and compute. Each model expects different data and a different runtime.
Processing thousands or tens of thousands of predictions at the speed research needs takes reusing computation that would otherwise be repeated, routing work according to the size of the input, and an operational foundation that keeps many compute resources busy. Turning a good model into a research result takes not only the model, but the inference infrastructure that keeps it running.
ai& has been building that infrastructure and providing the compute behind it, so that frontier models can be used in real research and development. Getting the capability of a model to the people who need it is part of what we are for.
So with JapanFold, ai& carries the cost of inference, and researchers, universities, and institutes in Japan can use every model free of charge. We believe that an environment where researchers here can try more ideas accelerates progress in biology, and leads to new treatments, new molecules, and the breakthroughs that follow them.
JapanFold is ai&'s first substantial effort in AI for Science.
01 How the workload differs from language-model serving
Serving a folding model and serving a language model sit at opposite ends of one physical constraint. A language model runs out of memory bandwidth; a folding model runs out of arithmetic.
A language model produces one token at a time, each depending on all the tokens before it. That property, autoregression, is what the serving stack is built around. Attention over a sequence of length L costs on the order of L², so recomputing the full history at every step is untenable. The KV cache keeps the keys and values already computed, and each new token attends to that stored state. The cost moves into memory, because the cache grows with every token generated and every request in flight.
Most of the past few years of serving work has gone into managing that trade: paged attention treating the cache like virtual memory, continuous batching admitting requests mid-flight, speculative decoding verifying several tokens behind a small draft model. More recent architectures compress the cache into a low-rank latent, or replace most of it with a fixed-size recurrent state. All of them work against one limit: producing a token reads the full weight set, roughly one floating-point operation for every byte moved.
A folding model is built differently. It is not autoregressive, so there is no sequential loop to accelerate and no state that grows as the request runs. In its place is a representation over pairs of residues, refined by repeated passes through the same stack, then a diffusion sampler that produces every atom coordinate at once and repeats a couple hundred times. The size of the expensive tensor is fixed the moment the target arrives.
The expensive object is not the weights. They are modest, a few hundred million parameters, read once and then reused for the entire request. It is the activation, the N×N pair grid. Updating the pair (i, j) requires a sum over every intermediate residue k, so each entry takes part in on the order of N operations before it is written back. At a few hundred tokens that is already hundreds of floating-point operations for every byte, two or three orders of magnitude above decode.
The grid supplies its own parallelism as well. One request presents hundreds of thousands of independent pair updates at once, where a decode step presents one token, so a folding request does not need a batch assembled around it to fill the machine. It arrives full. Which is why the techniques above have so little to attach to. They exist to rescue a processor starved of work, and this one is not starving.
02 Inside a folding model
A single request is hundreds of forward passes. Take Boltz-2, representative of the current open generation.
Tokenization works about the way it does in a language model, except the units are physical: one token per amino acid in the protein (a residue, in the literature), per base in a nucleic acid, and per atom of any small molecule bound to it (a ligand). A 512-residue protein with a drug-like molecule attached comes to a few hundred tokens.
The model then needs context, and none of it arrives in the request. It comes from a database search that pulls up thousands of related proteins from other species and lines them up position by position. Where two positions have changed together over evolutionary time, they are probably touching in the folded structure, and that correlation is most of the signal the model has to work with. The search is a CPU and disk job that runs before the accelerator does anything, and its quality bounds the quality of everything after it.
The trunk is where this stops resembling a transformer stack. Alongside the usual per-token embedding it carries a second tensor with an entry for every ordered pair of tokens, 128 channels deep, and sixty-four layers update both together. That pair tensor is not attention over a sequence but attention over a grid: the characteristic operation updates the cell (i, j) using cells (i, k) and (k, j), for every k. It is a triangle closure, and nearly all the arithmetic in the model goes there. The whole stack then runs again on its own output, several times over, because one pass does not settle the geometry. That is recycling, and it multiplies the trunk's cost by the number of passes.
Only then do coordinates appear, and they come out of a diffusion model, the same family that generates images, conditioned on the trunk rather than on a text prompt. It starts from random noise in three dimensions and denoises toward atom positions over a couple hundred steps, moving every atom on every step. Samples disagree and the first is rarely the best, so a real job draws several and ranks them. Two small heads then read the finished representations. One scores how far to trust the prediction, which is what lets a screen throw out weak results unattended. The other returns a binding probability and an affinity number, which is what a screening campaign actually consumes.
Those passes are not evenly costed. A large share of the wall-clock is the database search at the front, which never touches the accelerator, and most of the rest is the repeated denoiser.
03 The pair representation and how cost scales
Why is the pair representation the expensive part? Because of what a structure is.
To place atoms in space you need distances, and a distance is a property of two amino acids rather than of either one. A per-token embedding has nowhere to put it: whatever you store on token i cannot say how far i sits from j without naming j as well. So the model keeps a vector for every pair, and that is an N×N object before anyone has chosen an architecture.
The second cost follows because distances constrain one another. If residue i sits close to j, and j sits close to k, then i cannot be far from k. A model that ignores this proposes pairwise distances that no arrangement of atoms in three dimensions can satisfy. Enforcing the constraint means updating each pair using information from every third residue, and there are N³ such combinations. Triangle updates and triangle attention are that operation, and the name is literal.
Memory grows with N² because geometry is pairwise, and compute grows with N³ because consistency is three-way.
Boltz-2 keeps 128 channels for every ordered pair, in bf16. For a 512-token complex that is about 262,000 pairs and 34 million numbers, and the triangle operations need several copies resident at once. On a long complex the first thing to break is an out-of-memory error, well before latency becomes the constraint.
04 Serving inference at scale
A DGX B200 runs these models beautifully. Eight Blackwell GPUs carry 192 GB of HBM3e each at 8 TB/s, around 9 petaFLOPS of dense FP8 per GPU, and NVLink at 1.8 TB/s between them. The pair tensor and its working copies fit with room to spare. If you want to serve folding models today, it is great.
The trouble is the curve it sits on. A screening campaign multiplies one prediction by the ligand count, and a design campaign multiplies it again by the candidates it generates and filters.
From there it is a scheduling problem. Targets are batched only within a length bucket, the alignment search is cached per sequence and reused across a campaign, and the independent diffusion samples run at once rather than in sequence. The precision boundary the authors set, bf16 through the trunk and fp32 in the denoiser, is held exactly: a stack that quantized across that line would be serving a different model.
The fit is imperfect in one way. The B200's defining specification is that 8 TB/s of memory bandwidth, and it exists because autoregressive decode consumes it. A folding trunk does not, so a large share of what the machine costs, and of the power it draws, goes to a memory subsystem this workload barely exercises. That would be tolerable if HBM were cheap.
If the binding constraint is activation memory and dense arithmetic rather than weight bandwidth, the chip you want looks more like a large SRAM part than like a GPU: die area spent on arithmetic and on-chip memory rather than on a wide external memory interface, and links to its neighbours cheap enough that splitting a pair tensor across devices is not punished. Tenstorrent's Blackhole has that shape: small cores each holding their own SRAM, a native Ethernet fabric between chips, and GDDR6 in place of HBM.
Property | NVIDIA DGX B200 | Tenstorrent Blackhole Galaxy |
|---|---|---|
Accelerators per node | 8 Blackwell GPUs | 32 Blackhole ASICs |
Memory per accelerator | 192 GB HBM3e | GDDR6 |
On-chip memory | smaller on-die SRAM | large distributed SRAM |
Interconnect | NVLink 1.8 TB/s, NVSwitch | native Ethernet mesh |
Precision | bf16 trunk, fp32 denoiser | bf16 trunk, fp32 denoiser |
On a 512-residue protein, one Blackhole Galaxy reaches the throughput of a DGX B200 at about a fifth of the purchase price, roughly five times the throughput per dollar, holding workload and accuracy fixed. It is also why the demo endpoint carries no charge.
Chips of this shape suit this workload, and Tenstorrent have made one usable. The harder part is keeping any of them in production, which only pays off if the workload holds still long enough to amortize the work. Ours does. Anything that does not fit runs on GPUs.
05 Try it now
JapanFold is live at japanfold.aiand.com, and the three ways in share one backend. The workbench at workbench.japanfold.aiand.com needs nothing installed. The API takes the same request shape for every model at POST api.japanfold.com/v1/predictions. The skill, npx skills add japanfold/japanfold, hands that endpoint to a coding agent, so a screening loop is written once and runs unattended.
The models behind the endpoint differ. The guide marks whether each claim was measured by us or reported by the model's authors.
Boltz-2 affinity capability reported by authors
Use when you need a complex and a binding-affinity estimate from one call.
Avoid when the target is an antibody-antigen interface.
Tradeoff seconds against the hours a physics simulation takes, in exchange for an estimate.
ESMFold-2
Use when you are folding tens of thousands of single sequences for throughput.
Avoid when you need a complex or the last point of accuracy.
Tradeoff no alignment search means no CPU stage in front of the accelerator. Some accuracy goes with it.
OpenDDE DockQ rates reported by authors
Use when the interface is antibody-antigen. It is the strongest open model there: DockQ success of 51.0% on PXMeter-AB, 70.0% on FoldBench-AB, 66.4% on 2026ARK-AB.
Avoid when the target is a small-molecule ligand; the authors' own numbers put it behind AlphaFold3 and Protenix-v1 there.
BoltzGen binder rate reported by authors
Use when you are designing a binder: protein, peptide, or nanobody.
Tradeoff the authors report nanomolar binders for 66% of targets. Design is generative, so filtering and validation sit downstream.
The other four are narrower. Protenix-v2 predicts all-atom complexes in the AlphaFold3 lineage; reach for it, or Boltz-2, on ligand work. RFdiffusion3 generates backbones and scaffolds. ESMC (up to 6B) gives protein language-model embeddings. SaProt folds a structural token into the sequence, so structure rides inside the embedding.
Sources
Boltz-2: Towards Accurate and Efficient Binding Affinity Prediction: the bf16/fp32 precision boundary.
BoltzGen: Toward Universal Binder Design: the 66%-of-targets figure.
tt-bio benchmarks: throughput per dollar on Tenstorrent Blackhole Galaxy.
Author-reported figures are attributed inline; memory and compute scaling is derived; the roofline and latency-split figures are schematic.