Coding & ML Implementation — Interview Cheat Sheet

A: Softmax maps a real-valued vector to a probability distribution (non-negative and sums to 1). Subtracting the maximum is for numerical stability: when $z_i$ is large, $e^{z_i}$ causes floating-point overflow, giving inf and making the softmax output NaN. Subtracting a constant does not change the result mathematically (it cancels in numerator and denominator), but keeps the exponentiation in a safe range.

Follow-up: If all input logits are 0, what does softmax output? A: The uniform distribution $[1/C, 1/C, \dots]$.

A: It accepts logits (unnormalized scores). Internally it applies LogSoftmax followed by NLLLoss, which is mathematically equivalent to softmax → log → cross-entropy, but numerically more stable (avoids log(0)).

Follow-up: If you manually apply softmax first and then pass the result into CrossEntropyLoss, what happens? A: Performance degrades because CrossEntropyLoss treats the softmax output (already in (0,1)) as logits and applies log_softmax again; the compressed range ($e^{p_i}\approx 1$–$2.7$) produces semantically wrong loss values and biased gradients — it is not strictly a "sharpening" effect.

A: Think of a library search: $Q$ (Query) is your search question, $K$ (Key) is the index label on each book, $V$ (Value) is the book's content. Attention computes the similarity between $Q$ and each $K$ as weights, then takes a weighted sum of $V$ to produce "an aggregation of information most relevant to the query."

Follow-up: In self-attention, where do $Q$, $K$, $V$ come from? A: All three are derived from the same input $X$ through different linear projections: $Q = XW^Q$, $K = XW^K$, $V = XW^V$.

A: LoRA (Low-Rank Adaptation) is a parameter-efficient fine-tuning (PEFT) method. Core idea: the weight change $\Delta W$ during fine-tuning is low-rank, so it can be factored as the product of two small matrices $\Delta W = BA$. During training, the original weights $W_0$ are frozen; only $A$ and $B$ are trained, drastically reducing the number of trainable parameters.

Follow-up: How many trainable parameters does LoRA introduce? A: For a linear layer of shape $d_{\text{in}} \times d_{\text{out}}$, LoRA adds $r \times d_{\text{in}} + d_{\text{out}} \times r$ parameters, which is far fewer than $d_{\text{in}} \times d_{\text{out}}$ when $r \ll d$.

A: BatchNorm normalizes along the batch dimension (computing mean/variance of the same feature across different samples) and depends on batch size. LayerNorm normalizes along the feature dimension (computing statistics across different features within the same sample) and does not depend on batch size. Transformers use LayerNorm because sequence lengths vary and inference batch size may be 1.

Follow-up: How does BatchNorm behave at inference time? A: It uses the running mean and running variance accumulated during training rather than the statistics of the current batch.

A: Maintain a window $[l, r]$; the right pointer $r$ advances element by element, and when the window satisfies some condition the left pointer $l$ is shrunk; the answer is updated during expansion/contraction. The key is to clearly define the condition for the window being "valid" and the shrinking strategy. Time complexity is typically $O(n)$ because each element enters and exits the window at most once.

Follow-up: How do you tell whether a problem is suitable for a sliding window? A: The problem has "monotonicity" — when $[l, r]$ does not satisfy the condition, $[l, r+1]$ might; when $[l, r]$ satisfies the condition, $[l+1, r]$ definitely does not.

A: Closed-open: search interval $[lo, hi)$, initialize $hi = n$, loop condition $lo < hi$, hi = mid (no minus 1). Closed-closed: search interval $[lo, hi]$, initialize $hi = n - 1$, loop condition $lo \leq hi$, hi = mid - 1. The closed-open variant is cleaner for finding left/right boundaries and less prone to off-by-one errors.

Follow-up: How do you use binary search to find "the first position $\geq$ target"? A: Use the closed-open template: when arr[mid] < target set lo = mid + 1, otherwise set hi = mid; the final value of lo is the answer.

A: Top-k retains the $k$ tokens with the highest probability. Top-p (Nucleus) retains the smallest set of tokens whose cumulative probability first exceeds $p$ — the set size is adaptive: small when the distribution is concentrated, large when flat. Top-p is more flexible and avoids the problems of Top-k introducing noise when the distribution is sharp or truncating too aggressively when it is flat.

