Model Pruning for Low-Latency Trading Inference

Nuestro artículo de DeepLOB finaliza su sección de implementación con tres puntos (ONNX más TensorRT, cuantificación INT8, FPGA) y no trata ninguno de ellos. Este artículo es el tratamiento que falta para el primer problema subyacente a los tres: el modelo es más grande de lo necesario. La poda de la red neuronal elimina parámetros redundantes, y la afirmación interesante en la literatura no es que esto ahorre memoria sino que una subred que contenga entre el 10 y el 20 % de los pesos puede igualar la precisión del modelo denso.

Que la latencia importa es un argumento que el blog ya ha planteado: ZigBolt en la ruta de mensajería y el impuesto IPC con la aritmética del punto de equilibrio, y spread modeling ya posee el Compensación rápida pero ligeramente peor versus lenta pero mejor, completa con una tabla de aumento de gradiente versus aprendizaje profundo que tiene una fila de latencia de inferencia. Lo que ninguno de ellos cubre es cómo hacer que un modelo determinado sea más pequeño. De las etapas de un ciclo de cotización, la inferencia de modelos es la que está totalmente bajo nuestro control; los tramos de transporte están cubiertos, con números p50/p95/p99 reproducibles, en comunicación de datos para algotrading.

Qué es este artículo: el código matemático y de trabajo para la poda de magnitud, la poda estructurada, la poda de magnitud iterativa, la poda de movimiento, la destilación de conocimientos y la escasez semiestructurada de NVIDIA 2:4, aplicados a un MLP comercial.

Lo que no es: un resultado medido. Cada artículo empírico de este blog incluye una línea de procedencia o un repositorio complementario, y este aún no incluye ninguno de los dos. La curva de escasez versus precisión versus latencia se presenta a continuación como un experimento para ejecutar, no como una tabla para citar. Trate todo aquí como el método y los números como pendientes.

Lo que te compra la poda

La restricción es de tamaño uno. Considere un modelo de frecuencia media: un MLP de 4 capas con 2048 unidades ocultas sobre las características del libro de pedidos:

Para $d_{\text{in}} = 100$, $h = 2048$, $L = 4$, $d_{\text{out}} = 3$, es decir, aproximadamente 12,6 millones de parámetros, aproximadamente 48 MB en float32. L2 suele ser de 1 a 4 MB, por lo que los pesos no encajan; se transmiten desde más lejos en cada pase hacia adelante. Pode el 95% de ellos y tendrá aproximadamente 630.000 parámetros efectivos y 2,4 MB, lo cual sí encaja.

Que eso se traduzca en tiempo de reloj de pared depende de si el núcleo está ligado a la memoria, y esa es una cuestión de intensidad aritmética más que una cuestión de tamaño. La escalera de velocidad del motor de prueba retrospectiva trabaja el modelo de línea del techo (Williams, Waterman & Patterson) a través de un ejemplo medido en lugar de afirmar un factor de penalización; Aquí se aplica el mismo marco, y debería aplicarse la misma disciplina: medir los bytes movidos antes de reclamar la aceleración.

Fundamentos de poda

Poda no estructurada

El enfoque más simple: establecer los pesos individuales en cero según su magnitud. Dada una matriz de peso $\mathbf{W} \in \mathbb{R}^{m \times n}$, crea una máscara binaria $\mathbf{M} \in \{0, 1\}^{m \times n}$ tal que:

dónde $\tau$ es un umbral elegido para alcanzar el nivel de escasez deseado $s$:

La matriz podada es $\mathbf{W}' = \mathbf{W} \odot \mathbf{M}$, con $\odot$ el producto Hadamard. La intuición es que los pesos cercanos a cero contribuyen poco a la salida de la capa.

El problema, planteado claramente porque es fácil malinterpretar el número de escasez: la escasez no estructurada no se traduce en una aceleración en el hardware estándar. Una matriz con 90% de ceros todavía emite la misma cantidad de acumulaciones múltiples a menos que cambie a núcleos dispersos o hardware con soporte de escasez. Cuando se imprime el siguiente código Sparsity: 90.0%, eso es un recuento de ceros: no es 10x, y en un CPU GEMM denso tampoco es 1,01x. Los caminos que sí ganan tiempo son la poda estructurada (matrices más pequeñas) y la escasez semiestructurada 2:4 (soporte de hardware), ambas a continuación.

