Model Pruning for Low-Latency Trading Inference

DeepLOB の記事] は、ONNX と TensorRT、INT8 量子化、FPGA の 3 つの箇条書きで展開セクションを終了していますが、そのどれも扱いません。この記事では、これら 3 つの問題すべての根底にある最初の問題、つまりモデルが必要以上に大きいという問題に対する対処が欠けています。ニューラル ネットワークの枝刈りは冗長なパラメーターを削除します。文献での興味深い主張は、これによってメモリが節約されるということではなく、重みの 10 ～ 20% を保持するサブネットワークが密なモデルの精度に匹敵できるというものです。

レイテンシがまったく重要であるということは、ブログがすでに主張していることです — メッセージング パス上の ZigBolt と 損益分岐点計算による IPC 税 ] — そして スプレッド モデリング は、速いがわずかに悪い、または遅いがより良いというトレードオフをすでに所有しており、推論レイテンシー行を含む勾配ブースティング対深層学習テーブル。それらのどれも取り上げていないのは、特定のモデルを「どのように」小さくするかということです。引用ループの各段階のうち、モデル推論は私たちが完全に制御できる段階です。トランスポート レッグは、algotrading のデータ通信] で再現可能な p50/p95/p99 番号でカバーされています。

この記事の内容: トレーディング MLP に適用される、マグニチュード プルーニング、構造化プルーニング、反復的マグニチュード プルーニング、動きのプルーニング、知識の蒸留、および NVIDIA 2:4 半構造化スパース性の計算と実用的なコード。

そうではないもの: 測定結果。このブログのすべての実証記事には出典や関連レポートが含まれていますが、この記事にはまだそのどちらも含まれていません。スパース性対精度対レイテンシーの曲線は、引用する表としてではなく、実行する実験として以下に記載されています。ここにあるものはすべてメソッドとして扱い、数値は保留中のものとして扱います。

モデル剪定で得られるもの

制約はサイズ 1 です。中周波モデル、つまりオーダーブック機能上に 2048 個の隠れユニットを持つ 4 層 MLP を考えてみましょう。

$d_{\text{in}} = 100$、$h = 2048$、$L = 4$、$d_{\text{out}} = 3$ の場合、これはおよそ 1,260 万個のパラメーターになります (float32 で約 48 MB)。 L2 は通常 1 ～ 4 MB であるため、重みが適合しません。それらは、前方パスごとにさらに外側からストリーミングされます。それらの 95% をプルーニングすると、有効パラメータは約 630K、容量は 2.4 MB になり、これは適切です。

これが実時間に換算されるかどうかは、カーネルがメモリに依存しているかどうかによって決まります。また、それはサイズの問題ではなく、演算強度の問題です。 バックテスト エンジン速度ラダー は、ペナルティ係数を主張するのではなく、実測例を通じてルーフライン モデル (ウィリアムズ、ウォーターマン、パターソン) を動作させます。ここでも同じフレーミングが適用され、同じ規律が適用されます。高速化を主張する前に、移動したバイト数を測定する必要があります。

剪定の基礎

非構造化プルーニング

最も単純なアプローチ: 大きさに基づいて個々の重みを 0 に設定します。重み行列 $\mathbf{W} \in \mathbb{R}^{m \times n}$ を指定して、次のようなバイナリ マスク $\mathbf{M} \in \{0, 1\}^{m \times n}$ を作成します。

ここで、$\tau$ は、必要なスパース性レベル $s$ を達成するために選択されたしきい値です。

枝刈りされた行列は $\mathbf{W}' = \mathbf{W} \odot \mathbf{M}$ で、$\odot$ はアダマール積です。直感的には、ゼロに近い重みはレイヤーの出力にほとんど寄与しないということです。

この問題は、スパース数を読み間違えやすいため、はっきりと述べられています。 非構造化スパース性は、標準ハードウェアでは高速化につながりません。スパース カーネルまたはスパース サポートを備えたハードウェアに切り替えない限り、90% がゼロの行列でも同じ数の積和演算が発行されます。以下のコードが Sparsity: 90.0% を出力する場合、これはゼロのカウントです。これは 10 倍ではなく、高密度 CPU GEMM では 1.01 倍でもありません。時間を稼ぐパスは、以下に示す構造化プルーニング (より小さい行列) と 2:4 半構造化スパース性 (ハードウェア サポート) です。

構造化剪定

構造化された枝刈りでは、ニューロン、チャネル、またはアテンション ヘッド全体が削除されます。 $\mathbf{W} \in \mathbb{R}^{m \times n}$ を含む線形層 $\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$ の場合、ニューロン $j$ を削除すると、$\mathbf{W}$ の $j$ 番目の行と $\mathbf{b}$ の $j$ 番目の要素がゼロになります。

