Generative model engineer coding interview questions

As asked

Write a Python class that implements a linear DDPM noise scheduler. It should precompute the alpha_bar values for T timesteps, implement a method to add noise to a batch of images at an arbitrary timestep in a single step using the closed-form formula, and implement a method to compute the posterior mean and variance for one reverse step.

Sample answer outline

A strong implementation precomputes betas as a linspace, then alphas = 1 - betas, then alpha_bars as the cumprod. add_noise uses the reparameterization: x_t = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise. The reverse step uses the posterior mean formula from the DDPM paper. The candidate should handle tensor shapes correctly and not use loops for the forward process.

Reference implementation (python)

import torch
import torch.nn as nn

class DDPMScheduler:
    def __init__(self, num_timesteps: int = 1000,
                 beta_start: float = 1e-4,
                 beta_end: float = 0.02):
        self.T = num_timesteps
        # TODO: compute betas, alphas, alpha_bars
        pass

    def add_noise(self, x0: torch.Tensor,
                  noise: torch.Tensor,
                  t: torch.Tensor) -> torch.Tensor:
        # TODO: one-shot forward process q(x_t | x_0)
        pass

    def step(self, model_output: torch.Tensor,
             t: int, x_t: torch.Tensor) -> torch.Tensor:
        # TODO: one reverse step p(x_{t-1} | x_t)
        pass

Expect these follow-ups

• How would you modify this to support a cosine schedule?
• What changes are needed to support classifier-free guidance at inference time?

As asked

Write a PyTorch module that wraps an existing nn.Linear layer with LoRA adapters. The wrapper should freeze the original weights, add the low-rank matrices A and B with proper initialization (A random normal, B zeros), apply the scaled update during the forward pass, and expose a merge() method that absorbs BA into the base weight.

Sample answer outline

A strong answer freezes the base weight with requires_grad=False, initializes A with random normal (Gaussian) as in the original LoRA paper and B with zeros, so the initial LoRA delta is zero and training starts from the pretrained behavior. The forward pass returns base(x) + (x @ A.T @ B.T) * (alpha / rank). The merge() method adds the fused delta weight in-place and removes A and B. The candidate should handle dtype consistency and explain why B is initialized to zero and A to random normal rather than the reverse.

Reference implementation (python)

import torch
import torch.nn as nn
import math

class LoRALinear(nn.Module):
    def __init__(self, base_layer: nn.Linear,
                 rank: int = 4,
                 alpha: float = 1.0):
        super().__init__()
        self.base = base_layer
        self.rank = rank
        self.alpha = alpha
        d_out, d_in = base_layer.weight.shape
        # TODO: freeze base, init A and B
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # TODO: base forward + scaled LoRA delta
        pass

    def merge(self) -> nn.Linear:
        # TODO: absorb BA into base weight
        pass

Expect these follow-ups

• How would you extend this to support QLoRA where the base weight is stored in 4-bit NF4?
• What happens if you call merge() and then try to fine-tune again?

As asked

Given a 2D tensor of raw logits (batch x vocab_size), implement nucleus sampling that: applies temperature scaling, filters to the top-p cumulative probability mass, and returns sampled token indices. Handle the edge case where a single token already has probability greater than p.

Sample answer outline

Sort logits descending, compute softmax, cumsum, and mask tokens where cumsum exceeds p after shifting by one position to always keep at least the top token. Set masked logits to -inf, then sample with torch.multinomial. A strong answer correctly handles the off-by-one in the cumsum threshold so the nucleus always includes enough tokens to exceed p, not exactly reach it.

Reference implementation (python)

import torch
import torch.nn.functional as F

def top_p_sample(logits: torch.Tensor,
                 temperature: float = 1.0,
                 top_p: float = 0.9) -> torch.Tensor:
    """
    logits: (batch, vocab_size)
    returns: (batch,) sampled token indices
    """
    # TODO: scale by temperature
    # TODO: sort descending, compute cumulative probs
    # TODO: mask tokens outside nucleus
    # TODO: sample and return
    pass

Expect these follow-ups

• How would you vectorize this across the batch dimension efficiently?
• What is the difference between top-p filtering before vs after temperature scaling?

