Model Pruning for Low-Latency Trading Inference

Bài DeepLOB của chúng tôi kết thúc phần triển khai bằng ba gạch đầu dòng — ONNX cùng TensorRT, lượng tử hóa INT8, FPGA — nhưng không xử lý cụ thể mục nào. Bài này bổ sung phần còn thiếu về vấn đề đầu tiên nằm bên dưới cả ba: mô hình lớn hơn mức cần thiết. Pruning mạng nơ-ron loại bỏ các tham số dư thừa, và tuyên bố đáng chú ý trong tài liệu không phải là nó tiết kiệm bộ nhớ, mà là một subnetwork giữ 10-20% trọng số có thể đạt accuracy tương đương mô hình dense.

Latency quan trọng đến mức nào là điều blog đã lập luận — ZigBolt trên đường truyền messaging và thuế IPC với phép tính break-even — còn spread modeling đã sở hữu trade-off nhanh nhưng kém hơn một chút so với chậm nhưng tốt hơn, cùng bảng gradient boosting so với deep learning có hàng inference-latency. Không bài nào trong đó nói cách làm một mô hình cụ thể nhỏ hơn. Trong các giai đoạn của quoting loop, model inference là thứ hoàn toàn nằm trong quyền kiểm soát của chúng ta; các chặng transport được đề cập cùng số p50/p95/p99 có thể tái lập trong data communication cho algotrading.

Bài này là: toán học và code hoạt động cho magnitude pruning, structured pruning, Iterative Magnitude Pruning, movement pruning, knowledge distillation và sparsity bán cấu trúc NVIDIA 2:4, áp dụng cho một trading MLP.

Bài này không phải: một kết quả đã đo. Mỗi bài thực nghiệm trên blog đều có dòng provenance hoặc repo đi kèm, còn bài này hiện chưa có cả hai. Đường cong sparsity-so-với-accuracy-so-với-latency dưới đây được nêu như một experiment cần chạy, không phải bảng để trích dẫn. Hãy xem mọi thứ ở đây là phương pháp, còn các con số vẫn đang chờ.

Pruning mang lại điều gì

Ràng buộc nằm ở kích thước. Hãy xét một mô hình tần suất trung bình — MLP 4 lớp với 2048 hidden units trên các feature của sổ lệnh:

Với $d_{\text{in}} = 100$, $h = 2048$, $L = 4$, $d_{\text{out}} = 3$, có khoảng 12,6 triệu tham số — khoảng 48 MB ở float32. L2 thường chỉ 1-4 MB, nên trọng số không vừa; chúng được stream từ vùng bộ nhớ xa hơn trong mỗi forward pass. Prune 95% thì còn khoảng 630K tham số hiệu dụng và 2,4 MB, vừa với L2.

Việc đó có chuyển thành thời gian wall-clock hay không phụ thuộc kernel có bị giới hạn bởi bộ nhớ hay không, và đây là câu hỏi về cường độ tính toán chứ không phải kích thước. Thang tốc độ backtest engine trình bày mô hình roofline (Williams, Waterman & Patterson) qua một ví dụ đã đo thay vì khẳng định hệ số phạt; cùng cách nhìn đó áp dụng ở đây, và kỷ luật cũng vậy: hãy đo số byte di chuyển trước khi tuyên bố speedup.

Nền tảng của Pruning

Pruning không cấu trúc

Cách đơn giản nhất: đặt từng trọng số về zero dựa trên magnitude. Với ma trận trọng số $\mathbf{W} \in \mathbb{R}^{m \times n}$, tạo mask nhị phân $\mathbf{M} \in \{0, 1\}^{m \times n}$ sao cho:

trong đó $\tau$ là ngưỡng được chọn để đạt mức sparsity $s$ mong muốn:

Ma trận sau pruning là $\mathbf{W}' = \mathbf{W} \odot \mathbf{M}$, với $\odot$ là tích Hadamard. Trực giác là các trọng số gần zero đóng góp rất ít vào output của layer.

