← Amazon Interview Insights

Amazon·Machine Learning Engineer·Technical Phone Screen·Senior

Senior
Jun 2026

Summary

Amazon ML engineer technical screen, went deep on normalization techniques in transformers. Pretty brutal if you haven't thought carefully about the math behind LayerNorm recently.

Questions Asked (5)

Q1

Explain Layer Normalization in transformers, including the equation, the roles of gamma and beta, and what epsilon is doing there.

Technical Trade-offsSystem Design
Author's notes

I wrote out the equation fine but fumbled explaining epsilon.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by defining layer normalization and its purpose in stabilizing training. Then present the equation clearly, explaining each component (mean, variance, gamma, beta, epsilon). Finally, discuss why it's preferred over batch normalization in transformers, emphasizing independence from batch size and suitability for sequence data.

Pro tip: Mention that layer normalization is applied per sample and per feature, making it effective for variable-length sequences, and note that epsilon prevents division by zero and controls numerical stability.

1. Define Layer Normalization

Explain that layer normalization normalizes the inputs across the feature dimension for each sample independently, unlike batch normalization which normalizes across the batch.

2. Present the Equation

Write the equation: y = gamma * (x - mu) / sqrt(sigma^2 + epsilon) + beta, where mu and sigma^2 are the mean and variance computed over the features of a single sample.

3. Explain Gamma and Beta

Describe gamma (scale) and beta (shift) as learnable parameters that allow the network to restore representational power after normalization.

4. Explain Epsilon

Clarify that epsilon is a small constant added to the variance to avoid division by zero and improve numerical stability.

5. Discuss Role in Transformers

Highlight that layer normalization is crucial in transformers because it handles variable sequence lengths and small batch sizes effectively, and is typically applied before or after sub-layers (pre-norm vs post-norm).

Key Points to Mention

  • Layer normalization normalizes across features for each sample independently.
  • The equation includes mean and variance computed per sample.
  • Gamma and beta are learnable parameters for scaling and shifting.
  • Epsilon ensures numerical stability by preventing division by zero.
  • It is preferred over batch normalization in transformers due to batch independence.
  • Common placement: pre-norm (before sub-layer) or post-norm (after sub-layer).

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

Q2

Where is LayerNorm applied in a transformer block, and how does placing it before versus after the residual connection affect gradient flow and training stability?

Technical Trade-offsAlgorithms & Data Structures
Author's notes

Pre-norm vs post-norm.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by defining the two main placements: post-layer normalization (original Transformer) and pre-layer normalization (modern variants). Then explain how each affects gradient flow and training stability, referencing the residual connection and normalization order. Conclude with practical implications for training deep models.

Pro tip: Mention that pre-LN is now standard for large models like GPT because it avoids the need for learning rate warm-up and enables stable training without careful initialization. This shows awareness of current industry practices.

1. Define LayerNorm and residual connections

Briefly explain that LayerNorm normalizes activations across features and residual connections add the input to the output of a sublayer, aiding gradient flow.

2. Describe post-layer normalization

In the original Transformer, LayerNorm is applied after the residual connection (i.e., LayerNorm(x + Sublayer(x))). This can cause vanishing gradients in deep networks because the normalization is inside the residual branch.

3. Describe pre-layer normalization

In pre-LN, LayerNorm is applied before the sublayer and residual connection (i.e., x + Sublayer(LayerNorm(x))). This keeps the residual path clean, improving gradient flow and allowing training without warm-up.

4. Compare gradient flow and stability

Post-LN requires careful learning rate warm-up and can be unstable for deep models; pre-LN provides more stable gradients and faster convergence, but may slightly reduce expressivity.

5. Discuss practical implications

Mention that pre-LN is preferred for large-scale models (e.g., GPT) due to stability, while post-LN may still be used in some architectures with proper tuning.

Key Points to Mention

  • LayerNorm normalizes across the feature dimension, not the batch dimension.
  • Residual connections help gradients flow directly to earlier layers.
  • Post-LN: LayerNorm after residual addition; can cause gradient vanishing in deep networks.
  • Pre-LN: LayerNorm before sublayer; keeps residual path unnormalized, improving gradient flow.
  • Pre-LN often eliminates need for learning rate warm-up and allows higher learning rates.
  • Trade-off: Pre-LN may reduce model expressivity slightly compared to Post-LN.

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

Q3

How does LayerNorm differ from BatchNorm, and when would you prefer one over the other?

Technical Trade-offsSystem Design
Author's notes

This part went better.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by clearly defining what each normalization technique normalizes over, then contrast their behavior in terms of batch dependence, training vs. inference, and suitability for different architectures. Finally, discuss practical scenarios (e.g., batch size, sequence models, distributed training) where one is preferred over the other, tying back to Amazon-scale systems.

Pro tip: Emphasize that LayerNorm is often preferred in online serving and small-batch or variable-length sequence scenarios because it doesn't depend on batch statistics, which is critical for low-latency, high-throughput systems like those at Amazon.

1. Define the normalization axis

Explain that BatchNorm normalizes across the batch dimension for each feature, while LayerNorm normalizes across the feature dimension for each sample independently.

2. Discuss training vs. inference behavior

Highlight that BatchNorm uses batch statistics during training and running estimates during inference, which can cause discrepancies, whereas LayerNorm uses the same computation in both phases.

3. Address batch size and sequence length sensitivity