Follow-up: Can top-p and top-k be used together in practice? Which filter is typically applied first? A: Yes, they can be combined; typically top-k is applied first (to narrow the candidate set) and then top-p (cumulative probability cutoff).

A: BFS (Breadth-First Search) uses a queue, expands level by level, and is well-suited for finding shortest paths in unweighted graphs. DFS (Depth-First Search) uses a stack (or recursion), goes as deep as possible before backtracking, and is well-suited for finding connected components, topological sorting, and cycle detection. Time complexity of both is $O(V + E)$.

Follow-up: Can topological sorting be implemented with BFS? A: Yes — this is Kahn's algorithm: maintain a queue of nodes with in-degree 0; repeatedly dequeue a node and decrement the in-degree of its neighbors.

A: Assuming each element of $Q$ and $K$ is i.i.d. with mean 0 and variance 1, each element of $QK^\top$ has variance $d_k$. When $d_k$ is large, dot-product values are large in magnitude, pushing softmax gradients toward zero (outputs approach one-hot), making training difficult. Dividing by $\sqrt{d_k}$ stabilizes the variance near 1.

Follow-up: Are there other scaling approaches? A: Cosine attention uses $\frac{QK^\top}{\|Q\|\|K\|}$; some works also explore learnable temperature parameters.

A: $B = 0$ ensures that $\Delta W = BA = 0$ at the start of training, so the model behaves identically to the pretrained checkpoint and the initial state is not disrupted. If $A$ were also initialized to zero, $\Delta W$ would remain zero forever (zero gradient), and neither $A$ nor $B$ could learn anything. Therefore at least one matrix must be initialized non-zero.

Follow-up: What role does the scaling factor $\alpha / r$ play? A: $\alpha$ controls the absolute magnitude of the LoRA update. With $\alpha$ fixed, increasing $r$ relatively reduces the update magnitude (dividing by a larger $r$), which aids training stability. In practice $\alpha = r$ or $\alpha = 2r$ are common choices.

A: Multiple heads can attend to different types of patterns (e.g., syntactic, semantic, positional relationships), acting as a form of implicit regularization. A single head with large dimension is theoretically equivalent in expressive power, but harder to optimize — the model tends to learn only one type of pattern. Experiments show that multi-head performs better at the same total compute.

Follow-up: How different are the patterns learned by different heads? Are there visualization tools? A: Tools like BERTViz show that different heads do attend to different relationship types (e.g., some focus on adjacent words, others on subject-verb relationships). However, research also finds that some heads are redundant and can be pruned (head pruning).

A: Each iteration requires computing distances from $n$ samples to $K$ centroids ($O(nKd)$); with $i$ iterations the total complexity is $O(nKdi)$. K-Means++ selects initial centroids sequentially with probability $\propto D(x)^2$ ($D(x)$ = distance to the nearest already-chosen centroid), making the initial centroids as spread out as possible. Theoretically it guarantees an expected WCSS within an $O(\log K)$ approximation ratio.

Follow-up: How do you choose $K$? A: Common methods are the Elbow Method (plot WCSS vs $K$ and find the elbow) or the Silhouette Score (measures intra-cluster compactness vs. inter-cluster separation).

A: RMSNorm removes the mean-centering (re-centering) step, keeping only re-scaling. This eliminates the mean computation, reducing compute by roughly 10–15%. Experiments show that removing centering has little effect on performance, especially in large models; as a result mainstream models such as LLaMA and Qwen all adopt RMSNorm.

Follow-up: What is the difference between Pre-Norm and Post-Norm? Why do modern large models mostly use Pre-Norm? A: Pre-Norm applies normalization before the sub-layer; Post-Norm applies it after. Pre-Norm provides more stable gradients and typically converges more easily without requiring learning rate warmup. Post-Norm is theoretically more expressive but requires more careful hyperparameter tuning (e.g., warmup).

