ProSE: the architecture and design of a protein discovery engine

A model that reads protein sequences is structurally identical to one that reads English — the only differences are the weights and the fact that a protein is 300 to 2000 tokens where a sentence is 30 to 80 — and that purely quantitative change turns out to shift which operations dominate enough that GPUs and TPUs fall off a cliff; ProSE (Duke) builds an accelerator from systolic arrays of three different sizes, each matched to one recurring operation sequence, with no scratchpad at all.

The argument in one glance

Problem

The models that read proteins are too expensive to run widely

Language models trained on amino acid sequences now predict protein structure, stability and binding, work that would otherwise take a laboratory. The obstacle to using them everywhere is the cost of inference, which grows sharply with the length of the sequence.
What people did

Run them on the accelerators built for machine learning

Graphics processors with tensor cores and tensor processing units with large systolic arrays are the standard answer, and they are excellent at the models they were designed around.
Why it fell short

Both are provisioned for a different shape of problem

They target either convolutional models or language models with short inputs. A protein is an order of magnitude longer than a sentence, which changes what fraction of the time goes to matrix multiplication, and the matrices involved fall between what a large systolic array and a small tensor core each assume.
This paper's idea

Size the hardware to the operation sequences the model actually runs

Three sequences of operations account for roughly nine tenths of the work. Build three kinds of systolic array, one per sequence, each sized to the matrices it will see and each able to switch into an element-wise mode — and drop the scratchpad entirely.

Why this problem is worth solving

The same architecture reads two very different kinds of text.

A protein is a chain of amino acids, and there are about twenty of them, so a protein can be written as a string over a twenty-letter alphabet. That makes it look enough like language that the models built for language work on it: the same architecture, the same training procedure, with each amino acid taking the place of a word. Pre-training on vast numbers of known sequences produces a model that has learned something about how proteins are put together, which is then fine-tuned for specific questions — whether a protein will fluoresce, whether it will hold its shape, how strongly two proteins bind, or what three-dimensional structure it folds into. Nothing about the network differs from a language model except the weights and one other thing: a sentence is 30 to 80 tokens and a protein is 300 to 2000 or more, because functional units within a protein influence each other across long distances and cannot be modelled separately.

The machines built for this work are grids of multipliers with data flowing through them.

A systolic array is a square grid of small units, each holding a multiplier, an adder and an accumulator. Values enter along two edges and march through the grid a step at a time, and at every step every unit multiplies what it received and adds it to a running total. Nothing is fetched from memory in between, which is why the arrangement is so efficient. A design choice with large consequences is which operand stays put: the widely deployed tensor processing unit holds the weights in place and streams the data past them, which requires a large on-chip buffer to hold everything being streamed and the partial results coming out. That buffer is substantial silicon and substantial power. The alternative is to hold the output in place, letting the accumulator inside each unit be the only intermediate storage there is.

What people did before, and where it broke

Run the model on a graphics processor with tensor cores. The standard route, and the one with the mature software behind it. Where it broke. Its tensor cores are provisioned for very small matrix operations, so the matrices this model actually uses are far larger than the unit’s natural size and performance suffers; measured on a current part, efficiency falls below one inference per second per watt once sequences pass a few hundred tokens.
Run it on a tensor processing unit. Large systolic arrays are excellent for large matrix multiplications and this is nominally a matrix-multiplication workload. Where it broke. Its arrays are 128 units across and most of this model’s multiplications involve smaller matrices, so the array spends its time starting up and draining rather than computing, and it is left underutilised.
Accelerate the attention mechanism specifically. Several designs target the part of the model everyone identifies as the bottleneck. Where it broke. They accelerate one operation and leave the rest of the model on commodity hardware, which caps what the whole system can achieve.
Use an accelerator built for sparse matrix multiplication. These handle irregular matrices well and can be applied to transformers. Where it broke. They need substantial sparsity to pay off, and they accelerate only the matrix multiplications, which at long sequence lengths are barely a third of the runtime.
Use an accelerator built for convolutional networks. The most mature category of machine learning accelerator by far. Where it broke. A great deal of their silicon exists to make convolutional filters fast, and transformers do not have convolutional filters.
The gap. Every option was provisioned for a different shape of problem — either a different model family, or the same model family at a tenth of the sequence length. Nobody had asked what changes when the input gets ten times longer.

How sequence length changes what hardware you want

