Research Engineer Paper Implementation Interview: Turn Equations into Tested Code

A research-engineer interview may ask you to implement an idea from a paper. The hard part is rarely translating one displayed equation into three lines of code. It is deciding what each symbol means in a concrete tensor layout, handling masks and numerical behavior, and proving the implementation matches a small reference before optimizing it. A candidate who can narrate those decisions gives the reviewer something stronger than a notebook that happens to run.

The Attention Is All You Need paper provides a useful public example: scaled dot-product attention. PyTorch’s scaled_dot_product_attention reference documents a production API and important implementation details, while its reproducibility note explains why repeated GPU runs are not automatically bit-for-bit identical. These are official or primary technical sources. The interview exercise below is original practice. It is not a report that a named employer uses this exact prompt.

What must be extracted from a paper before coding?

Write down the operation, its input/output contract, and the ambiguities that the equation hides. For scaled dot-product attention, the core expression is softmax(QKᵀ / √d) V. That expression does not alone tell you the batch layout, where heads appear, which dimension is sequence length, what mask convention the task expects, or how dropout should behave in evaluation. Resolve those details from the prompt and the chosen reference before writing an optimized kernel.

A practical contract uses Q with shape [B, H, L, d], K with [B, H, S, d], and V with [B, H, S, d_v]. Here B is batch, H heads, L query positions, and S key/value positions. Q @ K.transpose(-2, -1) then has shape [B, H, L, S]. Softmax runs over the last dimension so each query position distributes weight across the available keys. Multiplying by V yields [B, H, L, d_v]. If an implementation returns [B, L, H, d_v], that might be a later layout conversion, but it should not happen silently.

State whether masks are boolean keep/disallow masks or additive biases, because libraries and prompts differ. State whether the task requires causal masking, padding masking, both, or neither. State the expected dtype and device behavior. These questions are not delay tactics; they prevent a superficially correct implementation from failing when L != S or the batch dimension is greater than one.

Build the simplest correct reference

The following is illustrative pseudocode, with exact mask conventions to be defined by the prompt:

scores = matmul(Q, transpose_last_two(K)) / sqrt(d)
scores = apply_allowed_mask(scores, mask)
weights = softmax(scores, axis=-1)
output = matmul(weights, V)
return output

The division by √d changes the scale of dot products before softmax. A mask that excludes a key should be applied to scores before softmax; zeroing a probability afterward without renormalization changes the probability mass. A stable softmax routine should be preferred to manually exponentiating raw scores. If every key is masked for a query, define the required behavior rather than assuming the library will choose the output you want.

PyTorch’s reference notes that the function may select different implementations and that fused kernels can yield different floating-point results. It also notes that dropout behavior follows the passed probability even in evaluation unless the caller supplies zero. For an interview baseline, begin with no dropout and compare within tolerances. Later, if dropout is required, test the training/evaluation contract explicitly.

Check one query by hand before large random tensors

Use a tiny original numerical example with one batch, one head, one query, and two keys. Let Q = [1, 0], K₁ = [1, 0], K₂ = [0, 1], V₁ = [2, 0], and V₂ = [0, 4]. The key width d is 2. The two scaled scores are 1/√2 ≈ 0.7071 and 0. Their softmax weights are approximately 0.6698 and 0.3302. The output is approximately [1.3395, 1.3210].

This toy case tests the axis of softmax and the order of the value vectors. It also reveals a transpose error: if a program compares the two query coordinates rather than the two keys, the output will not match. Use a tolerance rather than demanding exact decimal equality. Write the expected values by hand or independently calculate them before using your implementation as its own oracle.

Check	Input feature	Expected property	Defect it catches
One query, two keys	Toy vectors above	Output near [1.3395, 1.3210]	Wrong scale, axis, or value order.
Unequal lengths	L = 2, S = 3	Scores shape [B,H,2,3]	Accidental self-attention-only assumption.
Disallowed key	One masked key	Its weight is zero	Mask applied after softmax incorrectly.
Two batches, two heads	Distinct values per slice	No cross-slice mixing	Bad reshape or broadcasting.
All allowed keys equal	Equal scores	Uniform weights	Incorrect normalization dimension.

The table is a compact test plan, not evidence that a particular optimized backend was run. It makes failures interpretable. Random tests can follow, but they should not replace cases whose expected behavior you can explain.

How do you compare with a trusted implementation?

After the small hand case passes, compare your implementation against PyTorch’s torch.nn.functional.scaled_dot_product_attention for shapes, dtype, masks, and dropout settings that both implementations support. Use the same inputs and clearly map your mask semantics to the API’s semantics. Compare output values with a justified floating-point tolerance, not string equality. Test an unmasked case before adding a mask; then add one masking mode at a time.

