Masked Prediction

# Original sentence
tokens = ["The", "patient", "presented", "with", "severe", "chest", "pain"]
 
# Mask 15% of tokens randomly
masked_tokens = ["The", "patient", "[MASK]", "with", "severe", "chest", "pain"]
 
# Model predicts: "presented"
predictions = bert(masked_tokens)
loss = cross_entropy(predictions[2], original_tokens[2])

import torch
import torch.nn as nn
 
class BERT(nn.Module):
    """Bidirectional Encoder with masked language modeling"""
 
    def __init__(self, vocab_size, hidden_size=768, num_layers=12, num_heads=12):
        super().__init__()
 
        # Token and position embeddings
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = nn.Embedding(512, hidden_size)  # Max 512 tokens
 
        # Stack of bidirectional transformer encoders
        self.transformer_layers = nn.ModuleList([
            TransformerEncoderLayer(
                d_model=hidden_size,
                nhead=num_heads,
                dim_feedforward=hidden_size * 4
            )
            for _ in range(num_layers)
        ])
 
        self.layer_norm = nn.LayerNorm(hidden_size)
 
        # Masked language modeling head
        self.mlm_head = nn.Linear(hidden_size, vocab_size)
 
    def forward(self, token_ids, masked_positions):
        """
        Args:
            token_ids: [batch_size, seq_len] - input tokens (some masked)
            masked_positions: [batch_size, seq_len] - binary mask of positions to predict
 
        Returns:
            predictions: [batch_size, seq_len, vocab_size] - logits for masked positions
        """
        batch_size, seq_len = token_ids.shape
 
        # Create embeddings
        token_embed = self.token_embedding(token_ids)
        positions = torch.arange(seq_len, device=token_ids.device)
        position_embed = self.position_embedding(positions)
 
        # Combine embeddings
        x = token_embed + position_embed  # [batch_size, seq_len, hidden_size]
 
        # Pass through transformer layers
        for layer in self.transformer_layers:
            x = layer(x)
 
        x = self.layer_norm(x)
 
        # Predict all token positions
        predictions = self.mlm_head(x)  # [batch_size, seq_len, vocab_size]
 
        return predictions
 
 
def create_masked_lm_batch(tokens, mask_prob=0.15, vocab_size=30000):
    """Create masked tokens following BERT's 80-10-10 strategy"""
    masked_tokens = tokens.clone()
    targets = tokens.clone()
 
    # Sample 15% of positions to mask
    mask_positions = torch.rand(tokens.shape) < mask_prob
 
    for i in range(tokens.shape[0]):
        for j in range(tokens.shape[1]):
            if mask_positions[i, j]:
                rand = torch.rand(1).item()
 
                if rand < 0.8:
                    # 80%: Replace with [MASK] token
                    masked_tokens[i, j] = MASK_TOKEN_ID
                elif rand < 0.9:
                    # 10%: Replace with random token
                    masked_tokens[i, j] = torch.randint(0, vocab_size, (1,)).item()
                # else: 10% keep original token
 
    return masked_tokens, targets, mask_positions
 
 
# Training loop
def train_bert_mlm(model, data_loader, epochs=10):
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    model.train()
 
    for epoch in range(epochs):
        for batch_tokens in data_loader:
            # Create masked batch
            masked_tokens, targets, mask_positions = create_masked_lm_batch(
                batch_tokens, mask_prob=0.15
            )
 
            # Forward pass
            predictions = model(masked_tokens, mask_positions)
 
            # Compute loss only on masked positions
            predictions_masked = predictions[mask_positions]
            targets_masked = targets[mask_positions]
            loss = nn.CrossEntropyLoss()(predictions_masked, targets_masked)
 
            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
 
            if step % 100 == 0:
                print(f"Epoch {epoch}, Step {step}, Loss: {loss.item():.4f}")
 
    return model