A: .view() requires the tensor to be contiguous in memory; otherwise it raises an error; it returns a view sharing the same memory. .reshape() is more flexible: if memory is already contiguous it is equivalent to .view(); otherwise it automatically copies data to make it contiguous. In performance-sensitive code prefer .view() (zero-copy), but ensure .contiguous() is called first.

Follow-up: What operations can make a tensor non-contiguous? A: transpose(), permute(), certain expand() operations. These change the stride without moving data, so the logical order no longer matches the physical memory layout.

A: Temperature $T$ scales the logits: $p_i \propto e^{z_i / T}$. As $T \to 0$ the distribution degenerates to one-hot (greedy decoding); as $T \to \infty$ the distribution approaches uniform (fully random). $T < 1$ sharpens the distribution (more deterministic); $T > 1$ flattens it (more diverse).

Follow-up: What temperature values are typically used in practice? Does the choice vary by task? A: Deterministic tasks such as code generation often use $T \approx 0$–$0.2$; creative writing tasks often use $T \approx 0.7$–$1.0$. Values above 2.0 are generally avoided as output quality degrades.

A: (1) State definition — what does $dp[i]$ represent (e.g., the optimal solution over the first $i$ elements); (2) Transition equation — how $dp[i]$ is derived from previous states; (3) Base case and boundary — the value of $dp[0]$. The most critical step is choosing the right state; a good state definition makes the transition natural and efficient.

Follow-up: What is the difference between memoized search (top-down) and bottom-up tabulation? What are the pros and cons of each? A: Memoization uses recursion + caching, computes only needed states, and produces more intuitive code. Bottom-up uses iterative table-filling, has no recursion overhead, and is easier to space-optimize (rolling arrays). Time complexity is the same for both.

A: Let $p_i = \text{softmax}(z_i)$; then $\frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} - p_j)$ (the Jacobian matrix). When combined with Cross-Entropy $\mathcal{L} = -\log p_y$, the gradient simplifies to:

$$\frac{\partial \mathcal{L}}{\partial z_i} = p_i - \mathbb{1}[i = y]$$

That is, predicted probability minus the true label (one-hot) — an extremely clean form. This is also why PyTorch fuses the two into a single operator.

Follow-up: How does Label Smoothing change this gradient? A: Label smoothing replaces one-hot with $(1 - \epsilon)$ on the correct class and $\epsilon / (C-1)$ on other classes; the gradient becomes $p_i - \hat{y}_i$, preventing the model from becoming over-confident.

A: Main directions include:

• Sparse Attention: Compute attention only within a local window (e.g., Longformer's sliding window) or specific patterns (e.g., BigBird's random + global + local)
• Linear Attention: Approximate softmax with a kernel function and reorder the computation of $QK^\top V$ as $Q(K^\top V)$, reducing complexity to $O(Nd^2)$
• Flash Attention: Does not change the mathematical result; optimizes GPU HBM/SRAM access patterns via tiling and recomputation, reducing memory usage and wall-clock time
• Low-rank approximation: E.g., Linformer projects $K, V$ to a lower dimension

Follow-up: Does Flash Attention change the mathematical result of attention? A: No. Flash Attention is an IO-aware exact algorithm that uses block-wise computation (tiling) and an online softmax trick to produce a result mathematically equivalent to standard attention, while using GPU memory hierarchy more efficiently.

A: Rank $r$ governs the tradeoff between expressive capacity and parameter count. Common range: $r \in [4, 64]$. More complex tasks or more target layers may require larger $r$. $\alpha$ is the scaling factor; the effective learning rate is roughly proportional to $\alpha / r$. A common approach is to fix $\alpha$ (e.g., $\alpha = 16$) and tune $r$, or to set $\alpha = r$ so the scaling factor is 1.

Follow-up: Can different layers be assigned different ranks? What is the benefit? A: Yes — this is the idea behind methods like AdaLoRA. Different layers have different "importance" (analyzed via the singular value spectrum of $\Delta W$ through SVD); allocating higher rank to important layers and lower rank to unimportant ones achieves optimal allocation of the parameter budget.