Be careful with hidden differences. A library may use a fused kernel, a different accumulation order, or a backend-specific algorithm. A small numerical mismatch is not automatically a logic bug. Conversely, a very loose tolerance can conceal a wrong mask. Inspect the maximum absolute and relative errors and the inputs that produce them. For a reference comparison, disable dropout and control device/dtype as much as possible. The PyTorch reproducibility note is a reminder that deterministic settings are conditional, not a blanket guarantee.

If the prompt includes backward behavior, check gradients. Compare autograd gradients for a small differentiable case against a trusted implementation or a numerical finite-difference check where appropriate. Avoid testing only a saturated softmax input, where gradients can be tiny and uninformative. Separate forward-correctness tests from gradient-correctness tests so a failure has a clear cause.

What should you say about masks and edge cases?

A causal mask prevents a query position from attending to future key positions. A padding mask prevents attention to positions that are not meaningful tokens. Their shapes and broadcasting rules matter; a mask for [L, S] might broadcast across batches and heads, while a per-example padding mask might need explicit expansion. Do not rely on implicit broadcasting until you have stated which axes are intended to share a mask.

Check L != S, because cross-attention is not self-attention. Check a single key, where the sole allowed weight should be one. Check at least one masked key and a case where the masked key would otherwise have the highest score; that makes it obvious whether masking happened before softmax. If an entire query row has no allowed keys, the required output or error behavior must be decided by the task. Silent NaNs should not be treated as acceptable unless the specified contract explicitly allows them.

For numeric safety, use an established softmax implementation rather than building exp(scores) / sum(exp(scores)) naïvely. Large positive scores can overflow in naïve exponentiation. Mixed precision introduces additional differences, so make an accurate float32 or float64 baseline before tuning precision. If the prompt asks for an optimized implementation, preserve the baseline and run the same tests against the optimized path. Performance is meaningful only after you know which operation is being accelerated.

Explain the implementation as a reviewable change

Present a short plan before coding: “I will map the paper symbols to tensor shapes, implement a reference forward pass, test a hand-computable case and masks, then compare against the library reference. I will optimize only after those checks pass.” This gives the interviewer checkpoints to discuss. If time is limited, a correct reference and an honest list of unimplemented extensions is stronger than an untested high-performance sketch.

While coding, name dimensions clearly and assert the shape relationships that the operation requires. Explain why the softmax axis is keys, why the scale uses key width d, and why masks are applied to scores. Keep the original paper’s notation visible in comments or a short note, but translate it into the implementation’s actual tensor layout. A comment that repeats the equation without specifying axes is not enough.

When reviewing a failed implementation, begin with a shape trace before reading every line. Print or assert the shapes of scores, weights, and output. Then use the toy case. A shape error often identifies a transpose or reshape defect; a shape-correct but value-wrong output points toward scale, mask, or softmax axis. A failure only in backward tests points toward an in-place mutation, detached tensor, or custom gradient bug. This order makes debugging systematic.

What should the final interview answer include?

Summarize the paper operation in one sentence, state the tensor contract, identify the principal numerical and masking choices, and show evidence from at least two kinds of tests: one hand-computable case and one comparison or property test. If you did not test a GPU backend, say so. If you made a simplifying assumption, such as no dropout or self-attention only, label it. Do not claim paper fidelity from matching one sample output.

A useful extension is to measure speed and memory only after correctness. Benchmark under fixed shapes and hardware, include warm-up where appropriate, and compare the same dtype and mask conditions. An optimized path that silently changes precision or ignores padding is not a fair comparison. This is an inference about a sound evaluation workflow, not a claim that any specific hiring process requires benchmarking.

Practice questions for the equation-to-code workflow

Use the first one as the anchor, then vary the mask, layout, and backward-pass requirements so your reasoning transfers beyond a memorized snippet.

Practice question	What to rehearse
Implement and analyze custom attention	Translate the operation into shapes and outputs.
Implement the Forward Pass of Multi-Head Attention From Scratch	Keep head reshape and projection boundaries explicit.
Implement correct attention masking	Apply mask semantics before normalization.
Debug a transformer training pipeline	Localize shape, value, and training failures.
Implement attention and Transformer with backward pass	Add gradient checks after a correct forward pass.

A good paper-implementation response leaves a reviewer able to answer three questions: what does each tensor axis mean, what result should a tiny example produce, and which test would catch the most likely mistake? That is the bridge from reading a research paper to shipping reliable code.

Sources and further reading

• Vaswani et al.: Attention Is All You Need
• PyTorch: scaled_dot_product_attention
• PyTorch: Reproducibility