Model Pruning for Low-Latency Trading Inference

Наша статья о DeepLOB завершает раздел развертывания тремя пунктами — ONNX с TensorRT, квантование INT8 и FPGA, — не разбирая ни один из них. Эта статья рассматривает первую общую проблему: модель больше, чем нужно. Pruning нейросети удаляет избыточные параметры, а интересное утверждение литературы состоит не просто в экономии памяти: подсеть, содержащая 10–20% весов, может сохранить точность плотной модели.

То, что задержка имеет значение, уже обсуждалось в блоге — в ZigBolt для пути обмена сообщениями и в налоге на IPC с арифметикой безубыточности. В моделировании спреда уже разобран компромисс между быстрой, но немного худшей, и медленной, но лучшей моделью, а таблица градиентного бустинга и глубокого обучения содержит строку задержки вывода. Но ни один материал не объясняет, как уменьшить конкретную модель. В цикле исполнения вывод модели находится под нашим полным контролем, а транспортные ветви покрыты воспроизводимыми p50/p95/p99 в передаче данных для алгоритмической торговли.

О чем эта статья: математические основы и рабочий код для pruning по величине, структурного pruning, итеративного pruning по величине, movement pruning, дистилляции знаний и полуструктурированной разреженности NVIDIA 2:4 применительно к торговому MLP.

Чем это не является: измеренным результатом. В каждой эмпирической статье блога есть ссылка на источник или сопутствующий репозиторий, а здесь пока нет ни того, ни другого. Кривая "разреженность — точность — задержка" ниже приведена как эксперимент, который еще предстоит провести, а не как таблица для цитирования. Относитесь ко всему здесь как к методу, а к числам — как к ожиданиям.

Что дает вам обрезка

Главное ограничение — размер модели. Рассмотрим среднечастотную модель: 4-слойный MLP с 2048 скрытыми нейронами для книги ордеров:

При $d_{\text{in}} = 100$, $h = 2048$, $L = 4$, $d_{\text{out}} = 3$ это около 12,6 млн параметров, то есть примерно 48 МБ в float32. L2-кэш обычно имеет размер 1–4 МБ, поэтому веса не помещаются и при каждом прямом проходе подтягиваются издалека. Удаление 95% весов оставляет около 630 000 эффективных параметров и 2,4 МБ — уже подходящий размер.

Превратится ли это в ускорение по времени, зависит от того, упирается ли ядро в пропускную способность памяти. Это вопрос арифметической интенсивности, а не одного размера модели. Лестница скорости движка бэктеста применяет roofline-модель (Уильямс, Уотерман и Паттерсон) к измеренному примеру, а не к заявлению о коэффициенте штрафа. Здесь действует та же логика: прежде чем требовать ускорения, измерьте объем перемещаемых байтов.

Основы обрезки

Неструктурированная обрезка

Самый простой подход — обнулить отдельные веса по их величине. Для матрицы весов $\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 раз; на плотном процессоре GEMM это может не дать даже 1,01x. Реальные пути к ускорению — структурный pruning (меньшие матрицы) и полуструктурированная разреженность 2:4 с аппаратной поддержкой.

Структурированная обрезка

Структурированный pruning удаляет целые нейроны, каналы или головы внимания. Для линейного слоя $\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$ с $\mathbf{W} \in \mathbb{R}^{m \times n}$ удаление нейрона $j$ обнуляет $j$-ю строку $\mathbf{W}$ и $j$-й элемент $\mathbf{b}$:

Сначала удаляйте нейроны с наименьшими $\ell_2$-нормами. Это единственный вариант, который создает действительно меньшие матрицы, но только если слои физически перестроены. Обнуление строк при сохранении исходной формы тензора не меняет число FLOP; именно шаг перестройки из раздела реализации превращает маску в матрицу $512 \times n$.

Для сверточных слоев аналогом служит pruning фильтров. Для $\mathbf{W} \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times k \times k}$ важность выходного фильтра $i$ равна:

Удаление фильтра $i$ исключает весь выходной канал и пропорционально уменьшает число FLOP.

