GPT-NeoX¶

GPT-NeoX is EleutherAI's 20-billion-parameter autoregressive language model, released in 2022. It extended the architectural patterns established by GPT-J -- parallel residual connections and Rotary Position Embeddings -- to a significantly larger scale while introducing fused QKV projections and providing the community with one of the first fully open 20B-class models complete with training code and data documentation¹.

1. Architecture Overview¶

GPT-NeoX is a decoder-only transformer that shares GPT-J's parallel residual design but adds fused query-key-value projections and applies RoPE to 25% of the head dimension by default.

2. Key Innovations¶

2.1 Fused QKV Projection¶

Instead of three separate linear projections for Q, K, and V, GPT-NeoX fuses them into a single matrix multiplication:

2.2 Parallel Attention and FFN¶

Like GPT-J, GPT-NeoX computes attention and the feed-forward network from the same LayerNorm output:

This single-norm, dual-path design yields a single residual addition per block.

2.3 Advanced RoPE Scaling¶

GPT-NeoX introduced configurable RoPE parameters that later became standard:

The ability to adjust rotary_pct allows trading off positional sensitivity against pure content-based attention within each head.

3. Architecture Diagram¶

flowchart TD
    INPUT["Token IDs"] --> EMB["Token Embedding"]
    EMB --> BLOCKS

    subgraph BLOCKS["44 x NeoXBlock (Parallel Residual)"]
        LN["LayerNorm"]
        LN --> FUSED["Fused QKV Projection\n(single matmul)"]
        FUSED --> SPLIT["Split -> Q, K, V"]
        SPLIT --> ROPE["Partial RoPE\n(25% of head dims)"]
        ROPE --> MHA["Multi-Head Attention"]
        LN --> FFN["Feed-Forward\n(GELU activation)"]
        MHA --> SUM["Sum"]
        FFN --> SUM
        SUM --> RES["Residual Add with x"]
    end

    BLOCKS --> FINAL_LN["Final LayerNorm"]
    FINAL_LN --> LM_HEAD["LM Head"]

4. Configuration Parameters¶

5. Mathematical Formulation¶

5.1 Fused QKV Computation¶

For input \( x \in \mathbb{R}^{s \times d} \):

5.2 Partial RoPE¶

With rotary_pct = 0.25 and head dimension \( d_h = 96 \), the rotary dimension is \( d_{\text{rot}} = 24 \). For position \( m \):

5.3 Full Block Computation¶

where:

6. Zig Implementation¶

6.1 GPTNeoXConfig¶

pub const GPTNeoXConfig = struct {
    n_layers: u32 = 44,
    d_model: u32 = 6144,
    n_heads: u32 = 64,
    d_ff: u32 = 24576,
    vocab_size: u32 = 50432,
    max_seq_len: u32 = 2048,
    rotary_pct: f32 = 0.25,
    rotary_emb_base: f32 = 10000.0,
    parallel_residual: bool = true,
    norm_eps: f32 = 1e-5,
    activation: ActivationType = .gelu,

    pub fn headDim(self: GPTNeoXConfig) u32 {
        return self.d_model / self.n_heads;
    }

    pub fn rotaryDim(self: GPTNeoXConfig) u32 {
        return @intFromFloat(@as(f32, @floatFromInt(self.headDim()))
            * self.rotary_pct);
    }
};

6.2 Fused QKV Projection¶

pub const FusedQKVProjection = struct {
    w_qkv: Tensor(f32),    // [d_model, 3 * d_model]
    b_qkv: Tensor(f32),    // [3 * d_model]

    pub fn forward(self: *FusedQKVProjection, x: Tensor(f32)) !struct {
        q: Tensor(f32), k: Tensor(f32), v: Tensor(f32)
    } {
        // Single matmul: [seq, d] @ [d, 3d] -> [seq, 3d]
        const fused = try x.matmul(self.w_qkv, allocator);
        defer fused.deinit();

        const d = self.w_qkv.shape[0];
        return .{
            .q = try fused.slice(1, 0, d),
            .k = try fused.slice(1, d, 2 * d),
            .v = try fused.slice(1, 2 * d, 3 * d),
        };
    }
};

6.3 NeoX Block (Parallel Residual)¶

pub const NeoXBlock = struct {
    ln: LayerNorm,
    qkv: FusedQKVProjection,
    attention: MultiHeadAttention,
    ffn: FeedForward,
    out_proj: Linear,

    pub fn forward(self: *NeoXBlock, x: []f32, pos: u32) ![]f32 {
        const h = self.ln.forward(x);

        // Parallel path 1: Attention
        const qkv = try self.qkv.forward(h);
        applyPartialRoPE(&qkv.q, &qkv.k, pos, self.config.rotaryDim());
        const attn_out = try self.attention.forward(qkv.q, qkv.k, qkv.v);

        // Parallel path 2: FFN (from same normalized input)
        const ffn_out = try self.ffn.forward(h);

        // Single residual addition
        var output = try allocator.alloc(f32, x.len);
        for (0..x.len) |i| {
            output[i] = x[i] + attn_out[i] + ffn_out[i];
        }
        return output;
    }
};

7. Variants¶

8. Educational Value¶

9. References¶