Transformer Architecture: Residual Streams, Attention Blocks, and Autoregressive Generation

This is not a large-model training script; it explicitly runs the data flow through positional input, MHA, residual paths, FFN, and logits.

import numpy as np

def layer_norm(x, eps=1e-5):
    mean = x.mean(axis=-1, keepdims=True)
    var = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
    return (x - mean) / np.sqrt(var + eps)

def softmax(a, axis=-1):
    a = a - np.max(a, axis=axis, keepdims=True)
    exp = np.exp(a)
    return exp / exp.sum(axis=axis, keepdims=True)

def mha(x, weights, n_heads=2, mask=None):
    n, d_model = x.shape
    d_head = d_model // n_heads
    Q = x @ weights["Wq"]
    K = x @ weights["Wk"]
    V = x @ weights["Wv"]

    Q = Q.reshape(n, n_heads, d_head).transpose(1, 0, 2)
    K = K.reshape(n, n_heads, d_head).transpose(1, 0, 2)
    V = V.reshape(n, n_heads, d_head).transpose(1, 0, 2)

    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_head)
    if mask is not None:
        scores = np.where(mask[None, :, :], scores, -1e9)
    attn = softmax(scores, axis=-1)
    context = attn @ V
    context = context.transpose(1, 0, 2).reshape(n, d_model)
    return context @ weights["Wo"], attn

def ffn(x, weights):
    hidden = np.maximum(0, x @ weights["W1"] + weights["b1"])
    return hidden @ weights["W2"] + weights["b2"]

rng = np.random.default_rng(11)
n, d_model, d_ff, vocab = 5, 8, 24, 32
tokens = rng.integers(0, vocab, size=n)

embedding = rng.normal(size=(vocab, d_model)) / np.sqrt(d_model)
position = rng.normal(size=(n, d_model)) / np.sqrt(d_model)
H = embedding[tokens] + position

weights = {
    "Wq": rng.normal(size=(d_model, d_model)) / np.sqrt(d_model),
    "Wk": rng.normal(size=(d_model, d_model)) / np.sqrt(d_model),
    "Wv": rng.normal(size=(d_model, d_model)) / np.sqrt(d_model),
    "Wo": rng.normal(size=(d_model, d_model)) / np.sqrt(d_model),
    "W1": rng.normal(size=(d_model, d_ff)) / np.sqrt(d_model),
    "b1": np.zeros(d_ff),
    "W2": rng.normal(size=(d_ff, d_model)) / np.sqrt(d_ff),
    "b2": np.zeros(d_model),
    "W_vocab": rng.normal(size=(d_model, vocab)) / np.sqrt(d_model),
}

causal_mask = np.tril(np.ones((n, n), dtype=bool))
A, attention_weights = mha(layer_norm(H), weights, n_heads=2, mask=causal_mask)
U = H + A
H_next = U + ffn(layer_norm(U), weights)
logits = H_next[-1] @ weights["W_vocab"]
probs = softmax(logits)

print("hidden shape:", H_next.shape)
print("attention shape:", attention_weights.shape)
print("top next-token ids:", probs.argsort()[-5:][::-1])

• In StatsPAI-style work, a research question, data dictionary, regression table, footnotes, and reviewer comments are tokenized. Each token enters the residual stream through embedding plus position.
• Attention lets “treatment effect” attend to table headers, estimands, standard errors, sample restrictions, and identification-design paragraphs; FFN layers re-encode each local token representation nonlinearly.
• After many stacked layers, the model can generate explanations that combine definitions, table structure, and method context, but this is representational power rather than proof of causal validity.
• In an agent system, Transformers are good at reading and writing research material; reliable workflows still require tool calls, rerunnable code, data provenance, and human review checkpoints.

References

• Vaswani et al. (2017), Attention Is All You Needhttps://arxiv.org/abs/1706.03762
• NeurIPS Proceedings: Attention Is All You Needhttps://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html
• Devlin et al. (2018), BERT: Pre-training of Deep Bidirectional Transformers for Language Understandinghttps://arxiv.org/abs/1810.04805
• Radford et al. (2019), Language Models are Unsupervised Multitask Learnershttps://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf