← Amazon Interview Insights

Amazon·Data Scientist·Technical Phone Screen·Senior

SeniorPrefer not to say
May 2026Remote

Summary

Brutal technical screen for a Data Scientist role at Amazon. Eight questions covering everything from gradient descent derivations to LoRA fine-tuning on a single GPU. Whoever wrote this question set really wanted to see if you could go deep on ML fundamentals, RL, transformers, and LLM infra all in one sitting.

Questions Asked (8)

Q1

Derive the update rules for full-batch gradient descent and SGD for a mean loss over n samples. Compare their convergence behavior, gradient variance, and wall-clock efficiency. When does SGD actually win?

Technical Trade-offsAlgorithms & Data Structures
Author's notes

This is the kind of question that sounds easy until you're actually writing out the math live.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by mathematically deriving the update rules for full-batch gradient descent and SGD for a mean loss over n samples. Then compare their convergence behavior, gradient variance, and wall-clock efficiency, emphasizing the trade-offs. Conclude with practical scenarios where SGD outperforms full-batch, especially in large-scale machine learning.

Pro tip: Highlight that SGD's advantage is not just theoretical but depends on hardware, data size, and problem structure; mentioning that SGD often wins in early training and when data is redundant shows practical insight.

1. Derive update rules

Write the mean loss L(θ) = (1/n) Σ_{i=1}^n ℓ(θ; x_i, y_i). For full-batch GD, the update is θ_{t+1} = θ_t - η ∇L(θ_t) = θ_t - (η/n) Σ_{i=1}^n ∇ℓ_i(θ_t). For SGD, the update is θ_{t+1} = θ_t - η ∇ℓ_i(θ_t) where i is randomly sampled from {1,...,n}.

2. Compare convergence behavior

Full-batch GD has deterministic, smooth convergence but can be slow per iteration. SGD has noisy updates that can escape local minima and often converges faster initially, but may oscillate near the optimum; with decaying learning rate, it converges almost surely under convexity.

3. Analyze gradient variance

Full-batch gradient has zero variance (deterministic). SGD gradient is an unbiased estimator of the full gradient but has variance that depends on the sampling; variance decreases with batch size and can be reduced by mini-batching or variance reduction techniques.

4. Evaluate wall-clock efficiency

Full-batch requires processing all n samples per update, which is expensive for large n. SGD processes one sample per update, so each update is cheap, but many updates may be needed. Wall-clock time depends on hardware parallelism and data access patterns; SGD often wins when n is large and data is easily streamed.

5. Determine when SGD wins

SGD wins when n is very large, data is redundant, or when a quick approximate solution is sufficient. It also excels in online learning and when memory constraints prevent full-batch. However, for small n or when high precision is needed, full-batch or mini-batch may be better.

Key Points to Mention

  • Mathematical derivation of update rules for full-batch GD and SGD
  • Unbiasedness of SGD gradient and its variance
  • Convergence rates: full-batch GD has linear convergence for strongly convex; SGD has sublinear convergence with decaying learning rate
  • Wall-clock efficiency: per-iteration cost vs. number of iterations
  • Impact of batch size and variance reduction techniques
  • Practical scenarios: large-scale data, online learning, early stopping

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

Q2

Define batch size. Given 50,000 samples, 5 epochs, and batch size 200, how many update steps per epoch and in total? If you increase batch size to 2,000, what changes and how should you adjust the learning rate? When does the linear scaling rule break down?

Technical Trade-offsAlgorithms & Data Structures
Author's notes

The arithmetic is trivial but they clearly wanted the learning rate discussion.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by defining batch size clearly, then compute the update steps for the given scenario. Explain the effect of increasing batch size on steps and the need to adjust learning rate, referencing the linear scaling rule and its limitations.

Pro tip: Mention that the linear scaling rule is a heuristic and that you should monitor validation performance and consider warmup or layer-wise adaptation, especially for large batches.

1. Define batch size

Explain that batch size is the number of training examples used in one forward/backward pass to compute gradients and update model parameters.

2. Compute update steps

Calculate steps per epoch as total samples divided by batch size (50,000 / 200 = 250). Total steps = steps per epoch * epochs (250 * 5 = 1,250).

3. Analyze batch size increase

With batch size 2,000, steps per epoch = 50,000 / 2,000 = 25, total steps = 25 * 5 = 125. Fewer updates per epoch.

4. Adjust learning rate

Apply linear scaling rule: multiply learning rate by the factor of batch size increase (10x). So if original LR was η, new LR = 10η.

5. Discuss limitations

Explain that linear scaling breaks down for very large batches due to optimization difficulties, diminishing returns, and generalization gap; may need warmup, LARS/LAMB, or smaller scaling factor.

