Why total compute is invariant to h

Background

A surprising fact: adding heads doesn't increase the attention compute. Each head scores an $n\times n$ matrix over its $d_k = d_{\text{model}}/h$ dimensions, costing $n^2 d_k$. Summing over $h$ heads:

$$\text{total} = h \cdot n^2 d_k = h \cdot n^2 \frac{d_{\text{model}}}{h} = n^2 d_{\text{model}}$$

— the same as a single full-width head.

Problem statement

Implement multihead_compute(n, d_model, h) returning (per_head_ops, total_ops) for the score computation.

Input

• n — sequence length.
• d_model — model width.
• h — number of heads (divides d_model).

Output

Returns (per_head_ops, total_ops):

• per_head_ops = n * n * (d_model // h),
• total_ops = h * per_head_ops.

Examples

Example 1 — 8 heads

Input:  n = 10, d_model = 512, h = 8
Output: (6400, 51200)

Explanation: d_k = 64; per head 10*10*64 = 6400; total 8*6400 = 51200 = 10*10*512.

Example 2 — 1 head, same total

Input:  n = 10, d_model = 512, h = 1
Output: (51200, 51200)

Constraints

• d_k = d_model // h.
• per_head_ops = n*n*d_k; total_ops = h * per_head_ops.
• Return (int, int).

Notes

• The total_ops is n² · d_model regardless of h — you get multiple parallel attention views essentially for free (within constants).