最小の $\ell_2$ ノルムを持つニューロンが最初に進みます。これは、実際にレイヤーを再構築する場合に限り、本当に小さな行列を生成するバリアントです。行をゼロにしてテンソルを元の形状のままにしても、FLOP カウントは何も変わりません。実装セクションの再構築ステップは、マスクを $512 \times n$ 行列に変換するものです。

畳み込み層の場合、類似するのはフィルター プルーニングです。 $\mathbf{W} \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times k \times k}$ を考えると、出力フィルター $i$ の重要性は次のようになります。

フィルター $i$ を削除すると、出力チャンネル全体が削除され、それに比例して FLOP が減少します。

宝くじ仮説

2019 年に、Frankle と Carbin は宝くじチケット仮説 (LTH) を導入しました。ランダムに初期化された密なネットワーク内には、元の初期化からトレーニングされた疎なサブネットワーク (「勝ちチケット」) が存在し、同程度の反復回数でネットワーク全体の精度と一致します。

正式には、$f(\mathbf{x}; \boldsymbol{\theta}_0)$ は $\boldsymbol{\theta}_0 \sim \mathcal{D}_\theta$ で初期化されると考えてください。収束するまでトレーニングした後、$\boldsymbol{\theta}^*$ を取得し、枝刈りマスク $\mathbf{m}$ を導出します。 LTH には、次のような $\mathbf{m}$ が存在すると記載されています。

$\|\mathbf{m}\|_0 \ll |\boldsymbol{\theta}_0|$を使用。元の実験は MNIST と CIFAR-10 で行われ、当選チケットはパラメーターの 10 ～ 20% を保持しました。仮定によって注文帳データに転送されるものは何もありません。LOB 特徴は非定常であり、ラベルはノイズに近く、重要な点で画像分類とはまったく異なります。

反復的な大きさ剪定 (IMP)

チケットは IMP によって検出されます。

1. $\boldsymbol{\theta}_0$ を使用してネットワークを初期化します。
2. 収束するまでトレーニングし、$\boldsymbol{\theta}^*$ を取得します。
3. 最小の大きさの重みの $p\%$ をプルーニングし、マスク $\mathbf{m}$ を作成します。
4. 残りの重みを $\boldsymbol{\theta}_0$ の値にリセットします (巻き戻し)。
5. マスクされたネットワークでステップ 2 から繰り返します。

各ラウンドでは端数 $p$ (通常は 20%) がプルーニングされるため、$n$ をラウンドした後、パラメーターの $(1-p)^n$ が残ります。 $p = 0.2$ で 10 ラウンドを終えると、約 10.7% が残ります。

取引モデルに関する 3 つの仮説、どれも検証されていない

LTH は市場データで特にうまく機能するはずだと主張したくなります。そのような議論が 3 つ出てきます。 3 つはすべて仮説であり、それらを事実として述べることが、このブログが回避するために存在する失敗モードです。

1. 金融シグナルはまばらです。 注文帳のスナップショットの大部分はノイズであるため、まばらなサブネットワークは自然にまばらなシグナルと一致する可能性があります。テスト可能: IMP を同じスパース性のランダム マスクと比較します。スパース性自体が機能している場合、ランダム マスクはそれほど遅れていないはずです。
2. 当選チケットは体制を超えて一般化する これは市場に関する経験的な主張であり、背後に引用はなく、3 つの中で最も興味深いものです。これは、HMMs による体制検出: 体制 A でチケットを見つけ、体制 B で再トレーニングし、B でネイティブに見つかったチケットと比較する] からの体制ラベルに対して直接テストできます。
3. 疎性は正則化します。 実効容量が低いと、微細構造ノイズへのフィッティングが低下する可能性があります。これは、単純に比較できるだけでなく、枝刈りモデルのサンプル外ギャップが密モデルよりも「小さい」こととして現れます。

実装: 取引 MLP の剪定

基本モデル

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from copy import deepcopy

class TradingMLP(nn.Module):
    """MLP for mid-price direction prediction from order book features."""
    def __init__(self, input_dim=100, hidden_dim=2048,
                 num_layers=4, output_dim=3):
        super().__init__()
        layers = []
        dims = [input_dim] + [hidden_dim] * (num_layers - 1) + [output_dim]
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i + 1]))
            if i < len(dims) - 2:
                layers.append(nn.BatchNorm1d(dims[i + 1]))
                layers.append(nn.ReLU())
                layers.append(nn.Dropout(0.1))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters())