Key Points to Mention

  • Definition of batch size and its role in SGD
  • Calculation of steps per epoch and total steps
  • Linear scaling rule: LR ∝ batch size
  • Breakdown of linear scaling for large batches
  • Alternatives: warmup, layer-wise adaptive rates, gradient accumulation
  • Impact on convergence and generalization

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

Q3

For each of the following, say whether it's supervised or unsupervised and give one real use-case: logistic regression, SVM, k-NN, k-means, PCA, t-SNE, Isolation Forest.

Technical Trade-offs
Author's notes

Straightforward.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Quickly classify each algorithm as supervised or unsupervised, then provide a concise real-world use-case that highlights its practical value. Group similar algorithms (e.g., supervised classifiers, unsupervised dimensionality reduction) to show structured thinking and avoid repetition.

Pro tip: For Amazon, emphasize use-cases that tie to business impact, such as fraud detection, customer segmentation, or product recommendations, and mention scalability considerations for large datasets.

1. Clarify definitions

Briefly state the difference between supervised (labeled data) and unsupervised (unlabeled data) learning to set context.

2. Classify each algorithm

Go through the list and assign each to supervised or unsupervised, noting any that can be used in both settings (e.g., SVM, k-NN).

3. Provide a real use-case

For each algorithm, give one concrete, industry-relevant example that demonstrates its application.

4. Highlight trade-offs

Optionally mention key trade-offs (e.g., interpretability, scalability, need for labeled data) to show deeper understanding.

Key Points to Mention

  • Logistic regression: supervised, used for binary classification like spam detection or click-through rate prediction.
  • SVM: supervised, effective for high-dimensional data such as text classification or image recognition.
  • k-NN: supervised (can be unsupervised for clustering), used in recommendation systems or anomaly detection.
  • k-means: unsupervised, used for customer segmentation or document clustering.
  • PCA: unsupervised, used for dimensionality reduction and feature extraction in image or gene expression data.
  • t-SNE: unsupervised, used for visualization of high-dimensional data like word embeddings or single-cell RNA-seq.
  • Isolation Forest: unsupervised, used for anomaly detection in fraud detection or network intrusion.

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

Q4

How does reinforcement learning relate to supervised and unsupervised learning? Write out the REINFORCE policy gradient estimator and show how a baseline keeps the estimator unbiased while reducing variance. Then compute the gradient for a 3-step trajectory with returns [3, 1, -1] and a constant baseline equal to the mean of those returns.

Technical Trade-offsAlgorithms & Data Structures
Author's notes

The variance reduction proof is something I'd reviewed but the concrete numerical example caught me a little flat-footed.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by contrasting RL with supervised and unsupervised learning in terms of feedback, data, and objective. Then derive the REINFORCE estimator, explain the role of a baseline, and finally compute the gradient for the given trajectory with the specified baseline.

Pro tip: Emphasize that the baseline must not depend on the action to preserve unbiasedness, and show the variance reduction by comparing the magnitude of the gradient terms with and without the baseline.

1. Compare learning paradigms

Contrast RL with supervised learning (explicit labels, i.i.d. data) and unsupervised learning (no labels, structure discovery), highlighting RL's sequential decision-making, delayed rewards, and exploration-exploitation trade-off.

2. Write REINFORCE estimator

State the policy gradient theorem and the REINFORCE estimator: ∇J(θ) = E[∑ ∇log π(a|s) * G_t], where G_t is the return from time t.

3. Explain baseline and unbiasedness

Introduce a baseline b(s) that does not depend on the action, show that E[∇log π(a|s) * b(s)] = 0, hence the estimator remains unbiased, and explain how it reduces variance by centering returns.

4. Compute gradient for given trajectory

For the 3-step trajectory with returns [3, 1, -1] and constant baseline b = mean = 1, compute the adjusted returns [2, 0, -2] and the gradient as ∑ ∇log π(a_t|s_t) * (G_t - b).

Key Points to Mention

  • RL learns from scalar rewards via trial-and-error, unlike supervised learning which uses labeled examples.
  • Unsupervised learning finds patterns without labels, while RL optimizes long-term reward through sequential actions.
  • REINFORCE is a Monte Carlo policy gradient method that uses complete episode returns.
  • A baseline reduces variance without introducing bias as long as it does not depend on the action.
  • The gradient for the given trajectory is ∇log π(a1|s1)*2 + ∇log π(a2|s2)*0 + ∇log π(a3|s3)*(-2).
  • Using the mean return as baseline centers the returns, often leading to faster and more stable learning.

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

Q5

How do neural networks fit into RL? Specifically for DQN, why do target networks and experience replay stabilize training, and what goes wrong if you remove them?

Technical Trade-offsSystem Design
Author's notes

