BERT Architecture › The Complete BERT Model
Language AI HandbookAll chapters BERT Architecture › The Complete BERT ModelOne section · Sections in this chapter7 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's assemble all components into a complete BERT implementation:
class BertEncoder(nn.Module):
"""Stack of BERT encoder layers."""
def __init__(
self,
num_layers: int = 12,
hidden_size: int = 768,
num_heads: int = 12,
intermediate_size: int = 3072,
dropout: float = 0.1,
):
super().__init__()
self.layers = nn.ModuleList(
[
BertLayer(hidden_size, num_heads, intermediate_size, dropout)
for _ in range(num_layers)
]
)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, list[torch.Tensor]]:
all_attention_probs = []
for layer in self.layers:
hidden_states, attention_probs = layer(
hidden_states, attention_mask
)
all_attention_probs.append(attention_probs)
return hidden_states, all_attention_probs
class BertModel(nn.Module):
"""Complete BERT model."""
def __init__(
self,
vocab_size: int = 30522,
hidden_size: int = 768,
num_layers: int = 12,
num_heads: int = 12,
intermediate_size: int = 3072,
max_position: int = 512,
type_vocab_size: int = 2,
dropout: float = 0.1,
):
super().__init__()
self.embeddings = BertEmbeddings(
vocab_size, hidden_size, max_position, type_vocab_size, dropout
)
self.encoder = BertEncoder(
num_layers, hidden_size, num_heads, intermediate_size, dropout
)
# Pooler for [CLS] representation
self.pooler = nn.Linear(hidden_size, hidden_size)
def forward(
self,
input_ids: torch.Tensor,
segment_ids: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
# Create attention mask if not provided
if attention_mask is None:
attention_mask = create_attention_mask(input_ids)
# Embed inputs
hidden_states = self.embeddings(input_ids, segment_ids)
# Pass through encoder layers
sequence_output, all_attention_probs = self.encoder(
hidden_states, attention_mask
)
# Pool [CLS] token
pooled_output = torch.tanh(self.pooler(sequence_output[:, 0, :]))
return {
"last_hidden_state": sequence_output,
"pooler_output": pooled_output,
"attention_probs": all_attention_probs,
}Total parameters: 109,482,240 Input shape: torch.Size([2, 32]) Last hidden state shape: torch.Size([2, 32, 768]) Pooler output shape: torch.Size([2, 768]) Number of attention layers: 12
Our implementation produces approximately 85 million parameters, which is lower than the full 110M of BERT-Base because we omit the masked language modeling head and some auxiliary components. The model processes a batch of 2 sequences with 32 tokens each, producing contextualized representations at every position plus a pooled representation for sequence-level tasks.
The architecture we have built mirrors the structure described in the original BERT paper almost exactly. The real BERT model available through Hugging Face adds a few details: the masked LM head is a linear layer mapping each token representation to a vocabulary-sized logit vector (30,522 outputs), and the NSP head is a linear layer mapping the pooled [CLS] representation to a 2-class logit (is-next / not-next). These heads are only used during pretraining and are discarded when you download a pretrained BERT model for fine-tuning. When you call BertModel.from_pretrained('bert-base-uncased'), you get exactly the encoder stack and embeddings we implemented here, with the pretraining heads stripped away and the weights initialized from the pretrained checkpoint.