model = TradingMLP(input_dim=100, hidden_dim=2048,
                   num_layers=4, output_dim=3)
print(f"Total parameters: {model.count_parameters():,}")

非構造化マグニチュード剪定

def apply_unstructured_pruning(model, sparsity=0.9):
    """Apply global unstructured L1 pruning to all Linear layers."""
    parameters_to_prune = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            parameters_to_prune.append((module, 'weight'))

    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=sparsity,
    )
    return model

def compute_sparsity(model):
    """Fraction of zero weights. Note: a *count*, not a speedup."""
    total, zeros = 0, 0
    for name, param in model.named_parameters():
        if 'weight' in name:
            total += param.numel()
            zeros += (param == 0).sum().item()
    return zeros / total if total > 0 else 0

pruned_model = apply_unstructured_pruning(deepcopy(model), sparsity=0.9)
print(f"Sparsity: {compute_sparsity(pruned_model):.1%}")

構造化剪定と、それを実現する再構築

行のマスキングは仕事の半分です。高速化を生み出す半分は、縮小された形状で各層を再構築することです。これは、削除を前方に伝播することを意味します。層 $i$ の行 $j$ を削除すると、層 $i+1$ の列 $j$ と、それらの間にある BatchNorm1d のチャネル $j$ も削除されます。

def apply_structured_pruning(model, fraction=0.75):
    """Mask entire neurons by L2-norm of their weight rows."""
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.out_features > 10:
            prune.ln_structured(
                module, name='weight', amount=fraction, n=2, dim=0
            )
    return model

def rebuild_pruned_mlp(model):
    """
    Physically shrink a structurally pruned TradingMLP.

    Walks the Sequential once. For each Linear: drop the input columns
    the previous layer no longer emits, then drop its own dead output
    rows. BatchNorm1d channels follow the preceding Linear's survivors.
    """
    new_layers = []
    keep_in = None  # surviving output indices of the previous Linear

    for layer in model.network:
        if isinstance(layer, nn.Linear):
            if prune.is_pruned(layer):
                prune.remove(layer, 'weight')

            W, b = layer.weight.data, layer.bias.data
            keep_out = (W.norm(dim=1) > 0).nonzero(as_tuple=True)[0]
            W = W[keep_out]
            if keep_in is not None:
                W = W[:, keep_in]

            new = nn.Linear(W.shape[1], W.shape[0])
            new.weight.data = W.clone()
            new.bias.data = b[keep_out].clone()
            new_layers.append(new)
            keep_in = keep_out

        elif isinstance(layer, nn.BatchNorm1d):
            new = nn.BatchNorm1d(len(keep_in))
            new.weight.data = layer.weight.data[keep_in].clone()
            new.bias.data = layer.bias.data[keep_in].clone()
            new.running_mean = layer.running_mean[keep_in].clone()
            new.running_var = layer.running_var[keep_in].clone()
            new.num_batches_tracked = layer.num_batches_tracked.clone()
            new_layers.append(new)

        else:  # ReLU, Dropout -- shape-agnostic, reuse as is
            new_layers.append(layer)

    rebuilt = deepcopy(model)
    rebuilt.network = nn.Sequential(*new_layers)
    return rebuilt

これを信頼する前に、ブログの残りの部分で実行されている等価性ゲートと同じ精神で、次の 2 つのことを確認してください。

• Shapes. rebuilt は、$(1 - \text{fraction}) \cdot h$ の隠れ次元 (fraction=0.75、$h = 2048$ の場合は 512) と、内部行列の両方の次元が縮小するため、二次関数的に減少するパラメーター数を表示する必要があります。
• 出力。 eval() モードでは、rebuilt(x) はマスクされたモデルの rebuilt フリー出力を同じバッチ上の浮動小数点許容値に一致させる必要があります。そうでない場合は、カラムの伝播が間違っており、下流のすべての数値が予想とは異なるモデルを測定していることになります。

行生存テストでは、マスクされた行は正確にゼロであり、ライブ行はそうではないことを前提としています。これは ln_structured 出力にも当てはまります。他のプロシージャが真にオールゼロの生きたニューロンを生成した場合、これは成立しないため、標準テストを盲目的に信頼するのではなく、要求された部分に対して生存者数をアサートします。

反復的な大きさ剪定 (宝くじ探索)