Poda estructurada

La poda estructurada elimina neuronas, canales o cabezas de atención enteras. Para una capa lineal $\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$ con $\mathbf{W} \in \mathbb{R}^{m \times n}$, eliminando neurona $j$ pone a cero el $j$-ésima fila de $\mathbf{W}$ y el $j$-ésimo elemento de $\mathbf{b}$:

Neuronas con el menor $\ell_2$-la norma va primero. Esta es la variante que produce matrices realmente más pequeñas, pero sólo si realmente reconstruyes las capas. Poner a cero las filas y dejar el tensor en su forma original no cambia nada en el recuento de FLOP; El paso de reconstrucción en la sección de implementación es lo que convierte la máscara en una $512 \times n$ matriz.

Para capas convolucionales el análogo es la poda de filtro. Dado $\mathbf{W} \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times k \times k}$, la importancia del filtro de salida $i$ es:

Quitar filtro $i$ elimina un canal de salida completo, reduciendo los FLOP proporcionalmente.

La hipótesis del billete de lotería

En 2019, Frankle y Carbin introdujeron la hipótesis del billete de lotería (LTH): dentro de una red densa inicializada aleatoriamente existe una subred dispersa (un "billete ganador") que, entrenada desde su inicialización original, coincide con la precisión de la red completa en un número comparable de iteraciones.

Formalmente, considere $f(\mathbf{x}; \boldsymbol{\theta}_0)$ inicializado con $\boldsymbol{\theta}_0 \sim \mathcal{D}_\theta$. Después del entrenamiento para la convergencia obtenemos $\boldsymbol{\theta}^*$ y obtener una máscara de poda $\mathbf{m}$. La LTH afirma que existe $\mathbf{m}$ tal que:

con $\|\mathbf{m}\|_0 \ll |\boldsymbol{\theta}_0|$. Los experimentos originales se realizaron en MNIST y CIFAR-10, donde los billetes ganadores conservaban entre el 10 y el 20 % de los parámetros. Nada de eso se transfiere a los datos del libro de pedidos por suposición: las características LOB no son estacionarias y la etiqueta es casi ruido, lo cual es un régimen diferente de la clasificación de imágenes exactamente en las formas que podrían importar.

Poda de magnitud iterativa (IMP)

El ticket lo encuentra IMP:

1. Inicialice la red con $\boldsymbol{\theta}_0$.
2. Entrenar hacia la convergencia, obteniendo $\boldsymbol{\theta}^*$.
3. Pode el $p\%$ de pesos con magnitud más pequeña, creando máscara $\mathbf{m}$.
4. Restablecer los pesos supervivientes a sus valores en $\boldsymbol{\theta}_0$ (rebobinado).
5. Repita desde el paso 2 con la red enmascarada.

Cada fracción de ciruelas pasas redondas. $p$ (normalmente 20%), por lo que después $n$ rondas $(1-p)^n$ de los parámetros sobreviven. Después de 10 rondas en $p = 0.2$, queda aproximadamente el 10,7%.

Tres hipótesis sobre los modelos comerciales, ninguna de ellas probada

Es tentador argumentar que LTH debería funcionar especialmente bien con datos de mercado. Surgen tres de esos argumentos; las tres son hipótesis, y expresarlas como hechos es el modo de fracaso que este blog pretende evitar.

1. Las señales financieras son escasas. La mayor parte de una instantánea del libro de pedidos es ruido, por lo que una subred dispersa podría estar alineada naturalmente con una señal dispersa. Comprobable: compare IMP con una máscara aleatoria de la misma escasez; Si la escasez misma está haciendo el trabajo, la máscara aleatoria no debería quedarse atrás.
2. Los boletos ganadores se generalizan en todos los regímenes. Esta es una afirmación empírica sobre los mercados sin ninguna cita detrás de ella, y es la más interesante de las tres. Se puede comprobar directamente con las etiquetas del régimen de detección de régimen con HMM: busque el ticket en el régimen A, vuelva a entrenarlo en el régimen B y compárelo con un ticket encontrado de forma nativa en B.
3. La escasez se regulariza. Una capacidad efectiva más baja puede reducir el ajuste al ruido de la microestructura, lo que se manifestaría como que la brecha fuera de la muestra del modelo podado sería más pequeña que la del modelo denso, no simplemente comparable.

