ALBERT › Cross-Layer Parameter Sharing
Language AI HandbookAll chapters ALBERT › Cross-Layer Parameter SharingOne section · Sections in this chapter5 / 18 ↓
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
ALBERT's second major innovation is sharing parameters across transformer layers. Rather than each layer having its own attention and feed-forward weights, all layers share a single set of parameters. The model applies the same transformation repeatedly, with only the input changing between iterations.
A technique where transformer layers share the same weight matrices instead of learning independent parameters. This reduces model size proportionally to the number of layers while creating a form of iterative refinement where the same transformation is applied multiple times to progressively update representations.
class SharedTransformerBlock(nn.Module):
"""A single transformer block that can be applied multiple times."""
def __init__(self, hidden_dim, num_heads, intermediate_dim, dropout=0.1):
super().__init__()
# Multi-head self-attention
self.attention = nn.MultiheadAttention(
embed_dim=hidden_dim,
num_heads=num_heads,
dropout=dropout,
batch_first=True,
)
self.attention_norm = nn.LayerNorm(hidden_dim)
# Feed-forward network
self.ffn = nn.Sequential(
nn.Linear(hidden_dim, intermediate_dim),
nn.GELU(),
nn.Linear(intermediate_dim, hidden_dim),
nn.Dropout(dropout),
)
self.ffn_norm = nn.LayerNorm(hidden_dim)
def forward(self, x, attention_mask=None):
# Self-attention with residual
attn_out, _ = self.attention(x, x, x, key_padding_mask=attention_mask)
x = self.attention_norm(x + attn_out)
# FFN with residual
ffn_out = self.ffn(x)
x = self.ffn_norm(x + ffn_out)
return xclass ALBERTEncoder(nn.Module):
"""ALBERT encoder with shared transformer layers."""
def __init__(
self,
vocab_size,
embedding_dim,
hidden_dim,
num_heads,
intermediate_dim,
num_layers,
max_position=512,
dropout=0.1,
):
super().__init__()
# Factorized embeddings
self.token_embedding = nn.Embedding(vocab_size, embedding_dim)
self.position_embedding = nn.Embedding(max_position, embedding_dim)
self.segment_embedding = nn.Embedding(2, embedding_dim)
# Project to hidden dimension
self.embedding_projection = nn.Linear(embedding_dim, hidden_dim)
self.embedding_norm = nn.LayerNorm(hidden_dim)
self.embedding_dropout = nn.Dropout(dropout)
# Single shared transformer block
self.shared_layer = SharedTransformerBlock(
hidden_dim, num_heads, intermediate_dim, dropout
)
self.num_layers = num_layers
def forward(self, input_ids, segment_ids=None, attention_mask=None):
batch_size, seq_len = input_ids.shape
# Create position IDs
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
# Get embeddings
x = self.token_embedding(input_ids)
x = x + self.position_embedding(positions)
if segment_ids is not None:
x = x + self.segment_embedding(segment_ids)
# Project to hidden dimension
x = self.embedding_projection(x)
x = self.embedding_norm(x)
x = self.embedding_dropout(x)
# Apply shared layer multiple times
for _ in range(self.num_layers):
x = self.shared_layer(x, attention_mask)
return xLet's compare the parameter counts:
def count_parameters(model):
"""Count trainable parameters in a model."""
return sum(p.numel() for p in model.parameters() if p.requires_grad)
# Create ALBERT-style encoder
torch.manual_seed(42)
albert_encoder = ALBERTEncoder(
vocab_size=30522,
embedding_dim=128,
hidden_dim=768,
num_heads=12,
intermediate_dim=3072,
num_layers=12,
)Parameter Comparison (12 layers): BERT-Base: 109,482,240 parameters ALBERT-Base: 11,161,088 parameters Reduction: 89.8% ALBERT-Base Parameter Breakdown: Embeddings + projection: 4,073,216 Shared transformer layer: 7,087,872 Total: 11,161,088
