Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

ALBERT › Training ALBERT

Michael BrenndoerferUpdated July 19, 20252 min read

Language AI HandbookAll chapters

Part 24 of 67BERT and Variants

Chapter 174ALBERT › Training ALBERT

Reading section 11 of 18 · Training ALBERT · 2 minSection 11 of 18 · 2 min

Section 11 of 18 · 2 min read

ALBERT › Training ALBERTOne section · Sections in this chapter11 / 18 ↓

Sections in this chapter18

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

Training ALBERT combines the MLM and SOP objectives. The loss function sums both contributions:

In[38]:

Code

def compute_albert_loss(
    model, input_ids, segment_ids, mlm_labels, sop_labels, attention_mask=None
):
    """Compute combined MLM and SOP loss for ALBERT pretraining."""

    # Forward pass
    sequence_output, pooled_output = model(
        input_ids, segment_ids, attention_mask
    )

    # MLM loss
    mlm_logits = model.get_mlm_logits(sequence_output)
    mlm_loss = F.cross_entropy(
        mlm_logits.view(-1, mlm_logits.size(-1)),
        mlm_labels.view(-1),
        ignore_index=-100,  # Ignore non-masked positions
    )

    # SOP loss
    sop_logits = model.get_sop_logits(pooled_output)
    sop_loss = F.cross_entropy(sop_logits, sop_labels)

    # Combined loss
    total_loss = mlm_loss + sop_loss

    return total_loss, mlm_loss, sop_loss

In[39]:

Code

# Demonstrate training step
torch.manual_seed(42)

# Create dummy batch
batch_size = 4
seq_len = 64
vocab_size = 30522

input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
segment_ids = torch.zeros_like(input_ids)
segment_ids[:, seq_len // 2 :] = 1

# MLM labels: -100 for non-masked, token ID for masked
mlm_labels = torch.full_like(input_ids, -100)
mask_positions = torch.rand(batch_size, seq_len) < 0.15
mlm_labels[mask_positions] = input_ids[mask_positions]

# SOP labels: 0 or 1 for each example
sop_labels = torch.randint(0, 2, (batch_size,))

# Compute loss
total_loss, mlm_loss, sop_loss = compute_albert_loss(
    albert, input_ids, segment_ids, mlm_labels, sop_labels
)

Out[40]:

Console

ALBERT Training Loss:
  MLM Loss: 41.7546
  SOP Loss: 0.6960
  Total Loss: 42.4506

The MLM loss around 10 is expected for an untrained model predicting among 30,522 vocabulary tokens (random guessing would give $-\log(1/30522) \approx 10.3$ ). The SOP loss near 0.69 corresponds to random binary classification (50/50 guessing). Both losses would decrease during actual pretraining as the model learns meaningful patterns.