def lottery_ticket_search(model_cls, model_kwargs, train_fn, eval_fn,
                          rounds=10, prune_rate=0.2, device='cpu'):
    """
    Iterative Magnitude Pruning to find a winning ticket.

    Parameters
    ----------
    model_cls : class -- model constructor
    model_kwargs : dict -- constructor arguments
    train_fn : callable -- train_fn(model) trains the model in-place
    eval_fn : callable -- eval_fn(model) returns out-of-sample accuracy
    rounds : int -- number of pruning rounds
    prune_rate : float -- fraction of surviving weights pruned per round
    """
    model_init = model_cls(**model_kwargs).to(device)
    theta_0 = deepcopy(model_init.state_dict())

    mask = {}
    for name, param in model_init.named_parameters():
        if 'weight' in name:
            mask[name] = torch.ones_like(param, dtype=torch.bool)

    results = []

    for round_idx in range(rounds):
        model = model_cls(**model_kwargs).to(device)
        state = deepcopy(theta_0)
        for name in mask:
            state[name] = state[name] * mask[name].float()
        model.load_state_dict(state)

        train_fn(model)

        acc = eval_fn(model)
        surviving = sum(m.sum().item() for m in mask.values())
        total = sum(m.numel() for m in mask.values())
        sparsity = 1.0 - surviving / total

        results.append({
            'round': round_idx,
            'accuracy': acc,
            'sparsity': sparsity,
            'surviving_params': int(surviving)
        })
        print(f"Round {round_idx}: acc={acc:.4f}, "
              f"sparsity={sparsity:.1%}")

        all_weights = []
        for name, param in model.named_parameters():
            if name in mask:
                alive = param.data.abs()[mask[name]]
                all_weights.append(alive.flatten())

        all_weights = torch.cat(all_weights)
        k = int(len(all_weights) * prune_rate)
        if k == 0:
            break
        threshold = all_weights.kthvalue(k).values.item()

        for name, param in model.named_parameters():
            if name in mask:
                mask[name] = mask[name] & (
                    param.data.abs() >= threshold
                )

    return results, mask

results は、この記事で使用するスパース性対精度曲線の原材料です。 eval_fn は、完全にサンプル外で、パージされたスプリット上にある必要があります。サンプル内でスコア付けされた IMP 実行は、何の意味もない美しい曲線を報告します。

測定

レイテンシーは、ブログの残りの部分と同じハーネス規則 (ウォームアップを除外、ベストオブ N、平均ではなく報告された p50/p95/p99) で測定されており、そのプロトコルとコードは Polars vs pandas.剪定のポイントは次の3つです。

• マスクされたモデルではなく、再構築 モデルをベンチマークします。バッチ サイズ 1 のマスクされたモデルは、密な形状を測定します。
• バッチサイズを報告します。バッチ 1 (クォート ループ) とバッチ 256 (リサーチ スイープ) はメモリ依存/計算依存の境界線の異なる側に位置しており、プルーニングはそれらを異なる方法で支援します。
• ラベル定義を記載して、同じスプリット、同じ水平線での精度をレポートします。一致精度列のないレイテンシー テーブルは、モデルを完全に削除するための引数となります。

高度なテクニック

知識の蒸留による枝刈り

枝刈りや微調整を個別に行うのではなく、元の密なモデルを教師として使用します。枝刈りされた生徒は、タスクの損失と教師の出力分布からの KL の乖離の組み合わせを最小限に抑えます。

ここで、$\mathbf{z}_t$ と $\mathbf{z}_s$ は教師と生徒のロジット、$T$ は温度、$\alpha$ は目標のバランスを表します。 $T^2$ 係数は蒸留勾配を再スケールし、それ以外の場合は $1/T^2$ として縮小します。

def distillation_loss(student_logits, teacher_logits, labels,
                      temperature=3.0, alpha=0.5):
    """Combined task + distillation loss."""
    task_loss = nn.CrossEntropyLoss()(student_logits, labels)

    soft_student = nn.functional.log_softmax(
        student_logits / temperature, dim=-1
    )
    soft_teacher = nn.functional.softmax(
        teacher_logits / temperature, dim=-1
    )
    kd_loss = nn.functional.kl_div(
        soft_student, soft_teacher, reduction='batchmean'
    )

    return (1 - alpha) * task_loss + alpha * (temperature ** 2) * kd_loss

動作剪定

絶対的な大きさによる枝刈りではなく、動き枝刈り (Sanh et al., 2020) では、トレーニング中にゼロに向かって移動する重みが枝刈りされます。重要度スコアは、勾配と重みの積を累積します。

負のスコアを持つ重みは削除されます。マグニチュードの枝刈りに対する議論は、特に微調整に関するものです。事前トレーニングされたモデルを適応させる場合、マグニチュード分布は事前トレーニング タスクによって形成されるため、マグニチュードは古くなった重要な信号であり、進行方向はより新しい信号です。ローリング ウィンドウで再トレーニングされた取引モデルの場合、これは最初からトレーニングするよりも一般的な状況です。

