wenmin-wu/nlp-spatial-dropout-embedding

Spatial Dropout on Embeddings

Overview

Standard dropout on word embeddings randomly zeros individual elements, which can fragment the signal within a single embedding dimension. Spatial dropout drops entire channels (feature dimensions) consistently across all timesteps, forcing the model to not rely on any single feature. More effective for sequential models (LSTM, CNN) where temporal coherence matters.

Quick Start

import torch
import torch.nn as nn

class SpatialDropout(nn.Dropout2d):
    """Drop entire embedding dimensions across all timesteps."""
    def forward(self, x):
        # x: (batch, seq_len, embed_dim)
        x = x.unsqueeze(2)              # (B, T, 1, D)
        x = x.permute(0, 3, 2, 1)       # (B, D, 1, T)
        x = super().forward(x)           # drops entire D channels
        x = x.permute(0, 3, 2, 1)       # (B, T, 1, D)
        x = x.squeeze(2)                # (B, T, D)
        return x

# Usage in model
class TextModel(nn.Module):
    def __init__(self, embedding_matrix, p_drop=0.2):
        super().__init__()
        self.embedding = nn.Embedding.from_pretrained(
            torch.tensor(embedding_matrix, dtype=torch.float32))
        self.spatial_dropout = SpatialDropout(p_drop)
        self.lstm = nn.LSTM(300, 128, bidirectional=True, batch_first=True)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.spatial_dropout(x)
        x, _ = self.lstm(x)
        return x

Key Decisions

• p=0.2: typical range 0.1–0.3; higher for larger embeddings
• Before LSTM/CNN: apply after embedding, before sequential layer
• Leverages Dropout2d: reshaping trick reuses PyTorch's spatial dropout implementation
• Not for transformers: transformers use attention dropout instead; spatial dropout is for RNN/CNN pipelines

References

• Source: simple-lstm-pytorch-version
• Competition: Jigsaw Unintended Bias in Toxicity Classification