A: K-Means is guaranteed to converge (WCSS decreases monotonically and is lower-bounded by 0), but it converges only to a local optimum. Both the assignment and update steps decrease or preserve WCSS, so WCSS never increases. However, the objective is non-convex, and different initializations can lead to different local optima. In practice, multiple random restarts (multi-start) are common to pick the best result, or K-Means++ is used to improve initialization quality.

Follow-up: What shape constraints does K-Means impose on clusters? Are there alternative methods? A: K-Means assumes spherical (isotropic) clusters and performs poorly on non-spherical clusters. Alternatives include GMM (Gaussian Mixture Models, which allow elliptical clusters), DBSCAN (density-based, can discover arbitrarily shaped clusters), and Spectral Clustering.

A: During autoregressive generation, computing attention for a new token requires attending over all previous tokens. But the $K, V$ of previous tokens do not change when a new token is added, so they can be cached and reused. The KV Cache stores already-computed $K, V$ matrices; each step only needs to compute $Q, K, V$ for the new token and append the new $K, V$ to the cache. This reduces per-step compute from $O(n^2 d)$ to $O(nd)$ and is a key inference optimization.

Follow-up: How large is the KV Cache memory footprint? What compression methods exist? A: For a model with $L$ layers, $h$ heads, head dimension $d_k$, and sequence length $n$, the KV Cache size is $2 \times L \times h \times d_k \times n \times \text{bytes}$. Compression methods include: GQA (Grouped Query Attention — multiple Q heads share K/V), quantization (reducing KV from FP16 to INT8/INT4), and sparsification (discarding unimportant KV pairs).

A: float('-inf') becomes exactly 0 after softmax ($e^{-\infty} = 0$), so there is zero information leakage. A large negative number (e.g., $-10^9$) produces an extremely small but non-zero probability after softmax, theoretically leaking a trace of information. Additionally, -inf is numerically safer — the appropriate "large negative" threshold differs across precisions, whereas -inf behaves consistently at all precisions.

Follow-up: Besides the causal mask, what other masks are commonly used in attention? A: Padding mask (to mask out padding tokens), Local/Sliding window mask (to limit the attention range), Prefix mask (bidirectional for the prefix portion, causal thereafter), Block diagonal mask (independent attention for multiple sequences).

A: PyTorch's F.cross_entropy uses a more advanced numerically stable implementation — computing LogSumExp directly in log space, avoiding any explicit exp() + log() combination:

$$\log \sum_j e^{z_j} = m + \log \sum_j e^{z_j - m}$$

This is done in a single kernel, operating entirely in log space and avoiding overflow in intermediate results. A hand-written softmax + log + nll pipeline is step-by-step, and the intermediate exp can overflow.

Follow-up: What is the difference between torch.log_softmax and log(softmax(x))? A: Mathematically equivalent but numerically different. torch.log_softmax directly applies the LogSumExp trick internally, avoiding an intermediate softmax that could overflow. With log(softmax(x)), if softmax already overflows to inf, then log(inf) is inf, but the gradient becomes NaN.

A: Full-parameter fine-tuning must communicate the entire $\Delta W$ ($d_{\text{in}} \times d_{\text{out}}$), whereas LoRA only needs to communicate $A$ ($r \times d_{\text{in}}$) and $B$ ($d_{\text{out}} \times r$), totaling $r \times (d_{\text{in}} + d_{\text{out}})$ parameters — far less when $r \ll d$. This is especially important in bandwidth-constrained settings such as cross-device federated learning.

Follow-up: In federated learning aggregation with LoRA, do you aggregate $A, B$ or $\Delta W = BA$? A: Both approaches are viable. Aggregating $A, B$ directly has lower communication cost but is mathematically inexact (the mean of $BA$ does not equal the mean $B$ times the mean $A$). Aggregating $\Delta W$ is mathematically more correct but has higher communication cost. In practice the tradeoff depends on the use case.

A: LoRA assumes that the effective rank of the weight update $\Delta W$ is far smaller than the parameter dimensions — that is, task-adaptation information is concentrated in a small number of directions. The theoretical basis is the concept of "intrinsic dimensionality": in the high-dimensional parameter space of a pretrained model, task adaptation actually occurs on a low-dimensional submanifold.

Failure scenarios:

