L₀-регуляризация: как обучать разреженные нейросети напрямую

Почему L₀ не поддавался градиентному спуску

Разреженность нейросети — это доля весов, равных ровно нулю. Чем их больше, тем меньше памяти и вычислений нужно на инференс. Идеальная метрика для контроля разреженности — норма L₀, то есть просто счётчик ненулевых элементов. Проблема в том, что этот счётчик — кусочно-постоянная функция: её производная равна нулю почти всюду и не определена в точках скачка. Градиентному спуску не за что зацепиться.

Из-за этого на практике десятилетиями использовали суррогаты. L₁-регуляризация (Lasso) добавляет к функции потерь сумму модулей весов и загоняет часть из них к нулю, но одновременно тянет вниз и большие, полезные веса, искажая модель. L₂ (weight decay) вообще не создаёт настоящих нулей — только уменьшает значения. Магнитудный прунинг обрезает мелкие веса после обучения, но это отдельный этап, а не часть оптимизации.

Работа Кристоса Луизоса, Макса Веллинга и Диедерика Кингмы «Learning Sparse Neural Networks through L₀ Regularization» (2018) предложила способ включить L₀ прямо в обучаемую функцию потерь и оптимизировать её обычным SGD. Идея строится на стохастических гейтах и хитром перепараметрическом трюке.

Гейты вместо жёсткого обнуления

Каждый вес умножается на бинарный гейт z ∈ {0, 1}. Если гейт закрыт (z = 0), вес выключен независимо от своего значения. Число открытых гейтов — это и есть L₀-норма. Напрямую перебирать бинарные комбинации невозможно, поэтому гейты делают случайными: каждый z берётся из распределения Бернулли с параметром, который сеть учит сама.

Ожидаемая L₀-норма при таком подходе — это сумма вероятностей открытия всех гейтов. Уже гладкая величина. Но остаётся вторая проблема: сэмплирование из Бернулли недифференцируемо, градиент через него не пройдёт. Здесь и появляется главный вклад работы.

Hard concrete distribution

Авторы берут непрерывную релаксацию бинарного распределения — concrete (она же Gumbel-Softmax в бинарном случае) — и растягивают её носитель за пределы отрезка [0, 1], а затем зажимают результат обратно функцией hard-sigmoid. Получается hard concrete: распределение, которое даёт ненулевую вероятность точным нулям и точным единицам, но между ними ведёт себя непрерывно и дифференцируемо через reparameterization trick¹.

Практически это выглядит так. Для каждого гейта заводится обучаемый лог-параметр α. При прямом проходе сэмплируется шум u из равномерного распределения, через него и α вычисляется мягкое значение гейта. Итоговая штрафная добавка к лоссу — аналитическая формула, зависящая только от α, поэтому её градиент считается точно, а не оценивается сэмплами.

Ключевой сдвиг мышления: вы не обрезаете сеть после обучения, а платите за каждый включённый параметр прямо в функции потерь. Оптимизатор сам решает, какие связи стоят своей цены.

Что даёт метод на практике

Во время обучения используется мягкий стохастический гейт, а на инференсе — детерминированный: берётся ожидаемое значение гейта, зажатое в [0, 1], и веса с нулевым гейтом физически удаляются. Это отличает подход от L₁, где обнуление нужно ещё вручную дочищать порогом.

• Настоящие нули во время обучения. Разреженность появляется как побочный продукт оптимизации, а не отдельным постпроцессингом.
• Условные вычисления. Закрытые гейты можно не считать в прямом проходе, что ускоряет и само обучение, а не только инференс.
• Групповая разреженность. Один гейт можно навесить не на отдельный вес, а на целый нейрон или канал свёртки — тогда метод выключает структурные блоки, и модель реально уменьшается в размерах, а не просто становится «дырявой».

Как это соотносится с другими подходами

Псевдокод гейта

Логика hard concrete гейта на инференсе и обучении в упрощённом виде. Параметры beta, gamma, zeta — гиперпараметры релаксации, log_alpha — обучаемый параметр гейта.

import torch

def hard_concrete_gate(log_alpha, beta, gamma, zeta, training):
    if training:
        u = torch.rand_like(log_alpha).clamp(1e-6, 1 - 1e-6)
        s = torch.sigmoid((torch.log(u) - torch.log(1 - u) + log_alpha) / beta)
    else:
        # детерминированный проход: без шума
        s = torch.sigmoid(log_alpha)
    # растягиваем в [gamma, zeta] и зажимаем в [0, 1]
    s_stretched = s * (zeta - gamma) + gamma
    return torch.clamp(s_stretched, 0.0, 1.0)

def l0_penalty(log_alpha, beta, gamma, zeta):
    # аналитическая вероятность того, что гейт открыт (z > 0)
    return torch.sigmoid(log_alpha - beta * torch.log(torch.tensor(-gamma / zeta)))

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

Метод не универсальное лекарство. Он добавляет гиперпараметры (beta, gamma, zeta, лямбда), к которым чувствителен результат, и требует аккуратной инициализации гейтов, иначе сеть либо не разрежается вовсе, либо схлопывается слишком рано. Но как инструмент структурной обрезки во время обучения L₀-регуляризация остаётся одной из немногих схем, которые оптимизируют именно то, что вам нужно, — число живых параметров.