Problems
Loading...

04/08Scaled Dot-Product Attention

Medium

Implement scaled dot-product attention from the Transformer.

Attention⁡(Q,K,V)=softmax⁡ ⁣(QKTdk)V\operatorname{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V

Here, Q contains queries, K contains keys, V contains values, and (d_k) is the key width. Apply softmax along the key-position axis. Return a PyTorch float32 tensor with shape (batch, query_length, value_width).

Loading visualization...

Examples

Input: Q = [[[1.0, 0.0], [0.0, 1.0]]], K = [[[1.0, 0.0], [0.0, 1.0]]], V = [[[1.0, 2.0], [3.0, 4.0]]]

Output: [[[1.6605, 2.6605], [2.3395, 3.3395]]]

Explanation: The first query [1,0] aligns more with key [1,0], giving it weight ~0.66, so the output leans toward V[0]=[1,2]. The second query [0,1] aligns more with key [0,1], leaning toward V[1]=[3,4].

Input: Q = [[[1.0, 0.0, 0.0]]], K = [[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]], V = [[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]

Output: [[[0.4711, 0.2645, 0.2645]]]

Input: Q = [[[1.0, 1.0], [1.0, 1.0]]], K = [[[1.0, 0.0], [0.0, 1.0]]], V = [[[2.0, 0.0], [0.0, 2.0]]]

Output: [[[1.0, 1.0], [1.0, 1.0]]]

Hint 1

Transpose only the final two dimensions of K.

Hint 2

Use K.shape[-1] for the scale.

Hint 3

Multiply the attention weights by V.

Requirements

  • Use PyTorch tensor operations.
  • Scale scores by the square root of K's final dimension.
  • Apply softmax over the final score dimension.
  • Return the weighted sum of V without rounding.

Constraints

  • Q, K, and V have dtype torch.float32.
  • Q and K share their final dimension.
  • K and V share their sequence length.
  • Valid inputs are guaranteed.
Prev4 / 8Next