Model Pruning for Low-Latency Trading Inference

हमारे DeepLOB लेख का deployment section तीन बिंदुओं — ONNX plus TensorRT, INT8 quantization, FPGA — पर समाप्त होता है, लेकिन इनमें से किसी को विस्तार से नहीं समझाता। यह लेख उन तीनों के नीचे मौजूद पहली समस्या का छूटा हुआ उपचार है: मॉडल आवश्यकता से बड़ा है। Neural network pruning अनावश्यक parameters हटाती है, और साहित्य का रोचक दावा यह नहीं कि इससे memory बचती है, बल्कि यह कि weights के 10-20% वाले subnetwork की accuracy dense model से मेल खा सकती है।

यहाँ latency का महत्व है — ब्लॉग इसे पहले ही ZigBolt में messaging path पर और IPC tax में break-even arithmetic के साथ दिखा चुका है — और spread modeling fast-but-slightly-worse बनाम slow-but-better trade-off को पहले ही संभालता है, जिसमें inference-latency row वाली gradient-boosting-versus-deep-learning table भी है। इनमें से कोई यह नहीं बताता कि किसी दिए हुए मॉडल को कैसे छोटा बनाया जाए। Quoting loop के stages में model inference ही पूरी तरह हमारे नियंत्रण में है; transport legs को reproducible p50/p95/p99 numbers के साथ data communication for algotrading में कवर किया गया है।

यह लेख क्या है: magnitude pruning, structured pruning, Iterative Magnitude Pruning, movement pruning, knowledge distillation और NVIDIA 2:4 semi-structured sparsity का गणित और working code, trading MLP पर लागू।

यह क्या नहीं है: measured result नहीं। इस ब्लॉग का हर empirical article provenance line या companion repo रखता है, और इस लेख में अभी दोनों नहीं हैं। Sparsity-versus-accuracy-versus-latency curve नीचे चलाए जाने वाले experiment के रूप में दी गई है, quote करने के लिए table के रूप में नहीं। इसे method समझें और numbers को pending मानें।

Pruning से आपको क्या मिलता है

Constraint size का है। एक mid-frequency model पर विचार करें — order-book features पर 2048 hidden units वाला 4-layer MLP:

$d_{\text{in}} = 100$, $h = 2048$, $L = 4$, $d_{\text{out}} = 3$ के लिए लगभग 12.6 million parameters होते हैं — float32 में करीब 48 MB। L2 आम तौर पर 1-4 MB होता है, इसलिए weights फिट नहीं होते; हर forward pass पर उन्हें बाहर की अधिक दूर की memory से stream किया जाता है। उनमें से 95% prune करें तो लगभग 630K effective parameters और 2.4 MB बचते हैं, जो फिट हो जाता है।

क्या यह wall-clock time में बदलता है, यह इस पर निर्भर करता है कि kernel memory-bound है या नहीं, और यह size का नहीं बल्कि arithmetic-intensity का प्रश्न है। Backtest engine speed ladder roofline model (Williams, Waterman & Patterson) को किसी penalty factor का दावा करने के बजाय measured example से समझाता है; वही framing यहाँ लागू होती है, और वही discipline भी: speedup का दावा करने से पहले moved bytes मापें।

Pruning की बुनियाद

Unstructured Pruning

सबसे सरल तरीका: individual weights को उनके magnitude के आधार पर zero कर दें। Weight matrix $\mathbf{W} \in \mathbb{R}^{m \times n}$ के लिए binary mask $\mathbf{M} \in \{0, 1\}^{m \times n}$ बनाएँ, ताकि:

जहाँ $\tau$ desired sparsity level $s$ हासिल करने के लिए चुना गया threshold है:

Pruned matrix $\mathbf{W}' = \mathbf{W} \odot \mathbf{M}$ है, जहाँ $\odot$ Hadamard product है। Intuition यह है कि zero के पास वाले weights layer के output में बहुत कम योगदान देते हैं।

समस्या को सीधे कहें, क्योंकि sparsity number को गलत पढ़ना आसान है: unstructured sparsity standard hardware पर speedup में नहीं बदलती। 90% zeros वाली matrix भी उतने ही multiply-accumulates issue करती है, जब तक sparse kernels या sparsity-support वाले hardware पर switch न करें। जब नीचे का code Sparsity: 90.0% print करता है, वह zeros की गिनती है — यह किसी भी तरह का 10x नहीं है, और dense CPU GEMM पर 1.01x भी नहीं। जो paths वास्तव में समय बचाते हैं वे structured pruning (छोटी matrices) और 2:4 semi-structured sparsity (hardware support) हैं, दोनों नीचे दिए गए हैं।