Nêu thẳng vấn đề, vì con số sparsity rất dễ bị đọc sai: sparsity không cấu trúc không tự chuyển thành speedup trên phần cứng tiêu chuẩn. Ma trận có 90% zero vẫn thực hiện cùng số phép multiply-accumulate trừ khi bạn chuyển sang sparse kernel hoặc phần cứng hỗ trợ sparsity. Khi code dưới đây in Sparsity: 90.0%, đó là số lượng zero — không phải 10x dưới bất kỳ nghĩa nào, và trên dense CPU GEMM cũng không phải 1.01x. Những đường thực sự tiết kiệm thời gian là structured pruning (ma trận nhỏ hơn) và sparsity bán cấu trúc 2:4 (được phần cứng hỗ trợ), cả hai được trình bày bên dưới.

Pruning có cấu trúc

Pruning có cấu trúc loại bỏ toàn bộ neuron, channel hoặc attention head. Với linear layer $\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$ và $\mathbf{W} \in \mathbb{R}^{m \times n}$, loại neuron $j$ sẽ đưa row thứ $j$ của $\mathbf{W}$ và phần tử thứ $j$ của $\mathbf{b}$ về zero:

Neuron có chuẩn $\ell_2$ nhỏ nhất sẽ bị loại trước. Đây là biến thể tạo ra ma trận thực sự nhỏ hơn — nhưng chỉ khi bạn thật sự rebuild các layer. Đặt các row về zero rồi giữ tensor ở shape ban đầu không thay đổi FLOP count; bước rebuild trong phần triển khai mới biến mask thành ma trận $512 \times n$.

Với layer tích chập, analog là filter pruning. Cho $\mathbf{W} \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times k \times k}$, importance của output filter $i$ là:

Loại filter $i$ sẽ loại bỏ toàn bộ output channel, giảm FLOPs theo tỷ lệ.

Giả thuyết Lottery Ticket

Năm 2019, Frankle và Carbin đưa ra Lottery Ticket Hypothesis (LTH): trong một mạng dense khởi tạo ngẫu nhiên tồn tại một subnetwork thưa — một "winning ticket" — khi được train từ initialization ban đầu có thể đạt accuracy của toàn mạng trong số iteration tương đương.

Về hình thức, xét $f(\mathbf{x}; \boldsymbol{\theta}_0)$ được khởi tạo với $\boldsymbol{\theta}_0 \sim \mathcal{D}_\theta$. Sau khi train đến hội tụ, ta có $\boldsymbol{\theta}^*$ và suy ra pruning mask $\mathbf{m}$. LTH phát biểu rằng tồn tại $\mathbf{m}$ sao cho:

với $\|\mathbf{m}\|_0 \ll |\boldsymbol{\theta}_0|$. Các experiment ban đầu thực hiện trên MNIST và CIFAR-10, nơi winning ticket giữ lại 10-20% tham số. Không có lý do để mặc định điều đó chuyển sang dữ liệu sổ lệnh — feature LOB không dừng và label gần như noise, khác với phân loại ảnh đúng ở những điểm có thể quan trọng.

Iterative Magnitude Pruning (IMP)

Winning ticket được tìm bằng IMP:

1. Khởi tạo mạng với $\boldsymbol{\theta}_0$.
2. Train đến hội tụ, thu được $\boldsymbol{\theta}^*$.
3. Prune $p\%$ trọng số có magnitude nhỏ nhất, tạo mask $\mathbf{m}$.
4. Reset các trọng số còn sống về giá trị trong $\boldsymbol{\theta}_0$ (rewinding).
5. Lặp lại từ bước 2 với mạng đã mask.

Mỗi round prune một phần $p$ (thường là 20%), nên sau $n$ round còn $(1-p)^n$ tham số. Sau 10 round với $p = 0.2$, còn lại khoảng 10,7%.

Ba giả thuyết về mô hình giao dịch, chưa giả thuyết nào được kiểm thử

Dễ muốn lập luận rằng LTH phải đặc biệt hiệu quả trên dữ liệu thị trường. Có ba lập luận như vậy; cả ba đều là giả thuyết, và trình bày chúng như sự thật chính là failure mode mà blog này muốn tránh.

1. Tín hiệu tài chính thưa. Phần lớn snapshot sổ lệnh là noise, nên subnetwork thưa có thể tự nhiên phù hợp với tín hiệu thưa. Có thể kiểm thử: so sánh IMP với random mask cùng sparsity; nếu sparsity mới là thứ tạo ra kết quả, random mask không nên kém quá xa.
2. Winning ticket tổng quát hóa qua các regime. Đây là một claim thực nghiệm về thị trường nhưng không có citation, và là điều thú vị nhất trong ba giả thuyết. Có thể test trực tiếp với nhãn regime từ regime detection bằng HMM: tìm ticket trong regime A, retrain ở regime B và so sánh với ticket được tìm trực tiếp trong B.
3. Sparsity regularize. Capacity hiệu dụng thấp hơn có thể giảm việc fit noise microstructure — thể hiện ở khoảng cách out-of-sample của model đã prune nhỏ hơn model dense, chứ không chỉ tương đương.

