← TikTok Interview Insights

TikTok·Machine Learning Engineer·Technical Phone Screen·Senior

SeniorPrefer not to say
Jun 2026

Summary

TikTok ML engineer interview that went pretty deep on Transformer internals. They wanted pseudocode, not just hand-waving, which I wasn't fully prepared for. The concept questions at the end felt almost like a relief after writing out attention by hand.

Questions Asked (4)

Q1

Write pseudocode for scaled dot-product self-attention for a single head, including tensor shapes, Q/K/V projections, scaling, softmax, and optionally an attention mask.

Technical Trade-offsAlgorithms & Data Structures
Author's notes

This is where I spent most of my mental energy.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by defining the input tensors and their shapes, then walk through the linear projections to obtain Q, K, and V. Compute the scaled dot-product scores, apply optional masking, normalize with softmax, and finally multiply by V to get the output. Emphasize the scaling factor and mask handling as critical details.

Pro tip: Mention that the scaling factor 1/sqrt(d_k) prevents the dot products from growing too large, which would push softmax into regions with tiny gradients. Also, note that masking is typically done by adding a large negative value (e.g., -1e9) to the scores before softmax, and that this is essential for autoregressive decoding and padding.

1. Define input tensors and shapes

Specify the input tensor X of shape (batch_size, seq_len, d_model) and the projection weight matrices W_Q, W_K, W_V each of shape (d_model, d_k).

2. Compute Q, K, V projections

Perform linear projections: Q = X @ W_Q, K = X @ W_K, V = X @ W_V, resulting in tensors of shape (batch_size, seq_len, d_k).

3. Compute scaled dot-product scores

Calculate attention scores as S = (Q @ K^T) / sqrt(d_k), giving shape (batch_size, seq_len, seq_len).

4. Apply optional attention mask

If a mask is provided (e.g., for padding or causal attention), add a large negative value to masked positions in S to zero out their softmax probabilities.

5. Apply softmax and compute output

Compute attention weights A = softmax(S, dim=-1), then output O = A @ V, resulting in shape (batch_size, seq_len, d_k).

Key Points to Mention

  • Tensor shapes at each step, especially the batch and sequence dimensions.
  • The scaling factor 1/sqrt(d_k) and its purpose in stabilizing gradients.
  • How masking is implemented (e.g., adding -inf or a large negative number before softmax).
  • The softmax is applied along the last dimension (key dimension).
  • The final output shape matches the value dimension d_k.
  • Optionally, mention that this is for a single head and multi-head would involve concatenating or averaging heads.

AI-generated suggestions, not part of the candidate's original notes. May be inaccurate — verify before relying on them.

Q2

Write pseudocode for multi-class cross-entropy loss over a batch, handling both class index labels and one-hot labels, returning a scalar mean loss.

Algorithms & Data StructuresTechnical Trade-offs
Author's notes

Felt more comfortable here.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by clarifying the input shapes and label formats, then outline the mathematical steps: compute logits, apply log-softmax for numerical stability, gather the correct class log-probabilities, and average the negative log-likelihood. Write clean pseudocode that handles both class indices and one-hot labels, and mention edge cases like ignoring padding or class weights.

Pro tip: Emphasize numerical stability by using log-softmax instead of separate softmax and log operations, and explicitly state that you would use a gather operation for index labels to avoid unnecessary one-hot expansion. This shows you understand both correctness and efficiency in production ML code.

1. Clarify inputs and outputs

Confirm that logits have shape (batch_size, num_classes), labels can be either class indices (batch_size,) or one-hot (batch_size, num_classes), and the output is a scalar mean loss.

2. Compute log-probabilities stably

Use the log-softmax trick: subtract the max logit per example for numerical stability, then compute log_sum_exp and subtract it from each logit.

3. Extract correct class log-probabilities

If labels are indices, gather the log-probability at the true class index for each example. If labels are one-hot, compute the dot product between log-probabilities and labels.

4. Compute and average the loss

Take the negative of the extracted log-probabilities to get per-example losses, then return the mean over the batch as a scalar.

5. Handle edge cases and optimizations

Mention ignoring padding indices (e.g., -100), applying class weights if needed, and using efficient operations like gather instead of one-hot expansion for index labels.

Key Points to Mention

  • Numerical stability via log-softmax (subtracting max logit and using log-sum-exp)
  • Efficient handling of class index labels using gather instead of one-hot expansion
  • Support for both label formats with a conditional branch or unified formulation
  • Averaging over the batch to produce a scalar mean loss
  • Edge cases: ignoring padding indices, class weights, and reduction options
  • Vectorized operations for batch processing to avoid loops

