Architectures / NLP/LLM

🤖 Transformer Block

Transformer encoder block (simplified)

From Vaswani et al. (2017). Attention Is All You Need. NeurIPS 2017. This page is the graph, not the PDF: open it, edit it, verify it, export it.

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

8 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 43 single-graph checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()

        self.multiHeadAttention_1 = nn.MultiheadAttention(embed_dim=768, num_heads=8, batch_first=True)
        self.layerNorm_1 = nn.LayerNorm(768)
        self.feedForward_1 = nn.Sequential(
            nn.Linear(768, 3072),
            nn.ReLU(),
            nn.Linear(3072, 768)
        )
        self.layerNorm_2 = nn.LayerNorm(768)

    def forward(self, x):
        # Input shape: [512,768]
        multi_head_attention_tion_1 = self.multiHeadAttention_1(x, x, x)[0]
        add_add_1 = multi_head_attention_tion_1
        layer_norm_Norm_1 = self.layerNorm_1(add_add_1)
        feed_forward_ward_1 = self.feedForward_1(layer_norm_Norm_1)
        add_add_2 = feed_forward_ward_1
        layer_norm_Norm_2 = self.layerNorm_2(add_add_2)
        # Output
        return layer_norm_Norm_2


if __name__ == '__main__':
    model = TransformerBlock()
    model.eval()

    x = torch.randn(1, 512, 768)  # (batch, channels, length)
    with torch.no_grad():
        output = model(x)

    print(f'Input  shape : {tuple(x.shape)}')
    print(f'Output shape : {tuple(output.shape)}')
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

For agents

Also in NLP/LLM