Felt pretty comfortable here.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by framing neural networks as function approximators that enable RL to handle large state spaces, then focus on DQN as a specific algorithm. Explain the two key innovations—target networks and experience replay—and how each addresses a distinct instability: moving targets and correlated samples. Finally, describe the failure modes when they are removed, tying back to the bias-variance trade-off and convergence guarantees.

Pro tip: Emphasize that these techniques are not just engineering tricks but address fundamental issues in combining bootstrapping, off-policy learning, and function approximation (the deadly triad). Mentioning this shows deep understanding beyond surface-level definitions.

1. Set the context: Neural networks in RL

Briefly explain that neural networks serve as scalable function approximators for value functions or policies, enabling RL in high-dimensional state spaces like images. Mention that DQN is a value-based method that uses a neural network to approximate the Q-function.

2. Explain target networks

Describe how target networks provide a stable target for the Q-learning update by freezing the parameters of the target network for a fixed number of steps. This reduces the moving target problem where the network is chasing its own updates.

3. Explain experience replay

Explain that experience replay stores transitions in a buffer and samples mini-batches randomly, breaking temporal correlations and improving sample efficiency. This also allows the agent to learn from rare events multiple times.

4. Describe what goes wrong without them

Without target networks, the Q-values can oscillate or diverge due to the moving target. Without experience replay, the network overfits to recent experiences, leading to catastrophic forgetting and high variance in updates.

5. Connect to broader RL challenges

Summarize that these techniques mitigate the deadly triad of function approximation, bootstrapping, and off-policy learning. Mention that they are not perfect but are crucial for stable training in practice.

Key Points to Mention

  • Neural networks as function approximators for Q-values in DQN
  • Target network: periodic hard updates or soft updates (Polyak averaging) to stabilize targets
  • Experience replay: decorrelates samples, improves data efficiency, and reduces variance
  • Moving target problem: without target network, updates chase a non-stationary target
  • Correlated data problem: without replay, sequential updates lead to biased gradients and instability
  • Deadly triad: function approximation + bootstrapping + off-policy learning can diverge

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

Q6

Compare Transformers and RNNs on parallelism, long-range dependencies, and computational complexity. For sequence length 1024 and model dimension 512, what are the asymptotic time and memory costs of self-attention? Name two approaches to reduce the quadratic scaling.

System DesignTechnical Trade-offs
Author's notes

O(n^2 * d) time and O(n^2) memory for attention, that part is fine.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by contrasting Transformers and RNNs on the three axes: parallelism, long-range dependencies, and computational complexity. Then derive the asymptotic time and memory costs of self-attention for the given dimensions, and finally name two methods to reduce quadratic scaling. Keep the answer structured and quantitative.

Pro tip: Mention that while self-attention is O(n^2), it enables parallel training and better long-range modeling, which often outweighs the cost for moderate sequence lengths. Also, relate the quadratic scaling to practical limits like GPU memory.

1. Compare parallelism

Explain that Transformers process all positions in parallel, while RNNs are inherently sequential, limiting parallelization across time steps.

2. Compare long-range dependencies

Highlight that Transformers capture long-range dependencies directly via self-attention, whereas RNNs struggle due to vanishing/exploding gradients and sequential propagation.

3. Compare computational complexity

State that self-attention has O(n^2 * d) time and O(n^2) memory for sequence length n and dimension d, while RNNs are O(n * d^2) time and O(d) memory per step (or O(n*d) for all states).

4. Compute costs for given dimensions

For n=1024 and d=512, self-attention time is O(1024^2 * 512) ≈ 5.4e8 operations, and memory is O(1024^2) ≈ 1e6 attention scores.

5. Name reduction approaches

Mention two methods: sparse attention (e.g., Longformer, BigBird) and low-rank approximations (e.g., Linformer, Performer).

Key Points to Mention

  • Transformers enable parallel processing across sequence positions, unlike RNNs.
  • Self-attention directly models long-range dependencies, avoiding sequential bottlenecks.
  • Self-attention time complexity: O(n^2 * d); memory: O(n^2).
  • RNN time complexity: O(n * d^2); memory: O(d) per step (or O(n*d) for all hidden states).
  • For n=1024, d=512: self-attention requires ~1M attention scores (memory) and ~0.5B operations (time).
  • Sparse attention and low-rank approximations reduce quadratic scaling.

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

Q7

What are embeddings and what is polysemy? Propose a method to distinguish the word 'King' in a chess context versus a monarchy context. How would you evaluate it both intrinsically and extrinsically?

Technical Trade-offsProduct Analytics & Metrics
Author's notes

Contextual encoders like BERT handle polysemy naturally since the embedding changes per context.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by clearly defining embeddings and polysemy, then propose a contextual embedding method (e.g., fine-tuned BERT) to distinguish 'King' in chess vs. monarchy. Finally, outline intrinsic evaluation (e.g., word sense disambiguation accuracy) and extrinsic evaluation (e.g., downstream task performance like chess move prediction or monarchy-related QA).