Гипотеза о лотерейном билете

В 2019 году Франкл и Карбин предложили гипотезу лотерейного билета (LTH): внутри случайно инициализированной плотной сети существует разреженная подсеть — "выигрышный билет", которая при обучении с исходной инициализацией достигает точности всей сети за сопоставимое число итераций.

Формально пусть $f(\mathbf{x}; \boldsymbol{\theta}_0)$ инициализируется параметрами $\boldsymbol{\theta}_0 \sim \mathcal{D}_\theta$. После обучения до сходимости получаем $\boldsymbol{\theta}^*$ и маску pruning $\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$ параметров. После 10 раундов при $p = 0.2$ остается примерно 10,7%.

Три гипотезы о торговых моделях, ни одна из них не проверена

Соблазнительно утверждать, что LTH особенно хорошо подходит для рыночных данных. Возникают три аргумента; все три — гипотезы, и выдавать их за факты — именно тот путь к ошибке, которого блог старается избегать.

1. Финансовые сигналы разрежены. Большая часть снимка книги ордеров — шум, поэтому разреженная подсеть может естественно соответствовать разреженному сигналу. Проверка: сравнить IMP со случайной маской той же разреженности; если дело в самой разреженности, случайная маска не должна сильно отставать.
2. Выигрышные билеты переносятся между режимами. Это неподтвержденное эмпирическое утверждение о рынках и самое интересное из трех. Его можно напрямую проверить по меткам режимов из обнаружения режима с помощью HMM: найдите билет в режиме 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():,}")

Неструктурированный 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%}")

Структурированный pruning с перестройкой, который действительно ускоряет модель

Маскирование строк — только половина работы. Ускорение дает физическая перестройка каждого слоя в уменьшенной форме, а значит, удаление распространяется вперед: удаление строки $j$ слоя $i$ также удаляет столбец $j$ слоя $i+1$ и канал $j$ любого BatchNorm1d между ними.

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

Прежде чем доверять этому коду, проверьте две вещи в том же духе, что и остальные контрольные пороги эквивалентности в блоге:

• Формы. rebuilt должен иметь скрытые размеры $(1 - \text{fraction}) \cdot h$ — 512 при fraction=0.75 и $h = 2048$, — а число параметров должно уменьшиться квадратично, поскольку сокращаются оба измерения внутренних матриц.
• Выходы. В режиме eval() результат rebuilt(x) должен совпадать с результатом маскированной модели с точностью вычислений на одной партии. Если это не так, распространение удаления столбцов выполнено неверно, и все последующие измерения относятся уже не к той модели, о которой вы думаете.

Проверка выживания строк предполагает, что замаскированная строка равна нулю, а активная — нет. Для выхода ln_structured это верно. Другая процедура могла бы создать действительно нулевой, но живой нейрон, поэтому сравнивайте число выживших с запрошенной долей, а не слепо доверяйте стандартному тесту.

Итеративный pruning по величине (поиск лотерейных билетов)

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 против pandas. Для pruning важны три момента:

• Оценивайте перестроенную модель, а не маскированную. Маскированная модель с размером партии 1 по-прежнему измеряет плотную форму.
• Указывайте размер партии. Партия 1 (цикл котирования) и партия 256 (исследовательская очистка) находятся по разные стороны границы между памятью и вычислениями, поэтому pruning помогает им по-разному.
• Отчитывайтесь о точности на том же разбиении и горизонте, с явным определением метки. Таблица задержек без столбца точности — аргумент в пользу полного удаления модели.

Продвинутые методы

Сокращение с помощью дистилляции знаний

Вместо изолированного pruning и дообучения используйте исходную плотную модель как учителя. Обрезанный ученик минимизирует комбинацию потери задачи и 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

Обрезка движений

Вместо pruning по абсолютной величине movement pruning (Sanh et al., 2020) удаляет веса, которые во время обучения движутся к нулю. Оценка важности накапливает произведение изменения веса и градиента:

