AI research engineer coding interview questions

As asked

Write a function that computes scaled dot-product attention given Q, K, V tensors and an optional causal mask. Handle the scaling, masking, and softmax correctly. Then describe how you would test it.

Sample answer outline

The function should compute scores as Q @ K.transpose(-2, -1) / sqrt(d_k), add the causal mask (large negative values at masked positions before softmax), apply softmax over the key dimension, then multiply by V. Testing should include shape assertions, checking that the causal mask prevents future token attention by verifying upper-triangle outputs are zeroed post-softmax, and numerical comparison with torch.nn.functional.scaled_dot_product_attention on the same inputs.

Reference implementation (python)

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(
    q: torch.Tensor,  # (batch, heads, seq_q, d_k)
    k: torch.Tensor,  # (batch, heads, seq_k, d_k)
    v: torch.Tensor,  # (batch, heads, seq_k, d_v)
    causal: bool = False
) -> torch.Tensor:
    # TODO: implement this
    pass

Expect these follow-ups

• How would you extend this to multi-head attention?
• What changes are needed to make this numerically stable for very large or very small d_k values?

As asked

Write RMSNorm as a PyTorch Module. The module should take a hidden_size parameter, register a learnable scale parameter, and implement the forward pass correctly. Make sure gradients flow through the scale parameter.

Sample answer outline

The forward pass computes the root mean square of the input along the last dimension, divides the input by it (with a small epsilon for stability), then multiplies by the learned scale parameter. The scale should be initialized to ones and registered as nn.Parameter so PyTorch tracks it for gradient computation. A correctness test can check that all-ones input with unit scale returns all-ones output.

Reference implementation (python)

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        # TODO: register scale parameter
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # TODO: implement
        pass

Expect these follow-ups

• Why is the epsilon typically added inside the square root rather than outside?
• How would you fuse this operation into a single CUDA kernel for efficiency?

As asked

Given a tensor of logits for the next token (shape: vocab_size), implement top-k sampling and top-p (nucleus) sampling. Explain the tradeoff between the two in terms of diversity and quality of generated text.

Sample answer outline

Top-k zeros out all logits except the k highest and samples from the resulting distribution, so the sample set is fixed size. Top-p accumulates sorted probabilities until the cumulative sum exceeds p, keeping only those tokens, so the sample set size adapts to the distribution's sharpness. When the model is confident (peaked distribution), top-p uses fewer tokens; when uncertain, it uses more. A strong answer shows correct handling of sorting, cumulative sum, masking, and final sampling via torch.multinomial.

Reference implementation (python)

import torch

def top_k_sample(logits: torch.Tensor, k: int) -> int:
    """logits shape: (vocab_size,). Returns sampled token index."""
    # TODO
    pass

def top_p_sample(logits: torch.Tensor, p: float) -> int:
    """logits shape: (vocab_size,). Returns sampled token index."""
    # TODO
    pass

Expect these follow-ups

• What happens if you set top-p = 1.0 and temperature = 1.0 simultaneously?
• How does temperature scaling interact with top-p filtering, and should you apply temperature before or after the top-p filter?

As asked

Write a function that applies rotary positional embeddings to a query or key tensor of shape (batch, seq_len, d_head). Use the standard 2D rotation formulation and precomputed cos/sin tables.

Sample answer outline

Precompute frequencies as theta_i = base^(-2i/d) for i in 0..d/2, then compute position-angle products for each (position, frequency) pair. Build cos and sin tables of shape (max_seq_len, d/2). In the forward pass, split the head dimension into pairs, rotate each pair using the (cos, -sin; sin, cos) rotation, and concatenate. A correct implementation should produce relative dot products that depend only on position difference, verifiable by checking that (q at pos 3) dot (k at pos 1) equals (q at pos 5) dot (k at pos 3) for the same base vectors.

Reference implementation (python)

import torch

def precompute_freqs_cis(d_head: int, max_seq_len: int, base: float = 10000.0):
    # returns cos, sin tables of shape (max_seq_len, d_head // 2)
    pass

def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    # x: (batch, seq_len, d_head)
    pass

Expect these follow-ups

• Why do some implementations rotate pairs of adjacent dimensions rather than the first half and second half interleaved?
• What does the base frequency control and how does increasing it affect length generalization?

As asked

Write a function that clips the global gradient norm of all model parameters in place, matching the behavior of torch.nn.utils.clip_grad_norm_. Explain why global clipping is preferred over per-parameter clipping for transformer training.

Sample answer outline

Compute the global L2 norm by summing squared gradients across all parameters, then clip by multiplying each gradient by min(1, max_norm / total_norm). Global clipping preserves the direction of the gradient vector, just reducing its magnitude when it is too large. Per-parameter clipping changes the relative magnitudes between parameters and can distort the update direction, which is especially harmful in transformers where different parameter groups interact closely.

Reference implementation (python)

import torch
from typing import Iterable

def clip_grad_norm_(parameters: Iterable[torch.nn.Parameter], max_norm: float) -> float:
    """Clips gradients in place. Returns the total grad norm before clipping."""
    # TODO
    pass

Expect these follow-ups

• What does torch.nn.utils.clip_grad_norm_ return, and how do you use that return value for logging?
• Should you clip gradients before or after the optimizer step, and why?

As asked

Write a function that computes the KL divergence KL(p || q) for two distributions given as unnormalized logits. Make sure it handles numerical stability correctly. When would you use KL(p||q) versus KL(q||p) in ML?