Structured Pruning

Structured pruning पूरे neurons, channels या attention heads हटाती है। Linear layer $\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$ के लिए, जहाँ $\mathbf{W} \in \mathbb{R}^{m \times n}$ है, neuron $j$ हटाने पर $\mathbf{W}$ की $j$-वीं row और $\mathbf{b}$ का $j$-वाँ element zero हो जाता है:

सबसे छोटे $\ell_2$-norm वाले neurons पहले जाते हैं। यही वह variant है जो सचमुच छोटी matrices बनाता है — लेकिन केवल तब जब आप layers को वास्तव में rebuild करें। Rows को zero करके tensor को original shape पर छोड़ देने से FLOP count में कुछ नहीं बदलता; implementation section का rebuild step ही mask को $512 \times n$ matrix में बदलता है।

Convolutional layers में इसका analog filter pruning है। $\mathbf{W} \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times k \times k}$ दिए जाने पर output filter $i$ का importance है:

Filter $i$ हटाने से पूरा output channel हट जाता है और FLOPs proportionally कम हो जाते हैं।

Lottery Ticket Hypothesis

2019 में Frankle और Carbin ने Lottery Ticket Hypothesis (LTH) पेश की: randomly initialized dense network के भीतर एक sparse subnetwork — एक "winning ticket" — मौजूद है, जो original initialization से train किए जाने पर comparable iterations में पूरे network की accuracy से मेल खाता है।

औपचारिक रूप से, $f(\mathbf{x}; \boldsymbol{\theta}_0)$ को $\boldsymbol{\theta}_0 \sim \mathcal{D}_\theta$ के साथ initialized मानें। Convergence तक training के बाद हमें $\boldsymbol{\theta}^*$ मिलता है और pruning mask $\mathbf{m}$ निकाला जाता है। LTH कहती है कि कोई $\mathbf{m}$ ऐसा है कि:

जहाँ $\|\mathbf{m}\|_0 \ll |\boldsymbol{\theta}_0|$ है। मूल experiments MNIST और CIFAR-10 पर थे, जहाँ winning tickets ने parameters के 10-20% बनाए रखे। Order-book data पर यह अपने-आप लागू होगा, ऐसा मानने का कोई आधार नहीं — LOB features non-stationary हैं और label लगभग noise है, जो image classification से ठीक उन तरीकों में अलग regime है जो मायने रख सकते हैं।

Iterative Magnitude Pruning (IMP)

Ticket IMP से मिलता है:

1. Network को $\boldsymbol{\theta}_0$ से initialize करें।
2. Convergence तक train करें और $\boldsymbol{\theta}^*$ प्राप्त करें।
3. सबसे छोटे magnitude वाले $p\%$ weights prune करें और mask $\mathbf{m}$ बनाएँ।
4. Surviving weights को $\boldsymbol{\theta}_0$ की values पर reset करें (rewinding)।
5. Masked network के साथ step 2 से दोहराएँ।

हर round fraction $p$ (आम तौर पर 20%) prune करता है, इसलिए $n$ rounds के बाद parameters का $(1-p)^n$ हिस्सा बचता है। $p = 0.2$ पर 10 rounds के बाद लगभग 10.7% बचते हैं।

Trading Models के बारे में तीन Hypotheses, जिनमें से किसी का परीक्षण नहीं हुआ

यह कहना आकर्षक है कि LTH market data पर विशेष रूप से अच्छा काम करना चाहिए। ऐसे तीन arguments सामने आते हैं; तीनों hypotheses हैं, और उन्हें facts की तरह बताना वही failure mode है जिससे यह ब्लॉग बचना चाहता है।

1. Financial signals sparse हैं। Order-book snapshot का अधिकांश भाग noise है, इसलिए sparse subnetwork स्वाभाविक रूप से sparse signal के साथ aligned हो सकता है। Testable: IMP की तुलना same-sparsity random mask से करें; यदि काम sparsity कर रही है, तो random mask बहुत पीछे नहीं होना चाहिए।
2. Winning tickets regimes के पार generalize करते हैं। यह markets के बारे में empirical claim है, जिसके पीछे कोई citation नहीं, और तीनों में सबसे interesting है। इसे regime detection with HMMs के regime labels के विरुद्ध सीधे test किया जा सकता है: regime A में ticket खोजें, regime B में retrain करें और B में native रूप से मिले ticket से तुलना करें।
3. Sparsity regularize करती है। Lower effective capacity microstructure noise पर fitting कम कर सकती है — इसका संकेत यह होगा कि pruned model का out-of-sample gap dense model से छोटा हो, केवल comparable नहीं।