Implementación: poda de un MLP comercial

El modelo base

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

Poda de magnitud no estructurada

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

Poda estructurada, con la reconstrucción que la hace real

Ocultar filas es la mitad del trabajo. La parte que produce la aceleración es reconstruir cada capa en su forma reducida, lo que significa propagar la eliminación hacia adelante: soltar fila $j$ de capa $i$ también cae la columna $j$ de capa $i+1$ y canal $j$ de cualquier BatchNorm1d entre ellos.

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

Dos cosas para comprobar antes de confiar en esto, en el mismo espíritu que las puertas de equivalencia que se publican en el resto del blog:

• Formas. rebuilt debe mostrar dimensiones ocultas en $(1 - \text{fraction}) \cdot h$ — 512 para fraction=0.75, $h = 2048$ - y un recuento de parámetros que ha caído cuadráticamente, ya que ambas dimensiones de las matrices interiores se reducen.
• Salidas. En eval() modo, rebuilt(x) debe coincidir con el modelo enmascarado rebuilt-Salida libre a tolerancia de punto flotante en el mismo lote. Si no es así, la propagación de la columna es incorrecta y cada número posterior mide un modelo diferente al que cree.

La prueba de supervivencia de filas supone que una fila enmascarada es exactamente cero y una fila activa no lo es. Eso es válido para ln_structured producción; no sería válido si algún otro procedimiento produjera una neurona viva genuinamente totalmente cero, por lo que se debe afirmar el recuento de supervivientes frente a la fracción solicitada en lugar de confiar ciegamente en la prueba de la norma.

Poda de magnitud iterativa (búsqueda de billetes de lotería)

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 es la materia prima para la curva de escasez versus precisión que este artículo le debe. eval_fn tiene que estar genuinamente fuera de muestra, en divisiones purgadas: una ejecución IMP calificada dentro de la muestra informará una hermosa curva que no significa nada.

Medirlo

La latencia se mide con la misma convención de arnés que el resto del blog (calentamiento excluido, mejor de N, p50/p95/p99 informado en lugar de una media) y ese protocolo, con el código, se encuentra en Polars vs pandas. Tres puntos específicos de la poda:

• Compare el modelo reconstruido, no el modelo enmascarado. Un modelo enmascarado del tamaño de lote 1 mide la forma densa.
• Informar el tamaño del lote. El lote 1 (bucle de citas) y el lote 256 (barrido de investigación) se encuentran en lados diferentes de la línea ligada a la memoria/ligada a la computación, y la poda los ayuda de manera diferente.
• Informar la precisión en la misma división, en el mismo horizonte, con la definición de la etiqueta indicada. Una tabla de latencia sin la columna de precisión de coincidencia es un argumento para eliminar el modelo por completo.

Técnicas Avanzadas

Poda con destilación del conocimiento

En lugar de podar y ajustar de forma aislada, utilice el modelo denso original como maestro. El estudiante podado minimiza una combinación de pérdida de tareas y la divergencia KL de la distribución de producción del maestro:

dónde $\mathbf{z}_t$ y $\mathbf{z}_s$ son logits de profesores y estudiantes, $T$ es la temperatura y $\alpha$ equilibra los objetivos. El $T^2$ El factor cambia la escala de los gradientes de destilación, que de otro modo se reducen a medida que $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

Poda de movimiento

En lugar de podar por magnitud absoluta, la poda de movimiento (Sanh et al., 2020) poda los pesos que se mueven hacia cero durante el entrenamiento. La puntuación de importancia acumula el producto gradiente-peso:

Se podan los pesos con puntuaciones negativas. El argumento a favor de la poda de magnitud tiene que ver específicamente con el ajuste fino: cuando se adapta un modelo previamente entrenado, la distribución de magnitud fue moldeada por la tarea previa al entrenamiento, por lo que la magnitud es una señal de importancia obsoleta y la dirección de viaje es una más nueva. Para un modelo comercial reentrenado en ventanas móviles, esa es la situación más común que el entrenamiento desde cero.

