Implement Masked Cross-Entropy with Label Smoothing
Company: OpenAI
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
# Implement Masked Cross-Entropy with Label Smoothing
You are given a NumPy array of unnormalized class logits with shape (B, C), integer labels with shape (B,), and an optional Boolean mask with shape (B,). Implement numerically stable cross-entropy, then extend it with label smoothing.
Use epsilon in [0, 1). For each included example, assign target probability 1 - epsilon to the labeled class and distribute epsilon uniformly across the other C - 1 classes. Return the mean loss over examples whose mask is true. If no mask is supplied, include every example. Raise ValueError when the mask excludes every example.
For this practice version, the function also accepts an optional temperature. Temperature defaults to 1.0, must be finite and strictly positive, and scales logits as logits / temperature before log-softmax. Temperature was not a reported interview requirement; it is included only to make numerical behavior explicit.
Implement:
def masked_smoothed_cross_entropy(
logits,
labels,
mask=None,
epsilon=0.0,
temperature=1.0,
) -> float:
...
### Constraints & Assumptions
- B >= 1 and C >= 2.
- Every label is an integer in [0, C).
- Logits must be numeric and finite after conversion to NumPy float64. NaN and positive or negative infinity are rejected with ValueError.
- Computation is intentionally performed in float64 and the function returns a Python float. Inputs not exactly representable in float64 may round during conversion.
- If temperature scaling, max-shift differences, or the final float64 loss overflow, raise ValueError rather than returning NaN or infinity.
- A supplied mask must have Boolean dtype and exactly B entries.
- Do not call a library cross-entropy implementation.
### Clarifying Questions to Ask
- Should smoothing mass include or exclude the labeled class?
- Is reduction a sum, a batch mean, or a mean over unmasked examples?
- What should an all-false mask do?
- Are gradients required, or only the forward loss?
- Is temperature part of the required API, and what precision should calculations use?
### Hints
- Validate numeric inputs before relying on the max-shift stability argument.
- Apply temperature before subtracting the row maximum.
- Separate log-probability calculation, target construction, masking, and reduction.
- Use tiny hand-computable batches for unsmoothed and smoothed cases.
### What a Strong Answer Covers
- Correct shape, dtype, finiteness, epsilon, temperature, label, and mask validation.
- Stable log-softmax for finite scaled logits.
- Precise smoothing, masking, empty-selection, precision, and return-type semantics.
- A derivation connecting logits, log probabilities, target distributions, and reduction.
- Time and space complexity plus tests for extreme finite logits and invalid non-finite inputs.
- How cross-entropy relates to entropy and KL divergence.
### Follow-up Questions
1. Derive the gradient with respect to the original, unscaled logits.
2. How would the implementation change for token logits shaped (batch, sequence, classes)?
3. When can label smoothing improve generalization, and when can it hurt calibration?
4. How would you verify a hand-written gradient with finite differences?
Quick Answer: Implement numerically stable masked cross-entropy in NumPy with label smoothing and optional temperature scaling. Validate shapes, dtypes, finite logits, masks, and boundary cases while explaining the loss derivation and gradient.