AI-generated suggestions, not part of the candidate's original notes. May be inaccurate — verify before relying on them.

Q3

What does the position-wise feed-forward network in a Transformer block actually do, and why is it necessary given that attention already exists?

Technical Trade-offsSystem Design
Author's notes

I gave the textbook answer about attention mixing information across positions while the FFN applies a nonlinear transformation independently at each position.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by clearly defining the FFN's role as a position-wise transformation that adds non-linear capacity and feature mixing after attention aggregates context. Then explain why attention alone is insufficient: it's a weighted sum (linear) and lacks per-token non-linearity and channel-wise interactions. Finally, connect to practical benefits like increased model capacity, better representation learning, and efficiency via parameter sharing across positions.

Pro tip: Emphasize that the FFN operates independently on each position, which is crucial for parallelization and efficiency, and mention that the expansion factor (e.g., 4x) is a key hyperparameter that balances capacity and compute. This shows you understand both theory and practical trade-offs.

1. Define the FFN

Describe the FFN as two linear layers with a non-linear activation (e.g., ReLU, GELU) applied identically to each position. Mention the typical expansion factor (e.g., 4x) and that it's position-wise (shared weights across positions).

2. Explain its function

State that it introduces non-linearity and enables channel-wise interactions, transforming each token's representation independently. It acts as a key-value memory or feature enhancer after attention.

3. Contrast with attention

Highlight that attention is a linear weighted sum that mixes information across positions but lacks non-linear transformations. Without FFN, the model would be limited to linear combinations of value vectors, reducing expressive power.

4. Discuss necessity

Explain that the FFN adds model capacity and depth, allowing the network to learn complex functions. It also provides a bottleneck that can compress and expand representations, improving generalization.

5. Connect to practical implications

Mention that the FFN is computationally efficient due to parallelization across positions and that its design (e.g., expansion factor) is a trade-off between performance and resource usage, relevant for large-scale systems like TikTok.

Key Points to Mention

  • Position-wise operation: same FFN applied to each token independently, enabling parallelization.
  • Non-linearity: introduces activation functions (ReLU, GELU) that attention lacks.
  • Channel mixing: transforms features within each token, complementing attention's cross-token mixing.
  • Increased capacity: adds parameters and depth, crucial for learning complex patterns.
  • Expansion factor: typically 4x, balancing expressiveness and computational cost.
  • Ablation studies: removing FFN significantly degrades performance, proving its necessity.

AI-generated suggestions, not part of the candidate's original notes. May be inaccurate — verify before relying on them.

Q4

Why do we divide attention scores by the square root of the head dimension before applying softmax?

Technical Trade-offsAlgorithms & Data Structures
Author's notes

Knew this one cold.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by explaining that dividing by sqrt(d_k) prevents the dot products from growing too large in magnitude, which would push softmax into regions with tiny gradients. Then connect this to training stability and the mathematical reasoning behind the scaling factor, emphasizing its role in maintaining well-behaved attention distributions.

Pro tip: Mention that without scaling, the softmax saturates and gradients vanish, making training unstable—this shows you understand the practical implications beyond just the formula. Also, note that the scaling factor is derived from the variance of the dot product under the assumption of independent components, demonstrating deeper mathematical insight.

1. Define the problem

Explain that attention scores are computed as dot products between queries and keys, and that as the dimension d_k grows, the variance of these dot products increases.

2. Explain the softmax issue

Describe how large magnitude scores cause the softmax to saturate, leading to extremely small gradients and hindering learning.

3. Introduce the scaling factor

State that dividing by sqrt(d_k) normalizes the variance to approximately 1, keeping the scores in a range where softmax gradients are healthy.

4. Connect to training stability

Emphasize that this scaling is crucial for stable and efficient training of Transformer models, preventing vanishing gradients and allowing deeper architectures.

5. Mention practical implications

Note that this is a standard practice in Transformer implementations and that understanding it is key for debugging and optimizing attention mechanisms.

Key Points to Mention

  • Dot product variance grows with dimension d_k
  • Softmax saturation leads to vanishing gradients
  • Scaling by sqrt(d_k) normalizes variance to ~1
  • Derivation assumes independent components with zero mean and unit variance
  • Training stability and convergence speed
  • Standard practice in Transformer architectures (e.g., BERT, GPT)

AI-generated suggestions, not part of the candidate's original notes. May be inaccurate — verify before relying on them.