class MAE(nn.Module):
    """Masked Autoencoder for self-supervised vision pre-training"""
 
    def __init__(self, img_size=224, patch_size=16, in_channels=3,
                 encoder_embed_dim=1024, encoder_depth=24, encoder_num_heads=16,
                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16):
        super().__init__()
 
        self.patch_size = patch_size
        num_patches = (img_size // patch_size) ** 2
 
        # Patch embedding
        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_channels=in_channels,
            embed_dim=encoder_embed_dim
        )
 
        # Encoder: Deep and wide (processes only visible patches)
        self.encoder = VisionTransformer(
            embed_dim=encoder_embed_dim,
            depth=encoder_depth,
            num_heads=encoder_num_heads
        )
 
        # Decoder: Shallow and narrow (reconstructs all patches)
        self.decoder = VisionTransformer(
            embed_dim=decoder_embed_dim,
            depth=decoder_depth,
            num_heads=decoder_num_heads
        )
 
        # Learnable mask token for decoder
        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
 
        # Project encoder output to decoder dimension
        self.encoder_to_decoder = nn.Linear(encoder_embed_dim, decoder_embed_dim)
 
        # Reconstruct pixels from patches
        self.decoder_pred = nn.Linear(
            decoder_embed_dim,
            patch_size ** 2 * in_channels
        )
 
    def forward(self, imgs, mask_ratio=0.75):
        """
        Args:
            imgs: [batch_size, 3, H, W] - input images
            mask_ratio: float - fraction of patches to mask
 
        Returns:
            loss: reconstruction loss on masked patches
            pred: reconstructed image
            mask: binary mask showing which patches were masked
        """
        # Patchify image
        patches = self.patch_embed(imgs)  # [batch, num_patches, patch_dim]
        batch_size, num_patches, patch_dim = patches.shape
 
        # Random masking
        shuffle_indices = torch.rand(batch_size, num_patches).argsort(dim=1)
        num_keep = int(num_patches * (1 - mask_ratio))
 
        # Indices of visible and masked patches
        visible_indices = shuffle_indices[:, :num_keep]
        masked_indices = shuffle_indices[:, num_keep:]
 
        # Keep only visible patches for encoder (25% of patches)
        visible_patches = torch.gather(
            patches, dim=1,
            index=visible_indices.unsqueeze(-1).expand(-1, -1, patch_dim)
        )
 
        # Encode visible patches
        encoded = self.encoder(visible_patches)  # [batch, num_keep, encoder_dim]
 
        # Project to decoder dimension
        encoded = self.encoder_to_decoder(encoded)
 
        # Create mask tokens for masked positions
        mask_tokens = self.mask_token.repeat(
            batch_size, num_patches - num_keep, 1
        )
 
        # Combine encoded visible patches with mask tokens
        # (Restore original order with unshuffle operation)
        decoder_input = torch.cat([encoded, mask_tokens], dim=1)
        decoder_input = self.unshuffle(decoder_input, shuffle_indices)
 
        # Decode to reconstruct patches
        decoded = self.decoder(decoder_input)  # [batch, num_patches, decoder_dim]
        pred_patches = self.decoder_pred(decoded)  # [batch, num_patches, patch_dim]
 
        # Compute loss only on masked patches
        target_patches = patches
        mask = torch.zeros(batch_size, num_patches, device=imgs.device)
        mask.scatter_(1, masked_indices, 1)
 
        loss = F.mse_loss(
            pred_patches[mask.bool()],
            target_patches[mask.bool()]
        )
 
        return loss, pred_patches, mask
 
    def unshuffle(self, x, shuffle_indices):
        """Restore original patch order"""
        batch_size, num_patches, dim = x.shape
        unshuffle_indices = torch.argsort(shuffle_indices, dim=1)
        x = torch.gather(
            x, dim=1,
            index=unshuffle_indices.unsqueeze(-1).expand(-1, -1, dim)
        )
        return x

1. Split image into patches (e.g., 224×224 image → 196 patches of 16×16 pixels)
2. Randomly mask 75% of patches (keep only 49 visible patches)
3. Encode only visible patches through deep ViT encoder
4. Add learnable mask tokens for masked patch positions
5. Decode full sequence (49 encoded + 147 mask tokens) through shallow decoder
6. Reconstruct pixel values for all patches
7. Compute MSE loss only on masked patches (147 patches)

Original Image (16 patches):
[🐕] [🐕] [🐕] [🐕]
[🐕] [🐕] [🐕] [🐕]
[🐕] [🐕] [🐕] [🐕]
[🐕] [🐕] [🐕] [🐕]

After 75% masking (⬜ = masked):
[🐕] [⬜] [⬜] [⬜]
[⬜] [🐕] [⬜] [⬜]
[⬜] [⬜] [⬜] [🐕]
[⬜] [⬜] [🐕] [⬜]

MAE Reconstruction:
[🐕] [🐕] [🐕] [🐕]
[🐕] [🐕] [🐕] [🐕]
[🐕] [🐕] [🐕] [🐕]
[🐕] [🐕] [🐕] [🐕]

class MaskedEHRModel(nn.Module):
    """Masked prediction for electronic health records"""
 
    def __init__(self, vocab_size, hidden_size=256, num_layers=6):
        super().__init__()
 
        self.event_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = nn.Embedding(5000, hidden_size)  # Max sequence length
 
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=hidden_size,
                nhead=8,
                dim_feedforward=hidden_size * 4
            ),
            num_layers=num_layers
        )
 
        self.prediction_head = nn.Linear(hidden_size, vocab_size)
 
    def forward(self, event_sequence, mask_positions):
        """
        Args:
            event_sequence: [batch, seq_len] - sequence of medical event codes
            mask_positions: [batch, seq_len] - positions to predict
 
        Returns:
            predictions: [batch, seq_len, vocab_size] - predictions for all positions
        """
        batch_size, seq_len = event_sequence.shape
 
        # Embed events and positions
        event_embed = self.event_embedding(event_sequence)
        positions = torch.arange(seq_len, device=event_sequence.device)
        pos_embed = self.position_embedding(positions)
 
        x = event_embed + pos_embed
 
        # Transformer encoding (bidirectional)
        x = self.transformer(x.transpose(0, 1))  # Transformer expects [seq, batch, dim]
        x = x.transpose(0, 1)  # Back to [batch, seq, dim]
 
        # Predict events at all positions
        predictions = self.prediction_head(x)
 
        return predictions
 
 
def mask_ehr_events(event_sequence, mask_prob=0.15, vocab_size=10000):
    """Create masked EHR sequence following BERT strategy"""
    masked_sequence = event_sequence.clone()
    targets = event_sequence.clone()
 
    # Determine which positions to mask
    mask_positions = torch.rand(event_sequence.shape) < mask_prob
 
    # Apply 80-10-10 masking strategy
    for i in range(event_sequence.shape[0]):
        for j in range(event_sequence.shape[1]):
            if mask_positions[i, j]:
                rand = torch.rand(1).item()
 
                if rand < 0.8:
                    masked_sequence[i, j] = MASK_EVENT_ID
                elif rand < 0.9:
                    masked_sequence[i, j] = torch.randint(0, vocab_size, (1,)).item()
                # else: keep original
 
    return masked_sequence, targets, mask_positions