Maths for ML interview question · Linear algebra: matrices as transformations · mid level
A batch of 32 MNIST images goes into a dense layer with 128 units. What are the shapes, and what is the rule for when two matrices can be multiplied?
Free to read in full. This is 1 of the 120 questions in the Maths for ML and Deep Learning interview pack, with 147 diagrams. Get the pack, ₹999 · $20
Why it is asked
Shape errors are the most common bug in model code. The interviewer checks the inner-dimension rule and how frameworks store weights.
The answer
The input is 32 × 784, the weight matrix is 784 × 128, and the output is 32 × 128. The rule: an (m × k) matrix times a (k × n) matrix gives an (m × n) matrix, and the two inner sizes (k and k) must match.
Where the numbers come from:
- MNIST is a famous dataset of 28 × 28 pictures of handwritten digits 0-9. A dense layer with 128 units is a layer of 128 neurons, each connected to every input pixel.
- Each image is 28 × 28 = 784 gray pixels, flattened into one row of 784 features.
- 32 images stacked as rows give X of shape (32 × 784): batch × features.
- W has one column per unit: (784 × 128).
- Y = XW is (32 × 128): one row of 128 outputs per image. The bias, 128 numbers, is added to every row (broadcasting).
Why the inner sizes must match: entry (i, j) of the result is the dot product of row i of X and column j of W, and a dot product needs two vectors of the same length, here 784.
Two practical details:
- PyTorch's nn.Linear(784, 128) stores its weight as (128 × 784), that is (out × in), and computes Y = XWᵀ + b. Same math, transposed storage.
- With extra leading dimensions, such as (batch, tokens, features), where tokens are the word pieces of a text, the @ operator multiplies the last two dimensions and treats the rest as a batch: (8 × 50 × 768) @ (768 × 3072) gives (8 × 50 × 3072).
What the interviewer listens for
- (m × k)(k × n) = (m × n); inner sizes must match.
- Worked: (32 × 784)(784 × 128) = (32 × 128).
- Bias is broadcast across the batch rows.
- nn.Linear stores W as (out × in) and uses XWᵀ.
Follow-ups the interviewer may ask
- How many parameters does this layer have? Hint: 784 × 128 + 128 = 100,480.
- What shape does a (B × n × d) tensor have after @ (d × d_ff)? Hint: (B × n × d_ff).
- Why put the batch first? Hint: one large matrix multiply per batch keeps the GPU busy.
Common mistakes
- Writing the weight as (128 × 784) and forgetting the transpose in a hand-written layer.
- Flattening images to (784 × 32) and mixing up batch and feature axes.
Written by Roni Das, who builds distributed systems and AI systems for a living.
One question is not enough for the interview.
This is 1 of the 120 questions in the Maths for ML and Deep Learning pack
Each question has why it is asked, a full answer, what the interviewer listens for, the follow-ups, the common mistakes and a diagram (147 diagrams in all). A PDF you keep, junior to staff level.
India: UPI, cards or netbanking through Razorpay. Elsewhere: PayPal. Download at once from My packs.
More free questions from this pack
- Compute the cosine similarity of a = (3, 4) and b = (4, 3) by hand. What does the result mean? (junior)
- A disease affects 1% of people. A test catches 90% of sick people (sensitivity) and wrongly flags 9% of healthy people (false positive rate). You test positive. What is the chance you are sick? (mid)
- Why do we not use mean squared error for classification with a sigmoid output? Compare the gradients numerically. (staff)
- Compute scaled dot-product attention by hand for two tokens with d_k = 2. (senior)
Free questions from the other packs
- What do temperature and top-p control, and what values would you use for data extraction versus brainstorming? (AI Engineering)
- What is hybrid search, and how do you merge keyword and vector results with Reciprocal Rank Fusion? (AI Engineering)
- Users ask "list all X" questions and RAG returns only five items. Why, and what do you do? (AI Engineering)
- How do you check that an LLM judge agrees with human reviewers? (AI Engineering)
- Cut the cost of an LLM feature serving 10 million requests a day by 60% without hurting quality. Walk through your plan. (AI Engineering)
- Why are LLMs bad at counting letters or reversing a word? (GenAI and LLM)
- Write out scaled dot-product attention with the shape of every tensor. (GenAI and LLM)
- How does speculative sampling keep exactly the target model's distribution? (GenAI and LLM)
- What is GRPO, and how do verifiable rewards train reasoning models? (GenAI and LLM)
- Design a system that answers questions from 300-page contracts accurately. (GenAI and LLM)
- What happens, step by step, when an LLM server handles one request? (Inference Engineering)
- Derive the KV cache size formula and work it out for Llama 3 70B. (Inference Engineering)
- What is PagedAttention? (Inference Engineering)
- When does speculative decoding stop helping or make things worse? (Inference Engineering)
- Design a service to serve Llama 3 70B at 10 QPS with latency targets. (Inference Engineering)
- A customer says: "We want AI for our support team." What do you do in the first meeting? (Forward Deployed Engineer)
- The customer's admin offers you a global admin account "to save time". What do you ask for instead, and why? (Forward Deployed Engineer)
- Your connector writes tickets into the customer's system. A request timed out, the connector retried, and now there are duplicate tickets. How do you make the writes safe to retry? (Forward Deployed Engineer)
- The customer wants RAG over 10 million documents, all behind their firewall. Size it for me on the whiteboard. (Forward Deployed Engineer)
- One deployment of your AI platform must serve 50 business units of the same enterprise. How do you design multi-tenancy? (Forward Deployed Engineer)