Implementation: Trading MLP को Prune करना

Base Model

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

Unstructured Magnitude Pruning

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

Structured Pruning, With the Rebuild That Makes It Real

Rows को mask करना आधा काम है। Speedup देने वाला आधा काम हर layer को उसके reduced shape पर rebuild करना है — यानी removal को आगे propagate करना: layer $i$ की row $j$ हटाने पर layer $i+1$ का column $j$ और बीच के किसी BatchNorm1d का channel $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

इस पर भरोसा करने से पहले, ब्लॉग के बाकी equivalence gates की तरह दो बातें जाँचें:

• Shapes. rebuilt में hidden dimensions $(1 - \text{fraction}) \cdot h$ पर दिखने चाहिए — fraction=0.75 के लिए 512, $h = 2048$ — और parameter count quadratic रूप से गिरना चाहिए, क्योंकि interior matrices के दोनों dimensions छोटे हो जाते हैं।
• Outputs. eval() mode में rebuilt(x) का output masked model के rebuilt-free output से उसी batch पर floating-point tolerance तक match होना चाहिए। यदि नहीं होता, तो column propagation गलत है, और हर downstream number उस model को माप रहा है जो आपको लगता है उससे अलग है।

Row-survival test मानता है कि masked row बिल्कुल zero है और live row zero नहीं है। ln_structured output के लिए यह सही है; यदि कोई अन्य procedure सचमुच all-zero live neuron बनाए, तो ऐसा नहीं होगा। इसलिए norm test पर अंधा भरोसा करने के बजाय requested fraction के विरुद्ध survivor count assert करें।

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 इस लेख की sparsity-versus-accuracy curve के लिए raw material है। eval_fn को purged splits पर सचमुच out-of-sample होना चाहिए — in-sample score किया गया IMP run एक सुंदर curve देगा जिसका कोई अर्थ नहीं।

इसे मापना

Latency को ब्लॉग के बाकी हिस्से जैसी harness convention से मापें — warmup को बाहर रखकर, best-of-N, mean के बजाय p50/p95/p99 रिपोर्ट करें — और यह protocol code के साथ Polars vs pandas में है। Pruning के लिए तीन विशेष बिंदु:

• Masked model के बजाय rebuilt model को benchmark करें। Batch size 1 पर masked model dense shape को मापता है।
• Batch size रिपोर्ट करें। Batch 1 (quoting loop) और batch 256 (research sweep) memory-bound/compute-bound line के अलग-अलग sides पर होते हैं और pruning उन पर अलग असर डालती है।
• उसी split, उसी horizon और stated label definition पर accuracy रिपोर्ट करें। Matching accuracy column के बिना latency table पूरे model को हटाने का argument है।

Advanced Techniques

Knowledge Distillation से Pruning

Pruning और fine-tuning को अलग-अलग करने के बजाय original dense model को teacher की तरह उपयोग करें। Pruned student task loss और teacher के output distribution से KL divergence के संयोजन को minimize करता है:

जहाँ $\mathbf{z}_t$ और $\mathbf{z}_s$ teacher और student logits हैं, $T$ temperature है और $\alpha$ objectives के बीच balance करती है। $T^2$ factor distillation gradients को rescale करता है, जो इसके बिना $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

Movement Pruning

Absolute magnitude से prune करने के बजाय movement pruning (Sanh et al., 2020) उन weights को prune करती है जो training के दौरान zero की ओर बढ़ रहे होते हैं। Importance score gradient-weight product को accumulate करता है:

Negative scores वाले weights prune होते हैं। Magnitude pruning पर इसका argument खास तौर पर fine-tuning के बारे में है: pre-trained model को adapt करते समय magnitude distribution pre-training task से आकार ले चुकी होती है, इसलिए magnitude stale importance signal है और direction-of-travel अधिक ताज़ा है। Rolling windows पर retrain होने वाले trading model के लिए यह scratch से training की तुलना में अधिक सामान्य स्थिति है।

NVIDIA 2:4 Structured Sparsity