As asked

Write a function that applies Rotary Position Embedding to a batch of query or key tensors of shape (batch, seq_len, n_heads, head_dim). The function should precompute the rotation frequencies and apply the complex rotation in a vectorized way without using Python loops.

Sample answer outline

Precompute frequencies theta_i = 1 / (base^(2i/d)) for i in 0..d/2. Compute position-angle products via outer product of positions and freqs. Apply rotation using the identity [x1,x2] -> [x1*cos - x2*sin, x1*sin + x2*cos] on pairs of head dimensions. A strong answer reshapes tensors correctly to avoid any loop, handles the (batch, heads, seq, dim) vs (batch, seq, heads, dim) layout, and optionally uses torch complex number view for cleanliness.

Reference implementation (python)

import torch

def apply_rotary_emb(
    x: torch.Tensor,        # (batch, seq, n_heads, head_dim)
    positions: torch.Tensor, # (batch, seq)
    base: int = 10000
) -> torch.Tensor:
    """
    Apply RoPE in-place style, returns rotated x.
    head_dim must be even.
    """
    # TODO: compute per-dim frequencies
    # TODO: build cos/sin tables for each position
    # TODO: rotate pairs of dimensions
    pass

Expect these follow-ups

• How do you cache the cos/sin values to avoid recomputing them at every forward pass?
• What changes when you extend the base frequency for context length scaling?

As asked

Write the forward pass for grouped-query attention where the number of key-value heads (n_kv_heads) is less than the number of query heads (n_heads). The function should handle the head grouping correctly and remain compatible with a standard KV cache.

Sample answer outline

The key insight is that each KV head is shared across n_heads / n_kv_heads query heads. Implementation: project Q to (batch, n_heads, seq, head_dim), project K and V to (batch, n_kv_heads, seq, head_dim), then repeat_interleave K and V along the heads dim by the ratio, then proceed with standard scaled dot-product attention. A strong answer avoids materializing the expanded K/V in memory if possible by using einsum or batched matmul with appropriate reshaping.

Reference implementation (python)

import torch
import torch.nn.functional as F

def grouped_query_attention(
    q: torch.Tensor,   # (B, n_heads, T, head_dim)
    k: torch.Tensor,   # (B, n_kv_heads, T, head_dim)
    v: torch.Tensor,   # (B, n_kv_heads, T, head_dim)
    mask: torch.Tensor = None
) -> torch.Tensor:
    n_heads = q.shape[1]
    n_kv_heads = k.shape[1]
    assert n_heads % n_kv_heads == 0
    # TODO: expand k, v to match n_heads
    # TODO: scaled dot-product attention
    # TODO: return (B, n_heads, T, head_dim)
    pass

Expect these follow-ups

• Why does GQA reduce memory more during decoding than during prefill?
• How does MQA (multi-query attention with n_kv_heads=1) sit at the extreme of this design space?

As asked

Write a function that computes the Frechet Inception Distance between a set of real images and generated images. Use a pretrained Inception-v3 model to extract features, compute the mean and covariance of each set, then apply the Frechet distance formula. Handle the matrix square root numerically.

Sample answer outline

Extract pool3 features from Inception-v3 for both sets. Compute mu1, sigma1, mu2, sigma2 using np.cov with rowvar=False. The Frechet distance is ||mu1-mu2||^2 + trace(sigma1 + sigma2 - 2*sqrtm(sigma1@sigma2)). The matrix square root should use scipy.linalg.sqrtm and handle small imaginary parts from numerical error. A strong answer batches feature extraction and handles the edge case where sigma has near-zero eigenvalues by adding a small epsilon to the diagonal.

Reference implementation (python)

import numpy as np
from scipy.linalg import sqrtm
import torch

def compute_fid(real_features: np.ndarray,
                gen_features: np.ndarray) -> float:
    """
    real_features, gen_features: (N, 2048) Inception pool3 features
    """
    # TODO: compute means
    # TODO: compute covariances
    # TODO: matrix square root of product
    # TODO: return Frechet distance
    pass

Expect these follow-ups