From one large array for everything to three sizes matched to three sequences

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Making the input ten times longer does not make the same workload bigger; it makes it a different workload. This is the observation the design rests on, and it is easy to state and easy to overlook. A protein language model is structurally identical to a language model — the layers, the attention, the training procedure are the same, and only the weights differ — so the natural assumption is that hardware good at one is good at the other. Profiling across sequence lengths from 32 to 2048 says otherwise. Matrix multiplication, which everyone provisions for, falls from about half the runtime to about a third, while element-wise additions and divisions and the softmax climb to fill the gap. Meanwhile the multiplications themselves involve matrices that sit awkwardly between the two standard assumptions: too small for a systolic array 128 units across to use without spending its time filling and draining, too large for a tensor core built around a handful of elements. Both commodity platforms are mismatched at once, in opposite directions, and the result is that inference efficiency collapses past a few hundred tokens. Quantifying the gap between what a model spends its time on and what the accelerator was provisioned for is exactly the exercise the full-stack accelerator search work covered elsewhere on this site performs when it finds a tensor processing unit at 15 percent utilisation because operations taking 5 percent of the arithmetic consume 65 percent of the time.

If the model runs three chains of operations, build three kinds of array. Rather than one size fitting everything, the design groups roughly nine tenths of the work into three recurring sequences and provisions a distinct array type for each. The sequence handling the model’s largest matrices gets a small number of 64-wide grids. The one dominated by divisions and exponentials gets many 16-wide grids, and the reason is a ratio worth noticing: a grid of side n has n squared multipliers but only n element-wise units along its edge, so the smaller the grid the richer it is in exactly the units an element-wise-heavy chain needs. The choice of sizes and counts then comes from sweeping 238 configurations at a fixed budget of arithmetic units, which is how the final mix of two large, three medium and twenty small arrays is arrived at rather than asserted. Deriving a heterogeneous arrangement from measurements of what the workload actually asks for, instead of replicating one homogeneous tile, is the same move the reconfigurable array pruning work covered elsewhere on this site makes when it takes a uniform grid and cuts each tile down to what the application statistics justify.

Removing the scratchpad is where most of the efficiency comes from, and it is also the dependency. The conventional design holds weights stationary in the array and streams data past them, which requires a large on-chip buffer both to hold the streamed data and to catch partial results — substantial silicon, and power drawn continuously. This design inverts it: the output stays in place, the 32-bit accumulator inside each unit is the only intermediate storage, and everything else streams in from the host across a high-bandwidth link. That is what produces two orders of magnitude better efficiency per watt than a tensor processing unit, and the paper attributes it directly to eliminating that buffer. The same choice creates the design’s one hard requirement: with nothing stored locally, performance is set by how fast the host can feed it, and the larger configurations do not become compute-bound until the link reaches 360 gigabytes per second. Stripping the local bookkeeping out of a device and depending on the far side of a fast link to supply it is the same bargain the disaggregated memory work covered elsewhere on this site strikes when it removes per-client state from a memory box and pushes that work onto the clients.

Following one operation sequence through the array

A matrix multiply followed by a scaled addition, in one pass

The array is 16, 32 or 64 units on a side depending on which chain it serves. Nothing is written to a scratchpad because there is no scratchpad.

Step 1 of 5

Stream both operand matrices in from the host

Values enter along two edges and march through the grid, each unit multiplying what arrives and adding to its accumulator. Because the output stays in place, a large multiplication is decomposed into pieces the grid can hold and the partial sums build up where they are.

Two eight-deep buffers per array absorb the arrival latency of the link. An optional partial input buffer allows one operand to be reused across steps rather than re-streamed, which matters when the link is the constraint.

Why it works

Sequence length changes the mix, not just the size. Everyone provisions for matrix multiplication because it dominates at short lengths → but profiling from 32 to 2048 tokens shows it falling from about half the runtime to about a third while element-wise work and softmax rise → so hardware tuned for the short-input case is spending its transistors on the shrinking part of the problem.
Two standard assumptions bracket the actual matrix sizes. A large systolic array wants matrices at least as wide as itself and a tensor core wants very small ones → and this model’s multiplications sit between them, too small for one and too large for the other → so both commodity platforms are simultaneously mismatched, which is why a middle size wins.
Small grids are proportionally rich in element-wise units. A grid of side n holds n squared multipliers but only n units along the edge that can do element-wise work → so the ratio of element-wise capability to arithmetic improves as grids shrink → which is why the chain full of divisions and exponentials gets many 16-wide arrays while the chain of large multiplications gets a few 64-wide ones.
Rotating leftward is what makes the two modes interleave without stalling. Chaining a multiplication into element-wise arithmetic normally means draining results and fetching them back → but shifting the grid left feeds the column that finished first into the element-wise units immediately → so the element-wise work starts before the rest of the grid has completed and the pipeline never empties.
Nothing needs a scratchpad if the accumulators are the storage. A weight-stationary array must buffer both the streamed data and the partial results, which is a large and continuously powered structure → whereas an output-stationary array keeps every intermediate value in the accumulator where it was produced → giving two orders of magnitude better efficiency per watt, at the cost of depending on the host link for everything.

