Implement Stable Sigmoid, Softmax, and Scaled Dot-Product Attention
Company: Amazon
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
# Implement Stable Sigmoid, Softmax, and Scaled Dot-Product Attention
Implement the mathematical core of three common neural-network operations without calling library versions of those operations. Plain loops, elementary arithmetic, and `exp` are available. Explain both correctness and numerical stability.
### Clarifying Questions to Ask
- What tensor shapes and numeric type should the implementation accept?
- Should softmax operate on one vector or a chosen tensor axis?
- Is masking required for attention, and what should an entirely masked row produce?
- Is the task a readable reference implementation or a production-optimized kernel?
### Part 1: Sigmoid
Implement scalar sigmoid and explain how to avoid overflow for large positive or negative inputs.
#### What This Part Should Cover
- The definition `1 / (1 + exp(-x))`
- A sign-aware stable form
- Expected behavior at extreme magnitudes
### Part 2: Softmax
Implement softmax for a non-empty vector of logits. The output must be non-negative and sum to one within floating-point tolerance.
#### What This Part Should Cover
- Subtracting the maximum logit before exponentiation
- A single normalization denominator
- Invariance to adding the same constant to every logit
### Part 3: Attention
For this practice version, implement single-head scaled dot-product attention. Inputs are `Q` with shape `[Lq, d]`, `K` with shape `[Lk, d]`, and `V` with shape `[Lk, dv]`. Compute `softmax(QK^T / sqrt(d))V`. An optional Boolean mask has shape `[Lq, Lk]`; `true` means the key is allowed, and every query row is guaranteed to allow at least one key.
#### What This Part Should Cover
- Shape validation and the scale factor
- Row-wise stable softmax over keys
- Correct mask behavior without allowing masked positions to affect the output
- Time and memory cost in terms of `Lq`, `Lk`, `d`, and `dv`
### What a Strong Answer Covers
- Stable formulas rather than mathematically correct but overflow-prone code
- Clear shape and masking semantics
- Tests for extreme logits, normalization, and simple attention cases
- A distinction between a clear reference implementation and an optimized production kernel
### Follow-up Questions
- Why does subtracting the maximum not change softmax?
- Why is the dot product divided by `sqrt(d)`?
- How would causal masking change the mask construction?
- What memory bottleneck appears for long sequences, and what can an optimized kernel avoid materializing?
Quick Answer: Implement numerically stable sigmoid, softmax, and single-head scaled dot-product attention using elementary arithmetic and loops. The discussion covers extreme inputs, tensor shapes, row-wise normalization, optional masks, scaling, complexity, tests, and the distinction between reference code and optimized kernels.