Ampere और उसके बाद के NVIDIA GPUs hardware में 2:4 structured sparsity support करते हैं: हर 4 contiguous weights में ठीक 2 zero होने चाहिए।

यह fine-grained sparsity का वह एक रूप है जिसे hardware वास्तव में reward करता है, इसलिए unstructured pruning के 90%-zeros number से अधिक महत्वपूर्ण है। Constraint local है, global नहीं — उसे परवाह नहीं कि हर चार में से कौन-से दो बचते हैं — इसलिए global mask fix करने की तुलना में यह बहुत कमजोर restriction है, हालांकि उपलब्ध sparsity level केवल 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()

Speedup को साकार करने के लिए inference path को sparse tensor cores का उपयोग करना चाहिए — ONNX export plus TensorRT build, या torch.sparse.to_sparse_semi_structured। 2:4-masked model को dense runtime से export करने पर accuracy cost मिलती है और कोई लाभ नहीं।

Production Deployment

Validation

Pruned model compressed पुराना model नहीं, एक नया model है, और इसे किसी भी अन्य candidate जैसी acceptance gate से गुजरना होगा: walk-forward optimization के अनुसार rolling retraining और out-of-sample revalidation, साथ में deflated Sharpe ratio का selection-effect correction। यह correction यहाँ optional नहीं है — IMP candidate models का एक sequence बनाता है, इसलिए दस rounds में सबसे अच्छा दिखने वाला sparsity level search के तहत चुना गया और उसके Sharpe को effective number of trials से deflate करना होगा। "Sharpe drops more than 5% पर reject" जैसा flat rule इस arithmetic से नहीं बचता, इसलिए इस article में आपको ऐसा rule नहीं मिलेगा।

Quantization Stacking

Pruning quantization के साथ compose होती है। INT8 में quantized और 90% sparse model का compression ratio है:

48 MB का model 1.2 MB हो जाता है। यह storage claim है और इससे अधिक कुछ नहीं। क्या 1.2 MB model वही decisions produce करता है, यह अलग प्रश्न है, जिसका अलग उत्तर चाहिए, और GPU precision trap यही प्रश्न पूछने का कारण है: इस ब्लॉग में fp32 अकेले एक backtest computation में 211 का relative error पैदा कर चुका है जो पूरी तरह reasonable दिखता था। INT8 उससे कहीं अधिक aggressive reduction है। Quantized-and-pruned model को fp32 dense model के विरुद्ध quantified parity gate के पीछे ही ship करें — held-out period पर decision agreement rate और PnL delta, assurance नहीं।

Monitoring

Pruned models distribution shift के प्रति अधिक sensitive हो सकते हैं। इन चीज़ों पर नज़र रखें:

• Activation sparsity: यदि surviving neurons अधिकांशतः zeros emit करते हैं, तो effective model intended से छोटा और संभवतः degrading है।
• Retraining के दौरान gradient norms: exploding gradients संकेत देते हैं कि surviving subnetwork से हटाई गई चीज़ों की भरपाई बहुत आक्रामक ढंग से करने को कहा जा रहा है।
• Prediction entropy: noisy microstructure data पर overconfident होने वाला pruned model training regime को fit कर रहा हो सकता है।

निष्कर्ष

Methods अच्छी तरह स्थापित हैं और sweep चलने तक यह लेख केवल इतना ही दावा करता है। Unstructured pruning आपको sparsity number देती है, speed नहीं। Structured pruning speed देती है यदि — और केवल यदि — आप layers को mask करने के बजाय rebuild करें। Lottery Ticket Hypothesis संकेत देती है कि compact model overparameterized model के भीतर पहले से मौजूद है, हालांकि यह image benchmarks पर दिखाया गया है, order-book data पर नहीं; ऊपर दिए गए वे तीन कारण कि यह market data पर "काम करना चाहिए" hypotheses हैं और उनसे जुड़े experiments हैं, findings नहीं।

साहित्य से मिलने वाला practical heuristic है कि शुरुआत से छोटा design करने के बजाय बड़ा train करके prune down करें: बड़ा model loss landscape को अधिक प्रभावी ढंग से explore करता है और pruning महत्वपूर्ण pathways को बचाए रखती है। Trading model के लिए यह किस sparsity और किस accuracy cost पर सही है, यह एक IMP sweep दूर है — और sweep के बाद इस लेख को numbers के साथ फिर पढ़ा जाना चाहिए।