• The downstream task is far from the pretraining distribution, requiring modification of high-rank structure (e.g., cross-lingual or cross-modal transfer)
• In multi-task settings, the low-rank subspaces of different tasks do not overlap, so their superposition raises the effective rank
• The model itself is small, making the low-rank constraint a capacity bottleneck

Follow-up: Beyond simply increasing $r$, what structured approaches can break the low-rank limitation? A: AdaLoRA adaptively allocates rank per layer (higher rank for important layers); Adapter inserts a bottleneck with a nonlinear activation in the FFN, breaking the limitation of purely linear low-rank expressiveness; methods like (IA)^3 use vector scaling to provide additional expressiveness along a different efficiency-capacity curve.

A: BatchNorm computes mean and variance along the batch dimension; LayerNorm computes them along the feature dimension. The fundamental differences are:

1. Variable sequence length: In NLP, sequences within the same batch have different lengths; padding positions introduce noise into BatchNorm's statistics.
2. Autoregressive generation: At inference time, batch size is often 1, making BatchNorm's statistics extremely unstable.
3. Each token normalized independently: LayerNorm stabilizes the input distribution at each layer without depending on other samples in the batch.

Follow-up: What is the difference between Pre-LN and Post-LN? Why do modern large models mostly use Pre-LN? A: Post-LN: $\text{LN}(x + \text{Sublayer}(x))$; Pre-LN: $x + \text{Sublayer}(\text{LN}(x))$. Pre-LN allows gradients to flow back directly through the residual path (residual stream) without passing through the nonlinear transformation of LN, making training more stable and typically eliminating the need for learning rate warmup. In Post-LN, deep-network gradients must pass through LN, which can cause training instability.

A: The core is IO-aware block computation (tiling): divide $Q, K, V$ into small tiles, perform fused softmax + matrix multiplication on SRAM (on-chip cache), and avoid writing the $N \times N$ attention matrix back to HBM (high-bandwidth memory).

Key techniques:

• Online softmax: Exploit the decomposability of softmax — maintain a running max and running sum while processing each tile, scale results with a correction factor; the final output is mathematically equivalent to global softmax.
• Recomputation: During backward pass, do not store the attention matrix; recompute it from $Q, K, V$ — trading compute for memory.

Follow-up: How does FlashAttention's recomputation compare to general gradient checkpointing? A: Both embody the "trade compute for memory" principle. Gradient checkpointing is a general strategy: selectively not saving intermediate activations, then recomputing them via a forward pass during backprop. FlashAttention's recomputation is more specialized: it targets specifically the large attention matrix tensor and combines with the tiling strategy to reduce not only memory but also HBM access count (IO complexity drops from $O(N^2)$ to $O(N^2 d^2 / M)$, where $M$ is SRAM size). The two can be used together.

A: Full LayerNorm: $y_i = \gamma \cdot \frac{x_i - \mu}{\sigma} + \beta$; RMSNorm simplifies to: $y_i = \gamma \cdot \frac{x_i}{\text{RMS}(x)}$, removing $\mu$ and $\beta$.

Theoretical intuition for why this still works:

• After many linear transformations and activations in a deep network, feature means are often already in a reasonable range, or can be compensated by biases in subsequent layers.
• The aspect of normalization critical to training stability is re-scaling (controlling variance), not re-centering (subtracting the mean).
• Removing the mean computation eliminates one reduce operation; accumulated over a large-scale model this yields a meaningful efficiency gain.

Follow-up: In what situations might removing mean-centering be harmful? A: If a layer's input has a systematic bias that subsequent layers cannot easily compensate, the bias will propagate. The effect may be more pronounced in small or shallow networks. However, the residual connections and deep stacking in large-scale Transformers provide sufficient capacity to compensate, and degradation has almost never been observed in practice.

A: Head collapse leads to redundant multi-head representations — despite having $h$ heads' worth of parameters, the effective number of heads is far fewer than $h$, wasting compute and model capacity and effectively reducing the "effective rank" of the attention.

Detection methods:

• Compute KL divergence or cosine similarity between the attention distributions of different heads.
• Analyze the subspace overlap between $W^Q, W^K$ projection matrices (e.g., principal angles).

