Longformer › Memory and Speed Analysis

def measure_attention_memory(
    seq_lengths, embed_dim=768, num_heads=12, window_size=512
):
    """Estimate memory usage for attention matrices."""
    results = []

    for seq_len in seq_lengths:
        # Full attention: n x n attention matrix per head per layer
        full_attention_size = seq_len * seq_len * num_heads

        # Longformer: n x w for local + n x g for global (approximate)
        # Assuming 1 global token
        num_global = 1
        longformer_size = seq_len * (window_size + 2 * num_global) * num_heads

        # Convert to MB (assuming float32)
        bytes_per_float = 4
        full_mb = full_attention_size * bytes_per_float / (1024**2)
        longformer_mb = longformer_size * bytes_per_float / (1024**2)

        results.append(
            {
                "seq_len": seq_len,
                "full_attention_mb": full_mb,
                "longformer_mb": longformer_mb,
                "savings_ratio": longformer_mb / full_mb,
            }
        )

    return results


seq_lengths = [512, 1024, 2048, 4096, 8192]
memory_results = measure_attention_memory(seq_lengths)