Katkat
Back

Derive MLP shapes and explain PyTorch broadcasting

NVIDIASoftware Engineer · Technical Screen
MediumSoftware EngineerTechnical ScreenMachine Learning

Last updated

This question tests the ability to calculate matrix transformation dimensions in a Multilayer Perceptron (MLP) and explains the rules governing PyTorch broadcasting when performing operations between tensors of different shapes.

Part 1: MLP Dimensions

You are designing a simple MLP for image classification. Your input is a batch of 32 images, each flattened into a vector of size 784. The architecture is defined as follows:

  • Layer 1 (Linear): Maps 784 input features to 128 hidden units.
  • Activation: ReLU.
  • Layer 2 (Linear): Maps 128 hidden units to 10 output classes.

Calculate the required shape of the weight matrices (W1, W2) and bias vectors (b1, b2) for this network. Assume PyTorch-style linear layers where the operation is x @ W.T + b.

Part 2: Broadcasting Logic

Consider the following PyTorch code snippet:

import torch
a = torch.randn(5, 1, 4)
b = torch.randn(3, 4)
result = a + b
  1. Explain why this operation is valid.
  2. What is the final shape of result?
  3. Describe the underlying mechanism PyTorch uses to align these tensors.

Constraints

  • Assume tensors follow the convention: (batch_size, feature_dim).
  • Weights are typically stored as (out_features, in_features) in PyTorch Linear layers.

Hints

Hint 1

For part 1, recall that the weight matrix transpose operation `W.T` implies the stored weights must have the shape `(out_features, in_features)`.

Hint 2

For part 2, remember that broadcasting compares dimensions starting from the trailing (rightmost) dimension and moving left.

Hint 3

Broadcasting rules state that two dimensions are compatible if they are equal, or if one of them is 1.

Follow-up questions

  • How does broadcasting change if we introduce a dimension of size 0 versus a dimension of size 1?
  • If we add a bias vector of shape (10,) to an output matrix of shape (32, 10), how does PyTorch broadcasting handle the alignment?
  • Why might you prefer explicit dimension expansion (e.g., `unsqueeze` or `view`) over relying on implicit broadcasting in production code?