Point out that BatchNorm performance degrades with small batch sizes and variable sequence lengths, while LayerNorm is robust to these variations.

4. Relate to architecture and task

Mention that BatchNorm is common in CNNs for vision tasks, while LayerNorm is standard in Transformers and RNNs for NLP and sequence modeling.

5. Conclude with preference criteria

Summarize when to prefer each: BatchNorm for large-batch, fixed-length inputs (e.g., image classification); LayerNorm for small batches, variable-length sequences, and distributed or online serving.

Key Points to Mention

  • BatchNorm normalizes over the batch dimension; LayerNorm normalizes over the feature dimension per sample.
  • BatchNorm relies on batch statistics during training and running averages during inference, leading to potential train-test mismatch.
  • LayerNorm is independent of batch size and works well with small batches or online learning.
  • LayerNorm is preferred in Transformers and RNNs due to its per-sample normalization and stability with variable-length sequences.
  • BatchNorm is effective in CNNs for vision tasks where large batches are feasible.
  • Distributed training: BatchNorm requires synchronization of statistics across devices, while LayerNorm does not.

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

Q4

What is RMSNorm and how does it compare to LayerNorm in terms of computation and what it gives up?

Technical Trade-offsSystem Design
Author's notes

Blanked for a second.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by defining RMSNorm and contrasting it with LayerNorm in terms of computation, then discuss the trade-offs in normalization behavior and empirical performance. Emphasize the computational efficiency gains and what is sacrificed (mean centering and re-scaling invariance).

Pro tip: Mention that RMSNorm is used in large language models like LLaMA and that its success suggests that re-centering may not be crucial for many tasks, but be prepared to discuss scenarios where LayerNorm might still be preferred.

1. Define RMSNorm

Explain that RMSNorm normalizes the input by its root mean square (RMS) without subtracting the mean, and then scales by a learnable parameter.

2. Compare Computation

Detail the computational steps: LayerNorm computes mean and variance, while RMSNorm only computes RMS. Highlight that RMSNorm avoids mean computation and subtraction, reducing operations.

3. Discuss Trade-offs

Explain that RMSNorm gives up mean centering and the invariance to re-scaling and re-centering that LayerNorm provides. This may affect performance on tasks where such invariances are important.

4. Empirical Performance

Mention that despite the simplification, RMSNorm has been shown to perform comparably or better in some large-scale models, often with faster training and inference.

5. Application Context

Conclude by discussing when to choose RMSNorm (e.g., large transformers for efficiency) versus LayerNorm (e.g., when normalization robustness is critical).

Key Points to Mention

  • RMSNorm normalizes by root mean square, no mean subtraction.
  • LayerNorm computes mean and variance, then normalizes and scales.
  • RMSNorm reduces computation by avoiding mean calculation and subtraction.
  • RMSNorm loses re-centering and re-scaling invariance.
  • RMSNorm is used in models like LLaMA and shows competitive performance.
  • Trade-off: efficiency vs. potential loss of normalization benefits.

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

Q5

What are the implications of LayerNorm placement and design choices on inference latency and memory usage?

System DesignTechnical Trade-offs
Author's notes

Honestly the weakest part of my interview.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by clarifying that LayerNorm placement (pre-norm vs post-norm) and design choices (e.g., fused implementations, normalization axis) directly affect inference latency and memory through kernel launches, memory access patterns, and parallelization. Then, systematically compare trade-offs in latency and memory, using concrete examples from transformer architectures, and conclude with practical recommendations for optimizing inference in production systems.

Pro tip: Emphasize that in inference, LayerNorm is often memory-bound rather than compute-bound, so fusing it with adjacent operations (e.g., residual add) can yield significant latency improvements. Also, mention that pre-norm architectures typically have lower memory overhead during inference due to reduced activation storage, but may require more normalization layers.

1. Define LayerNorm and its role

Briefly explain what LayerNorm does and why its placement matters in neural networks, especially transformers.

2. Compare pre-norm vs post-norm

Discuss how pre-norm (LayerNorm before sublayer) and post-norm (after sublayer) affect gradient flow, training stability, and inference latency/memory.

3. Analyze design choices

Cover implementation details like fused kernels, normalization axis (e.g., feature vs. token), and precision (FP16/FP32) and their impact on latency and memory.

4. Quantify trade-offs

Provide concrete examples or estimates of latency and memory differences, referencing common architectures (e.g., BERT, GPT).

5. Summarize recommendations

Offer practical guidance on choosing LayerNorm placement and design for inference-optimized systems, considering hardware and deployment constraints.

Key Points to Mention

  • Pre-norm vs post-norm: pre-norm reduces activation memory and enables parallelization but may require more normalization layers; post-norm can be more stable but increases memory due to storing intermediate activations.
  • Fused LayerNorm implementations (e.g., in CUDA) reduce kernel launches and memory bandwidth, lowering latency.
  • Normalization axis: normalizing over the feature dimension (common in NLP) vs. spatial dimensions (vision) affects memory access patterns and cache efficiency.
  • Precision: using FP16/BF16 for LayerNorm can reduce memory footprint and speed up computation, but may introduce numerical instability.
  • Batch size and sequence length: larger batches/sequences amplify memory usage, making LayerNorm placement more critical.
  • Hardware considerations: on GPUs, memory-bound operations benefit from fusion; on CPUs, vectorization and threading impact performance.

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