Machine Learning System Design Interview #22 - The Softmax Trap
Why writing “clean code” for softmax regression silently destroys your training run, and how LogSumExp saves you.
You’re in a Senior Machine Learning interview at OpenAI. The interviewer hands you a marker and asks for a scratch implementation of Softmax Regression.
It feels like a “FizzBuzz” question. It isn’t. It’s a trap.
90% of candidates walk right into it by writing “clean,” modular code.
They write the forward pass like this:
probs = softmax(logits)
loss = -log(probs)
It makes sense logically. You compute the probabilities (the hypothesis), and then you compute the error (the loss). It separates concerns. It looks like “good” software engineering.
-----
𝐓𝐡𝐞 𝐓𝐮𝐫𝐧:
