ALBERT › Implementing a Complete ALBERT Model

class ALBERT(nn.Module):
    """Complete ALBERT model with factorized embeddings and shared layers."""

    def __init__(
        self,
        vocab_size=30522,
        embedding_dim=128,
        hidden_dim=768,
        num_heads=12,
        intermediate_dim=3072,
        num_layers=12,
        max_position=512,
        dropout=0.1,
    ):
        super().__init__()

        self.hidden_dim = hidden_dim

        # 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)

        # Projection 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

        # Pooler for sentence-level representations
        self.pooler = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim), nn.Tanh()
        )

        # MLM head (shares embedding weights for output projection)
        self.mlm_dense = nn.Linear(hidden_dim, embedding_dim)
        self.mlm_norm = nn.LayerNorm(embedding_dim)
        self.mlm_decoder = nn.Linear(embedding_dim, vocab_size)
        # Weight tying: reuse token embeddings for output
        self.mlm_decoder.weight = self.token_embedding.weight

        # SOP head
        self.sop_classifier = nn.Linear(hidden_dim, 2)

    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)

        # 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 and normalize
        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)

        # Sequence output
        sequence_output = x

        # Pooled output from [CLS] token (position 0)
        pooled_output = self.pooler(x[:, 0])

        return sequence_output, pooled_output

    def get_mlm_logits(self, sequence_output):
        """Get masked language modeling predictions."""
        x = self.mlm_dense(sequence_output)
        x = F.gelu(x)
        x = self.mlm_norm(x)
        logits = self.mlm_decoder(x)
        return logits

    def get_sop_logits(self, pooled_output):
        """Get sentence order prediction logits."""
        return self.sop_classifier(pooled_output)

# Create model and test
torch.manual_seed(42)
albert = ALBERT(
    vocab_size=30522,
    embedding_dim=128,
    hidden_dim=768,
    num_heads=12,
    intermediate_dim=3072,
    num_layers=12,
)

# Random input
input_ids = torch.randint(0, 30522, (2, 64))
segment_ids = torch.zeros_like(input_ids)
segment_ids[:, 32:] = 1  # Second half is segment B

sequence_output, pooled_output = albert(input_ids, segment_ids)
mlm_logits = albert.get_mlm_logits(sequence_output)
sop_logits = albert.get_sop_logits(pooled_output)