Dataford
Interview QuestionsInterview GuidesExperiencesMock InterviewsPricing
Get started

Self-Attention From Scratch

HardPython00:00
Practice interviewer
In session
5 left
00:00

Your question is Self-Attention From Scratch. Start with the requirements on the right.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

You need to log in / sign up to run or submit.

Problem

AppFolio Property Manager may use attention-based models to represent sequences such as maintenance updates or resident communications. Implement a single self-attention layer from scratch using only Python's standard library.

Given token embeddings and projection matrices, compute queries, keys, and values, then return scaled dot-product attention outputs. Do not use NumPy, PyTorch, TensorFlow, or other numerical libraries.

Formal Specification

Implement self_attention(tokens, w_q, w_k, w_v, causal), where tokens is an n x d_model matrix and each projection matrix is d_model x d_head. Compute Q = tokens @ w_q, K = tokens @ w_k, and V = tokens @ w_v. For every token pair (i, j), compute:

score(i, j) = dot(Q[i], K[j]) / sqrt(d_head)

Apply a row-wise softmax to the scores. When causal is true, position i may attend only to positions j <= i. Return the resulting n x d_head matrix, where each output row is the weighted sum of value rows.

The softmax must be numerically stable by subtracting the row maximum before exponentiation.

Constraints

  • 1 <= len(tokens) <= 256
  • 1 <= len(tokens[0]) <= 64
  • 1 <= len(w_q[0]) = len(w_k[0]) = len(w_v[0]) <= 64
  • tokens has shape n x d_model
  • Each projection matrix has shape d_model x d_head
  • All input values are finite real numbers

Function Signature

def self_attention(tokens, w_q, w_k, w_v, causal):
Your solutionPython 3
You need to log in / sign up to run or submit.
Run your code to see test output