Triển khai: Prune một Trading MLP

Mô hình cơ sở

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():,}")

Pruning magnitude không cấu trúc

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%}")

Pruning có cấu trúc, với bước rebuild khiến nó thực sự hiệu quả

Mask row mới chỉ là một nửa công việc. Nửa tạo ra speedup là rebuild từng layer về shape đã giảm — nghĩa là truyền việc loại bỏ về phía trước: bỏ row $j$ của layer $i$ cũng bỏ column $j$ của layer $i+1$ và channel $j$ của BatchNorm1d nằm giữa chúng.

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

Trước khi tin vào điều này, hãy kiểm tra hai điều theo cùng tinh thần với các equivalence gate mà phần còn lại của blog sử dụng:

• Shape. rebuilt phải có hidden dimension ở $(1 - \text{fraction}) \cdot h$ — 512 khi fraction=0.75, $h = 2048$ — và parameter count phải giảm theo bình phương, vì cả hai dimension của ma trận bên trong đều nhỏ đi.
• Output. Ở chế độ eval(), output của rebuilt(x) phải khớp output không có rebuilt của model đã mask trong sai số floating-point trên cùng batch. Nếu không, việc truyền column sai, và mọi con số downstream đang đo một model khác với bạn nghĩ.

Bài test row-survival giả định row đã mask bằng zero tuyệt đối còn row sống thì không. Điều này đúng với output của ln_structured; nếu quy trình khác tạo ra neuron sống thực sự toàn zero thì không đúng, vì vậy hãy assert số survivor theo fraction yêu cầu thay vì mù quáng tin vào norm test.

Iterative Magnitude Pruning (Lottery Ticket Search)

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 là nguyên liệu thô cho đường cong sparsity-so-với-accuracy mà bài viết này nợ bạn. eval_fn phải thực sự out-of-sample trên các split đã purge — một IMP run được chấm in-sample sẽ tạo ra đường cong đẹp nhưng vô nghĩa.

Đo lường

Đo latency theo cùng quy ước harness của blog — bỏ warmup, best-of-N, báo cáo p50/p95/p99 thay vì mean — và protocol này cùng code nằm trong Polars và pandas. Ba điểm riêng cho pruning:

• Benchmark model đã rebuild, không phải model đã mask. Model mask ở batch size 1 đo shape dense.
• Báo cáo batch size. Batch 1 (quoting loop) và batch 256 (research sweep) nằm ở hai phía khác nhau của ranh giới memory-bound/compute-bound, và pruning tác động khác nhau lên chúng.
• Báo cáo accuracy trên cùng split, cùng horizon, với định nghĩa label được nêu rõ. Bảng latency không có cột accuracy tương ứng là lập luận để xóa toàn bộ model.

Kỹ thuật nâng cao

Pruning kết hợp Knowledge Distillation

Thay vì prune và fine-tune riêng rẽ, hãy dùng model dense ban đầu làm teacher. Student đã prune tối thiểu hóa tổ hợp giữa task loss và KL divergence của phân phối output từ teacher:

trong đó $\mathbf{z}_t$ và $\mathbf{z}_s$ là logits của teacher và student, $T$ là temperature, còn $\alpha$ cân bằng các mục tiêu. Hệ số $T^2$ rescale gradient distillation, vốn sẽ giảm theo $1/T^2$ nếu thiếu nó.

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

Movement Pruning

Thay vì prune theo magnitude tuyệt đối, movement pruning (Sanh et al., 2020) loại các trọng số đang di chuyển về phía zero trong quá trình train. Importance score tích lũy tích gradient-trọng số:

Các trọng số có score âm sẽ bị prune. Lập luận ủng hộ nó so với magnitude pruning đặc biệt liên quan đến fine-tuning: khi adapt một model pre-trained, phân phối magnitude đã bị định hình bởi task pre-training, nên magnitude là tín hiệu importance lỗi thời còn hướng di chuyển mới hơn. Với model giao dịch retrain trên các rolling window, đây là tình huống phổ biến hơn train từ đầu.

