ホモジニアスな注意機構を超えて：フーリエ近似KVキャッシュによるメモリ効率的なLLM

コアコンセプト

FourierAttentionは、トランスフォーマーのヘッド次元が異種的な役割を持つことを発見することで、Key-Valueキャッシュからのメモリボトルネックに対処します。下位の次元はローカルコンテキストを処理し、上位の次元は長距離依存関係をキャプチャします。本手法は、直交フーリエ基底関数を使用してコンテキスト非依存な次元を圧縮し、優位なスペクトル係数のみを保持しながら長コンテキスト性能を維持します。

アーキテクチャ概要

• 次元分析ステージ: ノイズ摂動実験を通じて、完全保持が必要なヘッド次元と圧縮対象の次元を識別します
• フーリエ圧縮: プリフィリング中に平行移動フーリエ変換を使用してシーケンスを固定長スペクトル表現に変換します（k ≪ L係数）
• 選択的保持: 初期トークン、最近のローカルトークン、長コンテキスト依存次元は非圧縮のまま保持し、重要度の低い次元の中間範囲トークンを圧縮します
• カスタムTritonカーネル: デコード中のメモリ効率を最適化し、パフォーマンスを損なわない実装を提供します

実装

ステップ1: 注意ヘッドの次元を分析

ノイズ摂動を使用して、コンテキスト非依存で圧縮可能なヘッド次元を特定するための次元レベルの感度分析を実行します：

def analyze_head_dimensions(attention_module, sequence_length):
    """
    Identifies which dimensions can be compressed using noise perturbation.
    Dimensions with low sensitivity to noise are candidates for compression.
    """
    import torch

    num_heads = attention_module.num_heads
    head_dim = attention_module.head_dim
    sensitivity_scores = torch.zeros(head_dim)

    # Perturb each dimension and measure impact on output
    for dim in range(head_dim):
        original_output = attention_module(queries, keys, values)

        # Add noise to this dimension
        noisy_values = values.clone()
        noise = torch.randn_like(noisy_values[:, :, dim:dim+1]) * 0.1
        noisy_values[:, :, dim:dim+1] += noise

        noisy_output = attention_module(queries, keys, noisy_values)
        sensitivity_scores[dim] = torch.norm(original_output - noisy_output)

    # Lower sensitivity dims are compressible
    return sensitivity_scores

ステップ2: フーリエ基底圧縮を適用

KVキャッシュシーケンスを固定長フーリエ基底関数に投影します：

def fourier_compress_kv_cache(kv_sequence, num_coefficients):
    """
    Compresses KV cache using Fourier basis projection.
    Exploits shift-invariance for efficient single-pass computation.

    Args:
        kv_sequence: [batch, seq_len, head_dim] or [batch, seq_len, 2, num_heads, head_dim]
        num_coefficients: k, where k << seq_len
    """
    import torch
    import torch.fft as fft

    # Apply FFT along sequence dimension
    freq_domain = fft.rfft(kv_sequence, dim=1)

    # Retain only k dominant coefficients
    compressed = freq_domain[:, :num_coefficients]

    # Store metadata: original sequence length and coefficient indices
    return {
        'coefficients': compressed,
        'original_length': kv_sequence.shape[1],
        'num_coefficients': num_coefficients
    }

ステップ3: 選択的なトークン保持

重要度分析に基づいて特定のトークンを非圧縮のままで保持します：

def selective_preserve_tokens(kv_cache, preserve_initial=32,
                              preserve_recent=64, compress_middle=True):
    """
    Preserves initial tokens, recent tokens, and important dimensions.
    Compresses middle-range tokens in less critical dimensions.

    Returns mixed representation with full and compressed tokens.
    """
    batch_size, seq_len, head_dim = kv_cache.shape

    # Full representations for important ranges
    initial_tokens = kv_cache[:, :preserve_initial]
    recent_tokens = kv_cache[:, -preserve_recent:]

    # Compress middle tokens
    if compress_middle and seq_len > (preserve_initial + preserve_recent):
        middle_start = preserve_initial
        middle_end = seq_len - preserve_recent
        middle_tokens = kv_cache[:, middle_start:middle_end]

        # Apply Fourier compression with lower coefficient count
        compressed_middle = fourier_compress_kv_cache(
            middle_tokens,
            num_coefficients=int(middle_tokens.shape[1] * 0.3)
        )

        return {
            'initial': initial_tokens,
            'recent': recent_tokens,
            'middle_compressed': compressed_middle
        }

    return {'initial': initial_tokens, 'recent': recent_tokens}

ステップ4: デコード中の復元

フォワードパス中に圧縮表現を完全な次元に展開します：

def reconstruct_from_fourier(compressed_cache, original_length):
    """
    Reconstructs full KV cache from Fourier coefficients during decoding.
    """
    import torch
    import torch.fft as fft

    coefficients = compressed_cache['coefficients']

    # Pad coefficients to original length
    padded = torch.zeros(
        (coefficients.shape[0], original_length, coefficients.shape[2]),
        device=coefficients.device,
        dtype=coefficients.dtype
    )
    padded[:, :coefficients.shape[1]] = coefficients

    # Inverse FFT to reconstruct sequence
    reconstructed = fft.irfft(padded, n=original_length, dim=1)

    return reconstructed

実践的なガイダンス

• 閾値チューニング: ターゲット圧縮率（テスト済み76%）とシーケンス長に基づいてnum_coefficientsを調整します
• 次元割り当て: 感度スコアを使用してヘッド次元を完全保持グループと圧縮グループに分割します
• メモリ計算トレードオフ: フーリエ復元は高速です。主な節約はKVキャッシュサイズの削減から得られます
• ベンチマーク: 長コンテキストタスク（LongBench、Needle-In-A-Haystack）で検証し、品質保持を確認します
• 統合: プリフィリングとデコード中のメモリバウンド操作に対して効率的なカスタムTritonカーネルをデプロイします

参考文献

論文: arXiv:2506.11886 主要指標: 76%のKVキャッシュ圧縮、長コンテキストベンチマークにおける優れたパフォーマンス 関連研究: KVキャッシュ量子化、スパース注意、位置補間