Веса с отрицательными оценками удаляются. Преимущество перед pruning по величине относится именно к дообучению: при адаптации предварительно обученной модели распределение величин сформировано исходной задачей, поэтому величина — устаревший сигнал важности, а направление движения — более свежий. Для торговой модели, переобучаемой на скользящих окнах, это более типичная ситуация, чем обучение с нуля.

Структурированная разреженность NVIDIA 2:4

Графические процессоры NVIDIA Ampere и более поздних поколений аппаратно поддерживают структурированную разреженность 2:4: из каждых четырех соседних весов ровно два должны быть нулевыми.

Это единственная мелкозернистая форма разреженности, которую действительно вознаграждает оборудование, поэтому она важнее числа 90% нулей после неструктурированного pruning. Ограничение локальное, а не глобальное: ему не важно, какие два веса из четырех выживут. Поэтому оно гораздо мягче фиксированной глобальной маски, хотя 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()

Для ускорения путь вывода должен использовать разреженные тензорные ядра — экспорт ONNX со сборкой TensorRT или torch.sparse.to_sparse_semi_structured. Экспорт модели с маской 2:4 при плотном выполнении ухудшает точность и не дает преимуществ.

Производственное развертывание

Проверка

Обрезанная модель — новая модель, а не просто сжатая старая. Она проходит те же приемочные этапы, что и любой другой кандидат: последовательное переобучение и повторную вневыборочную проверку по walk-forward оптимизации, с поправкой на эффект отбора из дефлированного коэффициента Шарпа. Здесь поправка особенно важна: IMP создает последовательность моделей-кандидатов, и выбранный по лучшему виду за десять раундов уровень разреженности требует дефлятирования Sharpe на эффективное число испытаний. Простое правило вроде "отклонять при падении Sharpe более чем на 5%" такой арифметики не выдерживает, поэтому его здесь нет.

Наложение квантования

Обрезка сочетается с квантованием. Модель, разреженная на 90% и квантованная до INT8, имеет степень сжатия:

Модель размером 48 МБ превращается в 1,2 МБ. Это только требование к хранению. Дает ли модель 1,2 МБ те же решения — отдельный вопрос, который требует отдельного ответа. Ловушка точности графического процессора показывает, почему его нужно задавать: в блоге одна только fp32-арифметика дала относительную ошибку 211 в бэктесте, хотя результат выглядел правдоподобно. INT8 — гораздо более агрессивное сокращение. Отправлять квантованную и разреженную модель можно только после количественной проверки паритета с плотной fp32-моделью: доля совпавших решений и дельта PnL на отложенном периоде, а не обещание.

Мониторинг

Сокращенные модели могут быть более чувствительны к сдвигу распределения. Стоит посмотреть:

• Разреженность активаций: если выжившие нейроны выдают в основном нули, эффективная модель меньше ожидаемой и, вероятно, деградирует.
• Нормы градиентов при дообучении: взрывные градиенты означают, что выжившей подсети приходится слишком агрессивно компенсировать удаленное.
• Энтропия прогнозов: обрезанная модель, становящаяся чрезмерно уверенной на зашумленных микроструктурных данных, скорее всего, подогналась под режим обучения.

Заключение

Методы хорошо известны, но до проверки это все, что утверждает статья. Неструктурированный pruning дает разреженность, но не скорость. Структурированный pruning дает скорость, только если физически перестроить слои, а не просто замаскировать их. Гипотеза лотерейного билета предполагает, что компактная модель уже существует внутри переизбыточной, хотя это показано на эталонных изображениях, а не на данных книги ордеров. Три приведенные причины, почему это "должно" работать на рынке, — гипотезы с соответствующими экспериментами, а не выводы.

Практическая эвристика литературы — обучать большую модель и затем сокращать ее, а не проектировать маленькую с нуля: большая модель эффективнее исследует ландшафт потерь, а pruning сохраняет важные пути. Верно ли это для торговой модели, какая разреженность оптимальна и чего стоит один проход IMP — ответы появятся только после этого прохода. Тогда эту статью нужно будет перечитать уже с цифрами.