Crystal Clear about Dimensions in LLM Training

What is a neuron? A neuron is a learned function that takes multiple input values, applies its own weights and bias, and produces one output value.

In the simple house-price NN example,

in GPT-2, each token’s 768 numbers are fed into 3,072 neurons, producing 3,072 intermediate values, then compressed back to 768. overall,

MLP parameters across 12 blocks in GPT-2 are

X, the input lives in a tensor of B(batch size), T(token length) and C(dimensions), for one batch,

                 X
             T × 768
                 │
        ┌────────┼────────┐
        ▼        ▼        ▼
       Q         K        V
    T × 768   T × 768  T × 768
        │        │        │
        ▼        ▼
      split into 12 heads
        │
        ▼
    T × 64 per head
        │
        ▼
      Q Kᵀ
        │
        ▼
      T × T
        │
      softmax
        │
        ▼
        A
      T × T
        │
        │ × V
        ▼
      T × 64
        │
   12 heads
        │
        ▼
     concat
     T × 768
        │
        │ × Wo
        ▼
   T × 768
        │
        ▼
 AttentionOutput

Leave a Reply