Follow-up: How can head collapse be mitigated at training time? A:

• Diversity regularization: Add a regularization term to the loss that encourages diversity across head attention distributions (e.g., penalizing similarity).
• Attention dropout: Apply dropout to attention weights to prevent certain heads from dominating too early.
• Head pruning + retraining: Prune redundant heads after training and then fine-tune, forcing remaining heads to take on more responsibility; this also acts as implicit regularization.

A: In standard CE the target $q$ is one-hot ($q_y = 1, q_{j \neq y} = 0$); label smoothing sets $q_y = 1 - \epsilon, q_{j \neq y} = \frac{\epsilon}{C-1}$.

Gradient change ($\nabla_{z_i} \mathcal{L} = p_i - q_i$):

• Standard CE: $\nabla_{z_y} \mathcal{L} = p_y - 1$, $\nabla_{z_j} \mathcal{L} = p_j$
• Label smoothing: $\nabla_{z_y} \mathcal{L} = p_y - (1-\epsilon)$, $\nabla_{z_j} \mathcal{L} = p_j - \frac{\epsilon}{C-1}$

The model is no longer incentivized to push logits toward infinity, preventing over-confidence.

Generalization improvement: prevents the model from being overly confident in training labels (overfitting noisy labels) and implicitly imposes a soft constraint on logits.

Follow-up: What is the connection between label smoothing and knowledge distillation? A: Both fundamentally replace hard targets with "soft targets" during training. Label smoothing softens targets with a uniform distribution; distillation uses the teacher's output distribution as soft targets. Label smoothing can be viewed as "distillation without a teacher, targeting a uniform distribution." The advantage of distillation is that the teacher's soft targets encode structural information about inter-class similarity (e.g., the cat vs. dog logit gap is smaller than cat vs. airplane), rather than the unstructured uniform distribution.

A: K-Means is the hard-assignment special case of GMM-EM in the limit of isotropic, equal-covariance, $\sigma \to 0$.

Apply these constraints to GMM $p(x) = \sum_k \pi_k \mathcal{N}(x \mid \mu_k, \Sigma_k)$: all components share a fixed isotropic covariance $\Sigma_k = \sigma^2 I$ and equal weights $\pi_k = 1/K$. The soft responsibilities in the E-step become a softmax with temperature $\sigma^2$ (equal weights and covariances cancel normalization constants); letting $\sigma \to 0$ degenerates to hard assignment:

$$\gamma_{ik} = \frac{\exp\!\big(-\|x_i - \mu_k\|^2 / 2\sigma^2\big)}{\sum_j \exp\!\big(-\|x_i - \mu_j\|^2 / 2\sigma^2\big)} \;\xrightarrow{\;\sigma \to 0\;}\; \mathbf{1}\big[\,k = \arg\min_j \|x_i - \mu_j\|^2\,\big]$$

The two steps correspond one-to-one:

• E-step: GMM's soft responsibility $\gamma_{ik}$ → K-Means hard assignment to nearest centroid (the $\sigma \to 0$ limit above).
• M-step: GMM's weighted mean $\mu_k = \frac{\sum_i \gamma_{ik} x_i}{\sum_i \gamma_{ik}}$ degenerates under hard $\gamma$ to the arithmetic mean of points in the cluster — exactly K-Means centroid update; both update only the mean.
• Objective: In this limit, the GMM negative log-likelihood (up to constants involving $\sigma$) is proportional to the K-Means WCSS $\sum_i \min_k \|x_i - \mu_k\|^2$.

Follow-up: Since K-Means is a special case of GMM, when should you use GMM? A: When clusters are non-spherical, vary in size or density, or when soft (probabilistic) membership is needed. GMM additionally learns per-component covariance $\Sigma_k$ and weight $\pi_k$, fitting elliptical clusters of different orientations and scales and providing soft membership and uncertainty via $\gamma_{ik}$. K-Means, by assuming isotropic equal covariance, can only produce spherical (Voronoi) hard partitions. The tradeoff is that GMM has more parameters, is more sensitive to initialization and singular covariances (a component collapsing to a single point causes the likelihood to diverge), and often requires covariance lower bounds or regularization.