• Why does FID require at least 10k images to be statistically reliable?
• How would you adapt this to compute CLIP-FID instead of Inception-FID?

As asked

You have a training dataset with N domains of very unequal sizes. Write a PyTorch Dataset wrapper that samples from each domain with a configurable weight, oversampling small domains and undersampling large ones, without loading everything into memory.

Sample answer outline

A strong solution maintains a list of per-domain datasets and samples indices by first drawing a domain from a categorical distribution parameterized by the weights, then drawing a uniform index within that domain. __len__ should return the total weighted epoch size (typically the largest domain size times number of domains). This avoids materializing the full sample list in memory. The candidate should handle reproducibility via seeded RNG and explain why torch.utils.data.WeightedRandomSampler works differently (it requires the full weight list upfront).

Reference implementation (python)

import torch
from torch.utils.data import Dataset
from typing import List
import numpy as np

class WeightedMultiDomainDataset(Dataset):
    def __init__(self,
                 datasets: List[Dataset],
                 weights: List[float],
                 epoch_size: int):
        self.datasets = datasets
        self.weights = np.array(weights) / sum(weights)
        self.epoch_size = epoch_size

    def __len__(self) -> int:
        return self.epoch_size

    def __getitem__(self, idx: int):
        # TODO: sample domain, then sample item within domain
        pass

Expect these follow-ups

• How do you handle domain weights that should change during training (curriculum)?
• What is the risk of setting weights too aggressively toward small domains when they are lower quality?

As asked

Write the forward pass for a single attention layer that supports both prefill (processing a full sequence) and decode (generating one token at a time) modes. The function should update and use a pre-allocated KV cache in-place, returning the output token embedding.

Sample answer outline

Pre-allocate cache as (batch, n_heads, max_seq, head_dim) tensors for K and V. On each decode call, write the new K and V at position cache_pos, then attend over cache[:, :, 0:cache_pos+1]. During prefill, fill positions 0..T-1 in one pass with causal mask. A strong answer handles the attention mask correctly so decode tokens only attend to filled cache positions, and avoids any recomputation of past keys.

Reference implementation (python)

import torch
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class KVCache:
    k: torch.Tensor  # (batch, n_heads, max_seq, head_dim)
    v: torch.Tensor
    pos: int = 0

def attention_with_cache(
    q: torch.Tensor,     # (batch, 1, n_heads, head_dim) for decode
    k_new: torch.Tensor,
    v_new: torch.Tensor,
    cache: KVCache,
) -> torch.Tensor:
    # TODO: write k_new, v_new into cache at cache.pos
    # TODO: attend q over filled cache positions
    # TODO: increment cache.pos, return output
    pass

Expect these follow-ups

• How would you modify this to support paged KV cache where blocks are non-contiguous in memory?
• What are the memory layout implications of storing KV cache in (batch, seq, heads, dim) vs (batch, heads, seq, dim) order?

As asked

Diffusion models typically maintain an EMA copy of the weights for inference. Write an EMA wrapper class that updates a shadow copy of model parameters after each training step, with configurable decay, and exposes a method to temporarily swap EMA weights in for evaluation.

Sample answer outline

The EMA class holds a copy of all parameters as plain tensors (no grad). After each step, shadow = decay * shadow + (1 - decay) * param.data. A context manager or swap() method copies shadow params into the model and restores originals afterward. A strong answer mentions that decay should be warmed up from a lower value early in training (since EMA with high decay but few steps is biased toward initialization), and handles buffers (like BatchNorm running stats) separately.

Reference implementation (python)

import torch
import torch.nn as nn
from contextlib import contextmanager
from copy import deepcopy

class EMAModel:
    def __init__(self, model: nn.Module, decay: float = 0.9999):
        self.decay = decay
        self.shadow = deepcopy(model)
        self.shadow.requires_grad_(False)

    def step(self, model: nn.Module) -> None:
        # TODO: update shadow params with EMA
        pass

    @contextmanager
    def ema_scope(self, model: nn.Module):
        # TODO: swap EMA weights in, yield, restore
        pass

Expect these follow-ups

• Why is EMA with decay 0.9999 standard for diffusion models but 0.999 is typical for classifiers?
• How does bias correction for EMA work and when does it matter?