Disparidad estructurada NVIDIA 2:4

Las GPU NVIDIA Ampere y posteriores admiten una escasez estructurada de hardware 2:4: de cada 4 pesos contiguos, exactamente 2 deben ser cero.

Esta es la única forma de escasez detallada que el hardware realmente recompensa, razón por la cual es más importante que el número de ceros del 90% de la poda no estructurada. La restricción es local más que global (no le importa cuáles dos de cada cuatro sobrevivan), por lo que es una restricción mucho más débil que arreglar una máscara global, aunque el 50% es el único nivel de escasez que se ofrece.

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()

Lograr la aceleración requiere que la ruta de inferencia utilice los núcleos tensoriales dispersos: una exportación de ONNX más una compilación de TensorRT, o torch.sparse.to_sparse_semi_structured. Exportar un modelo enmascarado 2:4 a través de un tiempo de ejecución denso le brinda el costo de precisión y ningún beneficio.

Implementación de producción

Validación

Un modelo podado es un modelo nuevo, no uno antiguo comprimido, y pasa por la misma puerta de aceptación que cualquier otro candidato: reentrenamiento continuo y revalidación fuera de muestra por optimización de avance, con la corrección del efecto de selección de el índice de Sharpe deflactado. Esa corrección no es opcional aquí: IMP genera una secuencia de modelos candidatos, por lo que en la búsqueda se eligió el nivel de escasez que se ve mejor en diez rondas, y su Sharpe necesita desinflarse por el número efectivo de pruebas. Una regla plana como "rechazar si Sharpe cae más del 5%" no sobrevive a esa aritmética, razón por la cual no encontrará ninguna en este artículo.

Apilamiento de cuantificación

La poda compone con cuantización. Un modelo 90% disperso y cuantificado a INT8 tiene una relación de compresión de:

Un modelo de 48 MB pasa a ser de 1,2 MB. Ese es un reclamo de almacenamiento y nada más. Si el modelo de 1,2 MB produce las mismas decisiones es una pregunta separada con su propia respuesta, y la trampa de precisión de GPU es la razón para preguntarlo: en este blog, se ha demostrado que fp32 por sí solo produce un error relativo de 211 en un cálculo de prueba retrospectiva que parecía completamente razonable. INT8 es una reducción mucho más agresiva que eso. Envíe un modelo cuantificado y podado solo detrás de una puerta de paridad cuantificada frente al modelo denso fp32: tasa de acuerdo de decisión y delta de PnL en un período de espera, no una garantía.

Monitoreo

Los modelos podados pueden ser más sensibles al cambio de distribución. Vale la pena verlo:

• Escasez de activación: si las neuronas supervivientes emiten principalmente ceros, el modelo efectivo es más pequeño de lo previsto y probablemente se degrada.
• Normas de gradiente durante el reentrenamiento: los gradientes explosivos sugieren que a la subred superviviente se le está pidiendo que compense de manera demasiado agresiva lo que se eliminó.
• Entropía de predicción: un modelo podado que se vuelve demasiado confiado en datos de microestructura ruidosos probablemente se ajuste al régimen de entrenamiento.

Conclusión

Los métodos están bien establecidos y, hasta que se realice el barrido, eso es todo lo que afirma este artículo. La poda no estructurada le da un número de escasez y ninguna velocidad. La poda estructurada le proporciona velocidad si, y sólo si, reconstruye las capas en lugar de enmascararlas. La hipótesis del billete de lotería sugiere que el modelo compacto ya existe dentro del modelo sobreparametrizado, aunque eso se ha demostrado en puntos de referencia de imágenes y no en datos de la cartera de pedidos, y las tres razones por las que "debería" funcionar en los datos de mercado indicados anteriormente son hipótesis con experimentos adjuntos, no hallazgos.

La heurística práctica de la literatura es entrenar en grande y podar en lugar de diseñar en pequeño desde el principio: el modelo grande explora el panorama de pérdidas de manera más efectiva y la poda preserva los caminos que importaban. Si eso es válido para un modelo comercial, con qué escasez y con qué costo de precisión, está a un paso de IMP, y este artículo debe leerse nuevamente después de ese barrido, con números en él.