What it costs (trade-offs and scope)

Nothing was fabricated. The arrays are written in a hardware description language, synthesised at a 15-nanometre predictive process and scaled to 7 nanometres, with buffers modelled separately at 45 nanometres and scaled the same way. Timing comes from a cycle-accurate simulator the authors wrote.
The power comparison is asymmetric in a way that matters. The graphics processor’s power was measured at 395 watts; the tensor processing unit’s is taken from published thermal design power, because no measurement tool is publicly available. So the headline figures of 173 and 249 times better efficiency per watt compare a synthesised estimate against a datasheet number for a four-chip device. The paper states this, and does conservatively exclude host and interconnect power from the commodity platforms while including it for its own design.
The comparison is per instance, not per unit of hardware. The tensor processing unit instance measured has roughly sixteen times the arithmetic units of the design being compared to it. That is disclosed and it does mean the speedup figure is not a like-for-like efficiency of the arithmetic.
The most element-wise chain still leaves the accelerator. Exponentials are computed on the array but the summation and division of the softmax run on the host processor — described as trading performance for hardware simplicity, and it is the chain that grows most as sequences lengthen.
Everything rests on host bandwidth. With no local storage, the larger configurations remain bandwidth-bound until the link reaches 360 gigabytes per second, and the design space exploration statically partitions link lanes among the three array types. On a slower interconnect the argument does not hold.
The power argument assumes intermediate data stays in host cache. The claim is that intermediates mostly fit in the host’s last-level cache so memory is touched only on cold misses. At the batch size and sequence length evaluated, that is an assumption doing real work in the efficiency result.
The comparison against the tensor processing unit’s internals is a reconstruction. The paper says plainly that the exact microarchitectural steps are not published and that it extrapolates from block diagrams, so the step-by-step contrast that motivates several design choices is the authors’ reading.
The lookup table truncation is asserted rather than measured. Both functions are stored only over a limited exponent range, with the claim that this does not affect accuracy on the models studied and no accompanying numbers.
The biological validation is small and the paper says so. A downstream binding-affinity model trained on 39 antibody variants and tested on 35 achieves a rank correlation of 0.5161, described as not quite as accurate as a physics-based simulator but sufficient for validity, with the argument that better protein models will close the gap.
Inference only, and encoder only. Training is not addressed, which is where much of the cost of these models actually sits.
The protein framing motivates the length rather than the architecture. The paper is explicit that these models are structurally identical to language models, so what is being built is a long-sequence transformer accelerator; proteins supply the reason long sequences are obligatory rather than optional.

The evaluation

Three ways to run a transformer over a 512-token sequence

Property	Tensor cores on a graphics processor	One large systolic array with a buffer	Three sizes of array, no buffer
How well the arrays match the matrices
Where intermediate results live	Registers and on-chip memory
Element-wise work and activations	Separate execution paths
Speedup at 512 tokens
Efficiency per watt	The comparison point

A twelve-layer protein transformer at 512 tokens and batch 128, with the accelerator synthesised at a predictive 15-nanometre process and scaled to 7 nanometres, and timing from a cycle-accurate simulator including host communication.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

The efficiency figures deserve the most care. Nothing was fabricated: the design is synthesised at a predictive process and scaled, with timing from a simulator the authors wrote. More consequentially, the power comparison is asymmetric — the graphics processor’s consumption was measured at 395 watts while the tensor processing unit’s comes from published thermal design power, because no measurement tool exists publicly. So the striking figures of 173 and 249 times better efficiency per watt set a synthesised estimate against a datasheet number for a four-chip device, which the paper states while also noting it conservatively omits host and interconnect power from the commodity platforms. The speedup is likewise per instance rather than per unit of arithmetic, against a tensor unit holding roughly sixteen times the multipliers. Two design choices leave work undone: the softmax summation and division still run on the host, in the very chain that grows most as sequences lengthen, and the whole no-scratchpad argument requires a link fast enough that the larger configurations stay compute-bound only above 360 gigabytes per second, with the power claim additionally assuming intermediates fit in the host’s last-level cache. The step-by-step contrast against the tensor unit’s internals that motivates several choices is explicitly a reconstruction from block diagrams. The lookup table truncation is asserted not to affect accuracy without numbers. And the biological validation is modest by the paper’s own account — a rank correlation of 0.5161 from 39 training and 35 test sequences, less accurate than a physics-based simulator, with the case resting on protein models improving.