As asked

Write the DPO (Direct Preference Optimization) loss function given log probabilities of chosen and rejected completions under both the policy model and reference model. The function should return the DPO loss scalar and the implicit reward margin for logging.

Sample answer outline

The DPO loss is -log_sigmoid(beta * (log_ratio_chosen - log_ratio_rejected)) where log_ratio = log pi_theta(y|x) - log pi_ref(y|x). A strong implementation handles per-token log probabilities by summing them and normalizes by sequence length to prevent length bias. The reward margin is the inner term before the sigmoid, and logging it lets you track whether the policy is learning to prefer chosen completions.

Reference implementation (python)

import torch
import torch.nn.functional as F

def dpo_loss(
    policy_log_probs_chosen: torch.Tensor,   # (batch,) sum of token log probs
    policy_log_probs_rejected: torch.Tensor, # (batch,)
    ref_log_probs_chosen: torch.Tensor,      # (batch,)
    ref_log_probs_rejected: torch.Tensor,    # (batch,)
    beta: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Returns (loss, reward_margin)
    """
    # TODO: compute log ratio for chosen and rejected
    # TODO: DPO objective
    pass

Expect these follow-ups

• How does the IPO loss differ from DPO and what instability does it address?
• Why is it important to normalize log probabilities by sequence length in DPO?

As asked

Write a single denoising step function that applies classifier-free guidance. The function receives a diffusion model, a noise sample, the current timestep, a text conditioning tensor, and a guidance scale. It should run two forward passes (conditional and unconditional), combine them with CFG, and return the predicted noise to use for the DDIM update.

Sample answer outline

The implementation runs the model twice: once with the text conditioning and once with an empty/null conditioning vector, then computes: noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond). An optimized version batches both forward passes by concatenating along the batch dimension and splitting the output, halving the number of model calls. A strong answer notes the null conditioning is usually a zero-padded or BOS-only token sequence, not a literal zero tensor.

Reference implementation (python)

import torch

def cfg_step(
    model,
    x_t: torch.Tensor,          # (B, C, H, W) noisy latent
    t: torch.Tensor,            # (B,) timestep
    text_embeds: torch.Tensor,  # (B, seq, dim)
    null_embeds: torch.Tensor,  # (B, seq, dim)
    guidance_scale: float = 7.5,
) -> torch.Tensor:
    """
    Returns guided noise prediction.
    """
    # TODO: batch both forward passes
    # TODO: split and apply CFG formula
    pass

Expect these follow-ups

• How would you implement negative prompting on top of this CFG implementation?
• What is the maximum batch size implication of batching the two forward passes together?

As asked

Write a function that computes the alpha_bar values for a cosine noise schedule as described in the Improved DDPM paper. The function should return alpha_bar values for T timesteps such that the signal-to-noise ratio decreases smoothly from near 1 at t=0 to near 0 at t=T, with a configurable offset to avoid singularity near pure noise.

Sample answer outline

The cosine schedule defines alpha_bar(t) = f(t)/f(0) where f(t) = cos((t/T + s) / (1 + s) * pi/2)^2 and s is a small offset (typically 0.008) to prevent alpha_bar from being too small at t=0. The returned schedule should be clipped so beta_t = 1 - alpha_bar_t / alpha_bar_{t-1} does not exceed 0.999, preventing instability. A strong answer returns both alpha_bars and the corresponding betas for use in the forward and reverse processes.

Reference implementation (python)

import torch
import math

def cosine_alpha_bar_schedule(
    num_timesteps: int = 1000,
    s: float = 0.008,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Returns (alpha_bars, betas) each of shape (num_timesteps,).
    alpha_bars[0] ~= 1, alpha_bars[-1] ~= 0.
    """
    # TODO: compute f(t) = cos(...)^2
    # TODO: normalize to alpha_bar = f(t)/f(0)
    # TODO: derive betas, clip to [0, 0.999]
    # TODO: return (alpha_bars, betas)
    pass

Expect these follow-ups

• How does the cosine schedule compare to linear at t=T in terms of remaining signal?
• What is the zero-terminal-SNR modification and why does it improve inference-time alignment?

As asked

Write a Python class that implements a linear DDPM noise scheduler. It should precompute the alpha_bar values for T timesteps, implement a method to add noise to a batch of images at an arbitrary timestep in a single step using the closed-form formula, and implement a method to compute the posterior mean and variance for one reverse step.

Sample answer outline

A strong implementation precomputes betas as a linspace, then alphas = 1 - betas, then alpha_bars as the cumprod. add_noise uses the reparameterization: x_t = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise. The reverse step uses the posterior mean formula from the DDPM paper. The candidate should handle tensor shapes correctly and not use loops for the forward process.

Reference implementation (python)

import torch
import torch.nn as nn

class DDPMScheduler:
    def __init__(self, num_timesteps: int = 1000,
                 beta_start: float = 1e-4,
                 beta_end: float = 0.02):
        self.T = num_timesteps
        # TODO: compute betas, alphas, alpha_bars
        pass

    def add_noise(self, x0: torch.Tensor,
                  noise: torch.Tensor,
                  t: torch.Tensor) -> torch.Tensor:
        # TODO: one-shot forward process q(x_t | x_0)
        pass

    def step(self, model_output: torch.Tensor,
             t: int, x_t: torch.Tensor) -> torch.Tensor:
        # TODO: one reverse step p(x_{t-1} | x_t)
        pass

Expect these follow-ups

• How would you modify this to support a cosine schedule?
• What changes are needed to support classifier-free guidance at inference time?

As asked

Write a PyTorch module that wraps an existing nn.Linear layer with LoRA adapters. The wrapper should freeze the original weights, add the low-rank matrices A and B with proper initialization (A random normal, B zeros), apply the scaled update during the forward pass, and expose a merge() method that absorbs BA into the base weight.

Sample answer outline

A strong answer freezes the base weight with requires_grad=False, initializes A with random normal (Gaussian) as in the original LoRA paper and B with zeros, so the initial LoRA delta is zero and training starts from the pretrained behavior. The forward pass returns base(x) + (x @ A.T @ B.T) * (alpha / rank). The merge() method adds the fused delta weight in-place and removes A and B. The candidate should handle dtype consistency and explain why B is initialized to zero and A to random normal rather than the reverse.

Reference implementation (python)

import torch
import torch.nn as nn
import math

class LoRALinear(nn.Module):
    def __init__(self, base_layer: nn.Linear,
                 rank: int = 4,
                 alpha: float = 1.0):
        super().__init__()
        self.base = base_layer
        self.rank = rank
        self.alpha = alpha
        d_out, d_in = base_layer.weight.shape
        # TODO: freeze base, init A and B
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # TODO: base forward + scaled LoRA delta
        pass

    def merge(self) -> nn.Linear:
        # TODO: absorb BA into base weight
        pass

Expect these follow-ups

• How would you extend this to support QLoRA where the base weight is stored in 4-bit NF4?
• What happens if you call merge() and then try to fine-tune again?

As asked

Given a 2D tensor of raw logits (batch x vocab_size), implement nucleus sampling that: applies temperature scaling, filters to the top-p cumulative probability mass, and returns sampled token indices. Handle the edge case where a single token already has probability greater than p.

Sample answer outline

Sort logits descending, compute softmax, cumsum, and mask tokens where cumsum exceeds p after shifting by one position to always keep at least the top token. Set masked logits to -inf, then sample with torch.multinomial. A strong answer correctly handles the off-by-one in the cumsum threshold so the nucleus always includes enough tokens to exceed p, not exactly reach it.

Reference implementation (python)

import torch
import torch.nn.functional as F

def top_p_sample(logits: torch.Tensor,
                 temperature: float = 1.0,
                 top_p: float = 0.9) -> torch.Tensor:
    """
    logits: (batch, vocab_size)
    returns: (batch,) sampled token indices
    """
    # TODO: scale by temperature
    # TODO: sort descending, compute cumulative probs
    # TODO: mask tokens outside nucleus
    # TODO: sample and return
    pass

Expect these follow-ups

• How would you vectorize this across the batch dimension efficiently?
• What is the difference between top-p filtering before vs after temperature scaling?

As asked

Write a function that applies Rotary Position Embedding to a batch of query or key tensors of shape (batch, seq_len, n_heads, head_dim). The function should precompute the rotation frequencies and apply the complex rotation in a vectorized way without using Python loops.

Sample answer outline

Precompute frequencies theta_i = 1 / (base^(2i/d)) for i in 0..d/2. Compute position-angle products via outer product of positions and freqs. Apply rotation using the identity [x1,x2] -> [x1*cos - x2*sin, x1*sin + x2*cos] on pairs of head dimensions. A strong answer reshapes tensors correctly to avoid any loop, handles the (batch, heads, seq, dim) vs (batch, seq, heads, dim) layout, and optionally uses torch complex number view for cleanliness.

Reference implementation (python)

import torch

def apply_rotary_emb(
    x: torch.Tensor,        # (batch, seq, n_heads, head_dim)
    positions: torch.Tensor, # (batch, seq)
    base: int = 10000
) -> torch.Tensor:
    """
    Apply RoPE in-place style, returns rotated x.
    head_dim must be even.
    """
    # TODO: compute per-dim frequencies
    # TODO: build cos/sin tables for each position
    # TODO: rotate pairs of dimensions
    pass

Expect these follow-ups

• How do you cache the cos/sin values to avoid recomputing them at every forward pass?
• What changes when you extend the base frequency for context length scaling?

As asked

Write the forward pass for grouped-query attention where the number of key-value heads (n_kv_heads) is less than the number of query heads (n_heads). The function should handle the head grouping correctly and remain compatible with a standard KV cache.

Sample answer outline

The key insight is that each KV head is shared across n_heads / n_kv_heads query heads. Implementation: project Q to (batch, n_heads, seq, head_dim), project K and V to (batch, n_kv_heads, seq, head_dim), then repeat_interleave K and V along the heads dim by the ratio, then proceed with standard scaled dot-product attention. A strong answer avoids materializing the expanded K/V in memory if possible by using einsum or batched matmul with appropriate reshaping.

Reference implementation (python)

import torch
import torch.nn.functional as F

def grouped_query_attention(
    q: torch.Tensor,   # (B, n_heads, T, head_dim)
    k: torch.Tensor,   # (B, n_kv_heads, T, head_dim)
    v: torch.Tensor,   # (B, n_kv_heads, T, head_dim)
    mask: torch.Tensor = None
) -> torch.Tensor:
    n_heads = q.shape[1]
    n_kv_heads = k.shape[1]
    assert n_heads % n_kv_heads == 0
    # TODO: expand k, v to match n_heads
    # TODO: scaled dot-product attention
    # TODO: return (B, n_heads, T, head_dim)
    pass

Expect these follow-ups

• Why does GQA reduce memory more during decoding than during prefill?
• How does MQA (multi-query attention with n_kv_heads=1) sit at the extreme of this design space?

As asked

Write a function that computes the Frechet Inception Distance between a set of real images and generated images. Use a pretrained Inception-v3 model to extract features, compute the mean and covariance of each set, then apply the Frechet distance formula. Handle the matrix square root numerically.

Sample answer outline

Extract pool3 features from Inception-v3 for both sets. Compute mu1, sigma1, mu2, sigma2 using np.cov with rowvar=False. The Frechet distance is ||mu1-mu2||^2 + trace(sigma1 + sigma2 - 2*sqrtm(sigma1@sigma2)). The matrix square root should use scipy.linalg.sqrtm and handle small imaginary parts from numerical error. A strong answer batches feature extraction and handles the edge case where sigma has near-zero eigenvalues by adding a small epsilon to the diagonal.

Reference implementation (python)

import numpy as np
from scipy.linalg import sqrtm
import torch

def compute_fid(real_features: np.ndarray,
                gen_features: np.ndarray) -> float:
    """
    real_features, gen_features: (N, 2048) Inception pool3 features
    """
    # TODO: compute means
    # TODO: compute covariances
    # TODO: matrix square root of product
    # TODO: return Frechet distance
    pass

Expect these follow-ups

• Why does FID require at least 10k images to be statistically reliable?
• How would you adapt this to compute CLIP-FID instead of Inception-FID?

As asked

You have a training dataset with N domains of very unequal sizes. Write a PyTorch Dataset wrapper that samples from each domain with a configurable weight, oversampling small domains and undersampling large ones, without loading everything into memory.

Sample answer outline

A strong solution maintains a list of per-domain datasets and samples indices by first drawing a domain from a categorical distribution parameterized by the weights, then drawing a uniform index within that domain. __len__ should return the total weighted epoch size (typically the largest domain size times number of domains). This avoids materializing the full sample list in memory. The candidate should handle reproducibility via seeded RNG and explain why torch.utils.data.WeightedRandomSampler works differently (it requires the full weight list upfront).

Reference implementation (python)

import torch
from torch.utils.data import Dataset
from typing import List
import numpy as np

class WeightedMultiDomainDataset(Dataset):
    def __init__(self,
                 datasets: List[Dataset],
                 weights: List[float],
                 epoch_size: int):
        self.datasets = datasets
        self.weights = np.array(weights) / sum(weights)
        self.epoch_size = epoch_size

    def __len__(self) -> int:
        return self.epoch_size

    def __getitem__(self, idx: int):
        # TODO: sample domain, then sample item within domain
        pass

Expect these follow-ups

• How do you handle domain weights that should change during training (curriculum)?
• What is the risk of setting weights too aggressively toward small domains when they are lower quality?

As asked

Write the forward pass for a single attention layer that supports both prefill (processing a full sequence) and decode (generating one token at a time) modes. The function should update and use a pre-allocated KV cache in-place, returning the output token embedding.

Sample answer outline

Pre-allocate cache as (batch, n_heads, max_seq, head_dim) tensors for K and V. On each decode call, write the new K and V at position cache_pos, then attend over cache[:, :, 0:cache_pos+1]. During prefill, fill positions 0..T-1 in one pass with causal mask. A strong answer handles the attention mask correctly so decode tokens only attend to filled cache positions, and avoids any recomputation of past keys.

Reference implementation (python)

import torch
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class KVCache:
    k: torch.Tensor  # (batch, n_heads, max_seq, head_dim)
    v: torch.Tensor
    pos: int = 0

def attention_with_cache(
    q: torch.Tensor,     # (batch, 1, n_heads, head_dim) for decode
    k_new: torch.Tensor,
    v_new: torch.Tensor,
    cache: KVCache,
) -> torch.Tensor:
    # TODO: write k_new, v_new into cache at cache.pos
    # TODO: attend q over filled cache positions
    # TODO: increment cache.pos, return output
    pass

Expect these follow-ups

• How would you modify this to support paged KV cache where blocks are non-contiguous in memory?
• What are the memory layout implications of storing KV cache in (batch, seq, heads, dim) vs (batch, heads, seq, dim) order?

As asked

Diffusion models typically maintain an EMA copy of the weights for inference. Write an EMA wrapper class that updates a shadow copy of model parameters after each training step, with configurable decay, and exposes a method to temporarily swap EMA weights in for evaluation.

Sample answer outline

The EMA class holds a copy of all parameters as plain tensors (no grad). After each step, shadow = decay * shadow + (1 - decay) * param.data. A context manager or swap() method copies shadow params into the model and restores originals afterward. A strong answer mentions that decay should be warmed up from a lower value early in training (since EMA with high decay but few steps is biased toward initialization), and handles buffers (like BatchNorm running stats) separately.

Reference implementation (python)

import torch
import torch.nn as nn
from contextlib import contextmanager
from copy import deepcopy

class EMAModel:
    def __init__(self, model: nn.Module, decay: float = 0.9999):
        self.decay = decay
        self.shadow = deepcopy(model)
        self.shadow.requires_grad_(False)

    def step(self, model: nn.Module) -> None:
        # TODO: update shadow params with EMA
        pass

    @contextmanager
    def ema_scope(self, model: nn.Module):
        # TODO: swap EMA weights in, yield, restore
        pass

Expect these follow-ups

• Why is EMA with decay 0.9999 standard for diffusion models but 0.999 is typical for classifiers?
• How does bias correction for EMA work and when does it matter?

As asked

Write the DPO (Direct Preference Optimization) loss function given log probabilities of chosen and rejected completions under both the policy model and reference model. The function should return the DPO loss scalar and the implicit reward margin for logging.

Sample answer outline

The DPO loss is -log_sigmoid(beta * (log_ratio_chosen - log_ratio_rejected)) where log_ratio = log pi_theta(y|x) - log pi_ref(y|x). A strong implementation handles per-token log probabilities by summing them and normalizes by sequence length to prevent length bias. The reward margin is the inner term before the sigmoid, and logging it lets you track whether the policy is learning to prefer chosen completions.

Reference implementation (python)

import torch
import torch.nn.functional as F

def dpo_loss(
    policy_log_probs_chosen: torch.Tensor,   # (batch,) sum of token log probs
    policy_log_probs_rejected: torch.Tensor, # (batch,)
    ref_log_probs_chosen: torch.Tensor,      # (batch,)
    ref_log_probs_rejected: torch.Tensor,    # (batch,)
    beta: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Returns (loss, reward_margin)
    """
    # TODO: compute log ratio for chosen and rejected
    # TODO: DPO objective
    pass

Expect these follow-ups

• How does the IPO loss differ from DPO and what instability does it address?
• Why is it important to normalize log probabilities by sequence length in DPO?

As asked

Write a single denoising step function that applies classifier-free guidance. The function receives a diffusion model, a noise sample, the current timestep, a text conditioning tensor, and a guidance scale. It should run two forward passes (conditional and unconditional), combine them with CFG, and return the predicted noise to use for the DDIM update.

Sample answer outline

The implementation runs the model twice: once with the text conditioning and once with an empty/null conditioning vector, then computes: noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond). An optimized version batches both forward passes by concatenating along the batch dimension and splitting the output, halving the number of model calls. A strong answer notes the null conditioning is usually a zero-padded or BOS-only token sequence, not a literal zero tensor.

Reference implementation (python)

import torch

def cfg_step(
    model,
    x_t: torch.Tensor,          # (B, C, H, W) noisy latent
    t: torch.Tensor,            # (B,) timestep
    text_embeds: torch.Tensor,  # (B, seq, dim)
    null_embeds: torch.Tensor,  # (B, seq, dim)
    guidance_scale: float = 7.5,
) -> torch.Tensor:
    """
    Returns guided noise prediction.
    """
    # TODO: batch both forward passes
    # TODO: split and apply CFG formula
    pass

Expect these follow-ups

• How would you implement negative prompting on top of this CFG implementation?
• What is the maximum batch size implication of batching the two forward passes together?

As asked

Write a function that computes the alpha_bar values for a cosine noise schedule as described in the Improved DDPM paper. The function should return alpha_bar values for T timesteps such that the signal-to-noise ratio decreases smoothly from near 1 at t=0 to near 0 at t=T, with a configurable offset to avoid singularity near pure noise.

Sample answer outline

The cosine schedule defines alpha_bar(t) = f(t)/f(0) where f(t) = cos((t/T + s) / (1 + s) * pi/2)^2 and s is a small offset (typically 0.008) to prevent alpha_bar from being too small at t=0. The returned schedule should be clipped so beta_t = 1 - alpha_bar_t / alpha_bar_{t-1} does not exceed 0.999, preventing instability. A strong answer returns both alpha_bars and the corresponding betas for use in the forward and reverse processes.

Reference implementation (python)

import torch
import math

def cosine_alpha_bar_schedule(
    num_timesteps: int = 1000,
    s: float = 0.008,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Returns (alpha_bars, betas) each of shape (num_timesteps,).
    alpha_bars[0] ~= 1, alpha_bars[-1] ~= 0.
    """
    # TODO: compute f(t) = cos(...)^2
    # TODO: normalize to alpha_bar = f(t)/f(0)
    # TODO: derive betas, clip to [0, 0.999]
    # TODO: return (alpha_bars, betas)
    pass

Expect these follow-ups

• How does the cosine schedule compare to linear at t=T in terms of remaining signal?
• What is the zero-terminal-SNR modification and why does it improve inference-time alignment?