Sparsity có cấu trúc NVIDIA 2:4

GPU NVIDIA từ Ampere trở đi hỗ trợ sparsity có cấu trúc 2:4 trong phần cứng: cứ mỗi 4 trọng số liên tiếp phải có đúng 2 trọng số zero.

Đây là dạng sparsity fine-grained duy nhất mà phần cứng thực sự tận dụng, nên quan trọng hơn con số 90%-zero của pruning không cấu trúc. Ràng buộc mang tính local chứ không global — nó không quan tâm hai trọng số nào trong bốn trọng số còn sống — nên yếu hơn nhiều so với việc cố định global mask, dù mức sparsity duy nhất có thể chọn là 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()

Để hiện thực hóa speedup, inference path phải dùng sparse tensor cores — export ONNX cộng build TensorRT, hoặc torch.sparse.to_sparse_semi_structured. Export model có mask 2:4 qua runtime dense chỉ đem lại chi phí accuracy mà không có lợi ích.

Triển khai production

Xác thực

Model đã prune là một model mới, không phải model cũ được nén, và phải đi qua acceptance gate giống mọi candidate khác: rolling retraining và revalidation out-of-sample theo walk-forward optimization, cùng selection-effect correction từ deflated Sharpe ratio. Correction đó không tùy chọn ở đây — IMP tạo ra một sequence các model candidate, nên mức sparsity trông tốt nhất sau mười round đã được chọn trong quá trình search, và Sharpe của nó cần được deflate theo số trial hiệu dụng. Một rule phẳng như "reject nếu Sharpe giảm hơn 5%" không qua được phép tính đó, nên bạn sẽ không thấy rule như vậy trong bài.

Xếp chồng Quantization

Pruning kết hợp được với quantization. Model sparsity 90% và quantize INT8 có compression ratio là:

Model 48 MB trở thành 1,2 MB. Đây là một claim về storage và không hơn. Việc model 1,2 MB có tạo ra cùng quyết định hay không là câu hỏi riêng với câu trả lời riêng, và GPU precision trap là lý do cần hỏi: trên blog này, riêng fp32 đã cho relative error 211 trong một phép tính backtest trông hoàn toàn hợp lý. INT8 là mức giảm mạnh hơn nhiều. Chỉ ship model đã quantize và prune sau một parity gate định lượng so với model dense fp32 — decision agreement rate và PnL delta trên giai đoạn hold-out, không phải lời đảm bảo.

Giám sát

Model đã prune có thể nhạy hơn với distribution shift. Cần theo dõi:

• Activation sparsity: nếu các neuron còn sống phát ra chủ yếu zero, model hiệu dụng nhỏ hơn dự kiến và có thể đang suy giảm.
• Gradient norm khi retrain: gradient bùng nổ cho thấy subnetwork còn sống đang bị yêu cầu bù quá mạnh cho phần đã loại bỏ.
• Prediction entropy: model đã prune trở nên quá tự tin trên dữ liệu microstructure nhiễu có thể đang fit regime huấn luyện.

Kết luận

Các phương pháp đã được thiết lập tốt và cho đến khi sweep chạy xong, đó là toàn bộ điều bài viết này khẳng định. Pruning không cấu trúc cho bạn một con số sparsity nhưng không cho tốc độ. Pruning có cấu trúc cho tốc độ nếu — và chỉ nếu — bạn rebuild các layer thay vì mask chúng. Lottery Ticket Hypothesis gợi ý rằng model nhỏ gọn đã tồn tại bên trong model overparameterized, dù điều này mới được chứng minh trên benchmark ảnh chứ chưa trên dữ liệu sổ lệnh, và ba lý do cho rằng nó "phải" hoạt động trên dữ liệu thị trường ở trên đều là giả thuyết gắn với experiment, không phải phát hiện.

Heuristic thực tế từ tài liệu là train lớn rồi prune xuống thay vì thiết kế nhỏ ngay từ đầu: model lớn khám phá loss landscape hiệu quả hơn, còn pruning giữ lại các pathway quan trọng. Điều đó có đúng với model giao dịch hay không, ở mức sparsity nào và với chi phí accuracy bao nhiêu, chỉ còn cách một IMP sweep — và nên đọc lại bài viết này sau sweep đó với các con số cụ thể.