Sample answer outline

Convert logits to log-probabilities using log_softmax for numerical stability, then compute sum over classes of p * (log_p - log_q). KL(p||q) is zero-forcing: when p is zero the term vanishes, so the distribution q can have high probability where p is zero without penalty. KL(q||p) is mass-covering: it penalizes q for being zero where p is not, making it prefer distributions that cover all modes of p. KL(p||q) is used in RLHF to keep the policy close to the reference; KL(q||p) appears in variational inference as part of the ELBO.

Reference implementation (python)

import torch
import torch.nn.functional as F

def kl_divergence(p_logits: torch.Tensor, q_logits: torch.Tensor) -> torch.Tensor:
    """
    Compute KL(p || q) from unnormalized logits.
    Both tensors shape: (..., vocab_size)
    """
    # TODO
    pass

Expect these follow-ups

• Why does torch.nn.KLDivLoss expect log-probabilities as input rather than raw probabilities?
• What is the forward-reverse KL choice in variational autoencoders and why does it lead to mode-seeking versus mean-seeking behavior?

As asked

Implement beam search with a given beam width for autoregressive generation. Given a function next_token_logits(token_ids) that returns logits for the next token, generate a sequence up to max_len tokens and return the top beam.

Sample answer outline

Maintain a list of (score, token_ids) beams. At each step, call next_token_logits for each beam, take log_softmax, add the beam score, and collect all (beam_score + token_score, sequence) candidates. Keep the top-beam_width candidates by score, pruning the rest. Handle end-of-sequence tokens by adding completed beams to a finished list and not expanding them further. Return the sequence with the highest score from finished beams. A strong answer handles length normalization correctly.

Reference implementation (python)

from typing import Callable, List, Tuple
import torch

def beam_search(
    initial_ids: List[int],
    next_token_logits: Callable[[List[int]], torch.Tensor],
    beam_width: int = 4,
    max_len: int = 50,
    eos_token_id: int = 2,
) -> List[int]:
    """Returns the highest-scoring complete sequence."""
    # TODO
    pass

Expect these follow-ups

• Why does beam search sometimes produce repetitive output even with a large beam width?
• How does length penalty in beam search change what the optimal beam looks like?

As asked

Implement a LoRALinear module that wraps a frozen base Linear layer and adds trainable low-rank matrices A and B. The module should initialize correctly, forward correctly, and support merging the LoRA weights into the base weight for inference.

Sample answer outline

Freeze the base weight, initialize A with kaiming_uniform and B with zeros so the initial LoRA update is zero. Forward computes base(x) + (x @ A.T @ B.T) * (alpha / r). The merge method adds A.T @ B.T * (alpha / r) to the base weight and sets the LoRA matrices to zero (or a merged flag). A correct implementation should verify that unmerged and merged outputs are numerically identical.

Reference implementation (python)

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int, alpha: float = 1.0):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False
        # TODO: add A, B matrices
        self.r = r
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # TODO
        pass

    def merge_weights(self):
        # TODO
        pass

Expect these follow-ups

• Why must B be initialized to zero and not A?
• How do you apply LoRA only to the Q and V projections and not to K and the output projection?

As asked

Implement cross-entropy loss with label smoothing from scratch using PyTorch tensor operations, without using nn.CrossEntropyLoss directly. Explain what label smoothing does to the loss surface and why it helps generalization.

Sample answer outline

Label smoothing replaces hard one-hot targets with a soft distribution: (1 - epsilon) mass on the correct class and epsilon / (num_classes - 1) on each wrong class. Compute log_softmax of logits, then take the negative dot product with the soft target distribution summed over the class dimension and averaged over the batch. Label smoothing prevents the model from becoming overconfident, reduces the magnitude of the logit for the correct class that the model is pushed toward, and improves calibration. A correct implementation matches nn.CrossEntropyLoss(label_smoothing=epsilon) numerically.

Reference implementation (python)

import torch
import torch.nn.functional as F

def cross_entropy_label_smoothing(
    logits: torch.Tensor,   # (batch, num_classes)
    targets: torch.Tensor,  # (batch,) integer class indices
    epsilon: float = 0.1
) -> torch.Tensor:
    """Returns scalar mean loss."""
    # TODO
    pass

Expect these follow-ups

• When would you NOT use label smoothing, for example in a distillation setup?
• How does label smoothing interact with temperature scaling at inference time for calibration?

As asked

Given a query matrix of shape (n, d) and a database of shape (m, d), compute the top-k nearest neighbors for each query using cosine similarity. Write an efficient vectorized PyTorch implementation and describe its memory complexity.

Sample answer outline

L2-normalize both query and database tensors along the feature dimension, then compute similarity as the matrix product (n, d) @ (d, m) = (n, m). Call torch.topk with dim=1 on the result to get the top-k indices and scores for each query. Memory complexity is O(n * m) for the full similarity matrix, which can be prohibitive for large m. A strong answer notes that batching the query dimension or using approximate nearest neighbor libraries (FAISS) is necessary at scale.

Reference implementation (python)

import torch

def cosine_top_k(
    queries: torch.Tensor,   # (n, d)
    database: torch.Tensor,  # (m, d)
    k: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """Returns (scores, indices) each of shape (n, k)."""
    # TODO
    pass

Expect these follow-ups

• How would you handle m being too large to fit in GPU memory?
• What is the tradeoff between exact cosine similarity and approximate nearest neighbor search with FAISS?