GPT-2 Small · Masked multi-head self-attention

Transformer block › Attention

Prepare Q, K, V

learnedattn · model.py 69–112

Projects the normalized input into queries, keys and values, and divides each into attention heads.

\begin{aligned}&qkv={\htmlData{light=attn-input}{x}}W_{\mathrm{QKV}}+b_{\mathrm{QKV}}\end{aligned}

in: $x$ [B batch, T positions, 768 d]
out: $Q$ [B batch, 12 heads, T positions, 64 d_h] $K$ [B batch, 12 heads, T positions, 64 d_h] $V$ [B batch, 12 heads, T positions, 64 d_h]
parameters: 1,771,776 · × 12 blocks = 21,261,312

K and V, head-first, are stacked into present, this call’s cache entry (model.py 104); with past, the cached keys and values are prepended along the position axis (106–108).

Parameters

1,771,776 parameters per block, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Fused Q, K, V projection	$768 \times 2{,}304 + 2{,}304$	1,771,776
Reshape into attention heads	—	0
Per-instance total		1,771,776

Across 12 independent blocks: $1{,}771{,}776 \times 12 = 21{,}261{,}312$ parameters.

The group references these owners; it adds no weights of its own.

Explanation

One learned projection produces the fused field of 2,304 features: $qkv=xW_{\mathrm{QKV}}+b_{\mathrm{QKV}}$ . Splitting it gives $Q$ , $K$ and $V$ of 768 features each, and each is reshaped into 12 heads of 64 features. Only the projection owns parameters; splitting and reshaping reorganize values without changing them.

Operations: Fused Q, K, V projection · Q, K, V split · Reshape into attention heads
Intermediate states: Query field · Key field · Value field

Codemodel.py 50–103 · modeling_gpt2.py 185–191

model.py 50–56

def conv1d(x, scope, nf, *, w_init_stdev=0.02):
    with tf.variable_scope(scope):
        *start, nx = shape_list(x)
        w = tf.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
        b = tf.get_variable('b', [nf], initializer=tf.constant_initializer(0))
        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])
        return c

model.py 102–103

        c = conv1d(x, 'c_attn', n_state*3)
        q, k, v = map(split_heads, tf.split(c, 3, axis=2))

modeling_gpt2.py 185–191

            query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
            shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
            key_states = key_states.view(shape_kv).transpose(1, 2)
            value_states = value_states.view(shape_kv).transpose(1, 2)

        shape_q = (*query_states.shape[:-1], -1, self.head_dim)
        query_states = query_states.view(shape_q).transpose(1, 2)

model.py 75

    def split_heads(x):

model.py 77

        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])

model.py 103

        q, k, v = map(split_heads, tf.split(c, 3, axis=2))

Sources4

Evidence for the operations this group contains. The group itself is a presentation and has no implementation of its own.

Contextual evidence from Fused Q, K, V projection

The released projection helper owns a learned matrix and output bias and applies them to the last feature axis.

linear-operation · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 50–56

Read code excerpt

def conv1d(x, scope, nf, *, w_init_stdev=0.02):
    with tf.variable_scope(scope):
        *start, nx = shape_list(x)
        w = tf.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
        b = tf.get_variable('b', [nf], initializer=tf.constant_initializer(0))
        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])
        return c

Contextual evidence from Fused Q, K, V projection

The released attention module creates one 3× projection and splits it into Q, K, and V head fields.

openai · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 101–104

Read code excerpt

        c = conv1d(x, 'c_attn', n_state*3)
        q, k, v = map(split_heads, tf.split(c, 3, axis=2))

Contextual evidence from Fused Q, K, V projection

The maintained implementation applies c_attn once, splits the three fields, and reshapes each to heads.

transformers · huggingface/transformers / src/transformers/models/gpt2/modeling_gpt2.py
Pinned revision 71c6f699ac9b ↗ · lines 185–191

Read code excerpt

            query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
            shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
            key_states = key_states.view(shape_kv).transpose(1, 2)
            value_states = value_states.view(shape_kv).transpose(1, 2)

        shape_q = (*query_states.shape[:-1], -1, self.head_dim)
        query_states = query_states.view(shape_q).transpose(1, 2)

Contextual evidence from Q, K, V split

The fused field is split three ways and each result is reshaped into the head axis.

openai · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 75–103

Read code excerpt

    def split_heads(x):

… lines 76–76 omitted …

        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])

… lines 78–102 omitted …

        q, k, v = map(split_heads, tf.split(c, 3, axis=2))

Contextual evidence from Reshape into attention heads

The released split_heads function introduces a head axis and moves it before the position axis.

openai · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 75–79

Read code excerpt

    def split_heads(x):

… lines 76–76 omitted …

        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])

GPT-2

GPT-2 attention computation flow