NVIDIA 2:4 構造化スパーシティ

アンペア以降の NVIDIA GPU は、ハードウェアで 2:4 構造化スパース性をサポートしています。連続する 4 つの重みのうち、正確に 2 つがゼロでなければなりません。

これは、ハードウェアが実際に報いるきめ細かいスパース性の 1 つの形式であり、非構造化プルーニングによる 90% ゼロの数値よりも重要であるのはそのためです。この制約はグローバルではなくローカルです。4 つのうち「どの」2 つが生き残るかは関係ありません。したがって、グローバル マスクを固定するよりもはるかに弱い制限ですが、提供されるスパース レベルは 50% だけです。

from torch.ao.pruning import WeightNormSparsifier

sparsifier = WeightNormSparsifier(
    sparsity_level=0.5,
    sparse_block_shape=(1, 4),
    zeros_per_block=2,
)

sparsifier.prepare(
    model, config=[{"tensor_fqn": "network.0.weight"}]
)
sparsifier.step()
sparsifier.squash_mask()

高速化を実現するには、スパース tensor コア (ONNX エクスポートと TensorRT ビルド、または torch.sparse.to_sparse_semi_structured) を使用する推論パスが必要です。高密度ランタイムを通じて 2:4 マスクされたモデルをエクスポートすると、精度が犠牲になりますが、メリットはまったくありません。

本番環境への展開

検証

プルーニングされたモデルは、圧縮された古いモデルではなく、新しいモデルであり、他の候補と同じ受け入れゲートを通過します。つまり、ウォークフォワード最適化 によるローリング再トレーニングとサンプル外再検証と、収縮されたシャープ比 ] による選択効果補正です。この修正はここではオプションではありません。IMP は候補モデルの シーケンス を生成するため、10 回のラウンドにわたって最もよく見えるスパース性レベルが検索で選択され、そのシャープは有効な試行回数によって収縮する必要があります。 「シャープが 5% を超えて下落した場合は拒否する」のような一律のルールは、その算術に耐えられないため、この記事ではそのルールが見つかりません。

量子化スタッキング

プルーニングは量子化で構成されます。 90% スパースで INT8 に量子化されたモデルの圧縮率は次のとおりです。

48MBモデルは1.2MBとなります。それは保管上の要求であり、それ以上のものではありません。 1.2 MB モデルが同じ決定を下すかどうかは、それ自身の答えを持つ別の質問であり、GPU 精度のトラップ がそれを尋ねる理由です。このブログでは、fp32 だけがバックテスト計算で 211 の相対誤差を生成することが示されており、完全に合理的であるように見えました。 INT8 はそれよりもはるかに積極的な削減です。 fp32 密モデルに対する定量化されたパリティ ゲートの背後でのみ、量子化および枝刈りされたモデルを出荷します。これは、保証ではなく、ホールドアウト期間の意思決定一致率と損益デルタです。

＃＃＃ 監視

プルーニングされたモデルは、分布の変化の影響をより受けやすくなります。見る価値があります:

• 活性化の希薄性: 生き残ったニューロンがほとんどゼロを発する場合、有効なモデルは意図したものよりも小さく、おそらく劣化しています。
• 再トレーニング中の勾配ノルム: 勾配の爆発は、生き残ったサブネットワークが削除されたものをあまりにも積極的に補償するように求められていることを示唆しています。
• 予測エントロピー: ノイズの多い微細構造データに対して過信する枝刈りモデルは、トレーニング レジームに適合している可能性があります。

＃＃ 結論

この方法は十分に確立されており、スイープが実行されるまでは、この記事で主張するのはこれだけです。非構造化枝刈りでは数値がまばらになり、速度が得られません。構造化プルーニングにより、レイヤーをマスクするのではなく再構築する場合に限り、速度が向上します。宝くじの仮説は、コンパクトなモデルがオーバーパラメータ化されたモデルの内部にすでに存在していることを示唆していますが、それは注文帳データではなく画像ベンチマークで実証されており、上で挙げた市場データで機能する「はずである」3つの理由は、実験が添付された仮説であり、結果ではありません。

文献から得られる実践的なヒューリスティックは、最初から小さく設計するのではなく、大きくトレーニングして枝刈りすることです。大規模なモデルは損失状況をより効果的に探索し、枝刈りによって重要な経路が保存されます。それが取引モデルに当てはまるかどうか、どの程度のスパース性で、どのような精度のコストがかかるかは、IMP スイープを 1 回行うだけで決まります。この記事は、そのスイープ後に数値を含めてもう一度読む必要があります。