Background

Why length is the variable that matters

It is worth understanding why sequence length is such a sharp lever, because the effect is not linear. In the attention mechanism at the heart of these models, every token is compared against every other token, so the work and the intermediate storage both grow with the square of the sequence length while the rest of the model grows linearly. At 30 to 80 tokens that quadratic term is small enough to ignore and the model looks like a pile of large matrix multiplications, which is the workload every accelerator was designed around. At 300 to 2000 tokens the quadratic term takes over, and with it comes a change in character: the score matrix it produces is large, short-lived, and passed through a softmax, so a growing share of the cost is moving that matrix and applying element-wise operations to it rather than multiplying anything. That is why the attention dataflow work covered elsewhere on this site concentrates on when to materialise that matrix at all rather than on making the multiplications faster. For proteins the length is not a choice — functional regions of a protein influence one another across hundreds of positions, so the sequence must be modelled whole — which makes this the regime rather than an edge case, and makes an accelerator built for the short-input regime the wrong tool.

If you remember only three things

1. A model reading protein sequences is structurally identical to one reading English, differing only in weights and in the fact that a protein is 300 to 2000 tokens where a sentence is 30 to 80 — and that purely quantitative change shifts matrix multiplication from about half the runtime to about a third while element-wise work and softmax rise, so hardware provisioned for short inputs is spending its transistors on the shrinking part of the problem.
2. The multiplications themselves fall between the two standard assumptions, too small for a systolic array 128 units across to use without filling and draining and too large for a tensor core, so both commodity platforms are mismatched at once; grouping roughly nine tenths of the work into three recurring operation chains and provisioning arrays of 64, 32 and 16 units for them — chosen by sweeping 238 configurations at a fixed arithmetic budget — matches the hardware to what the model asks for.
3. Each grid switches into rotating its contents leftward so the column that finished its multiplication first immediately feeds element-wise units at the edge, chaining a multiply into the arithmetic that follows it in one pass; and with output held stationary the accumulators are the only intermediate storage, which removes the large on-chip buffer entirely and is where two orders of magnitude of efficiency per watt come from — at the price of depending on a very fast host link for everything.

Check yourself

Check your understanding

1. Why do commodity accelerators handle protein transformers badly when they handle language transformers well?
2. Why are the smallest arrays assigned to the chain full of divisions and exponentials?
3. What does removing the on-chip buffer accomplish, and what does it require?
4. Which aspect of the evaluation deserves the most caution?

Sources

The paper is “ProSE: The Architecture and Design of a Protein Discovery Engine” by Eyes Robson, Ceyu Xu and Lisa Wu Wills, published at ASPLOS 2022, pages 655-668, DOI 10.1145/3503222.3507722. The authors are at Duke University.

This explainer is written from the full 14-page text. Everything above comes from the paper’s own sections: the background contrasting human language and protein transformers, with their identical structure and their differing token lengths and downstream tasks; the software validation on an antibody binding-affinity task with its rank correlation and its comparison against a physics-based simulator; the profiling of the model across sequence lengths from 32 to 2048 showing the shifting operation mix, and the observations about matrix sizes relative to a tensor processing unit array and a graphics tensor core; the experiment comparing homogeneous and heterogeneous array collections at equal arithmetic budget; the identification of three recurring operation chains covering roughly nine tenths of the work and their mapping onto three array types; the design of the dual-mode array with its matrix-multiply and element-wise rotating modes and the reasoning for rotating leftward; the sizing argument relating grid side length to the ratio of element-wise units to multipliers; the processing element microarchitecture with its streaming buffers, accumulator-based intermediate storage and optional partial input buffer; the step-by-step contrast against a reconstructed tensor processing unit microarchitecture for both a matrix multiply and a scaled addition; the truncated lookup tables for the two special functions with their exponent ranges and sizes; the multithreaded orchestration across array types; the implementation and evaluation methodology covering the hardware description toolchain, the synthesis processes and scaling, the simulator, and the power measurement approach for each platform; and the design space exploration over 238 configurations, the sweep over arithmetic budgets and link bandwidths, and the resulting speedup and efficiency comparisons.