Pro tip: Emphasize that the choice of embedding method should align with business goals—at Amazon, extrinsic metrics tied to customer experience often matter more than intrinsic benchmarks.

1. Define embeddings and polysemy

Explain that embeddings are dense vector representations of words capturing semantic and syntactic properties, and polysemy refers to a word having multiple meanings depending on context.

2. Propose a method for context disambiguation

Suggest using contextual embeddings like BERT or ELMo, which generate different vectors for 'King' based on surrounding words, or fine-tune a model on domain-specific data to create separate sense embeddings.

3. Design intrinsic evaluation

Evaluate the method intrinsically by measuring how well it distinguishes senses, e.g., through word sense disambiguation accuracy on a labeled dataset or by clustering embeddings of 'King' in different contexts.

4. Design extrinsic evaluation

Evaluate extrinsically by applying the embeddings to downstream tasks: for chess, predict legal moves or game outcomes; for monarchy, answer questions or classify texts about royal families, and compare performance against a baseline.

5. Discuss trade-offs and business impact

Highlight trade-offs between model complexity, latency, and accuracy, and tie evaluation metrics to business outcomes like improved search relevance or recommendation quality.

Key Points to Mention

  • Embeddings: dense vectors capturing semantic similarity; examples: Word2Vec, GloVe, BERT.
  • Polysemy: one word, multiple meanings; 'King' in chess vs. monarchy.
  • Contextual embeddings (e.g., BERT) dynamically adjust representations based on context.
  • Intrinsic evaluation: word sense disambiguation accuracy, similarity judgments, clustering.
  • Extrinsic evaluation: downstream task performance (e.g., chess move prediction, monarchy QA).
  • Trade-offs: computational cost vs. accuracy; need for labeled data; alignment with business metrics.

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

Q8

You have one 24GB GPU and a 7B parameter model. Design a fine-tuning plan that fits within those constraints. Walk through your choices for quantization, adapter method, optimizer, and learning rate schedule. Estimate how many trainable parameters you'd actually have with LoRA, assuming hidden size ~4096 and ~32 layers.

System DesignTechnical Trade-offs
Author's notes

This one was the most fun.

Create a free account to read the full note

AI HintsAI Generated

Suggested Approach

Start by framing the memory budget: 24GB must hold model weights, activations, gradients, and optimizer states. Propose QLoRA (4-bit NF4 quantization) with LoRA adapters, gradient checkpointing, and 8-bit AdamW to fit. Then estimate LoRA parameters and justify hyperparameters like rank, learning rate, and schedule.

Pro tip: Mention that you would first run a quick memory profiling with a small batch to verify headroom, and that you would use gradient accumulation to simulate larger batches without increasing memory.

1. Assess memory constraints

Break down the 24GB budget: model weights (7B params), activations, gradients, and optimizer states. Explain that full fine-tuning is infeasible due to optimizer states (e.g., AdamW needs 2x model size).

2. Choose quantization and adapter method

Propose 4-bit quantization (NF4) to reduce base model memory to ~3.5GB, and LoRA adapters to train only a small fraction of parameters. Mention that QLoRA combines both and is ideal for this scenario.

3. Estimate trainable parameters with LoRA

Calculate LoRA parameters: for each target module (e.g., q, v), add two low-rank matrices of size (hidden_size x r) and (r x hidden_size). With hidden size 4096, 32 layers, and rank r=8, parameters per module = 2 * 4096 * 8 = 65,536. If targeting 2 modules per layer, total = 32 * 2 * 65,536 = 4,194,304 (~4.2M).

4. Select optimizer and learning rate schedule

Use 8-bit AdamW to save memory, with a learning rate around 1e-4 to 3e-4 (higher than full fine-tuning). Employ a cosine schedule with warmup (e.g., 3-5% of steps) to stabilize training.

5. Address remaining memory and training efficiency

Enable gradient checkpointing to reduce activation memory, use mixed precision (bf16) for forward/backward, and consider gradient accumulation to increase effective batch size. Monitor GPU memory and adjust batch size accordingly.

Key Points to Mention

  • QLoRA: 4-bit NormalFloat quantization + LoRA adapters
  • Memory savings from quantization (7B model ~3.5GB in 4-bit)
  • LoRA parameter calculation: rank, target modules, and total trainable parameters (~4M for r=8, q&v)
  • Optimizer: 8-bit AdamW to reduce optimizer state memory
  • Learning rate: higher LR (1e-4 to 3e-4) with cosine schedule and warmup
  • Gradient checkpointing and mixed precision to fit within 24GB

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