Теперь у нас получилась многослойная архитектура компьютерной нейронной сети. Например такая:
ЛинейныйСлой1 -> Активация1 -> ЛинейныйСлой2 -> Активация2 -> ЛинейныйСлой3
Подобная нейросеть называется MLP (multi-layer perceptron) и вполне способна со средненькой уверенностью предсказывать цифру на показанном ей изображении или отличать кошку от собаки.

Теперь главное, ради чего писалась вся эта банальщина.

Все описанные выше допущения хороши тем, что в своё время они позволяли упростить сложное, получить работающее решение или ещё что-то полезное. Они наслаивались друг на друга и привели нас сейчас к многослойным нейросетям, которые отошли от своего прототипа настолько далеко, что слабо понятно, почему они всё ещё называются нейронными.

Похожи ли процессы в современных компьютерных нейросетях на происходящее в нервной системе животных? Абсолютно нет. И на самом деле это здорово, ибо открывает для нас довольно широкое окно возможностей по созданию других прототипов нервной системы животных.

Думал, с чего бы простого начать собирать новую картину мира нейросетей...

Можно, наверное, с загадки. Как обучить нейронку в 10 000 слоёв?

Предположим, что нам показалось, что такая большая глубина нейросети может дать accuracy 100% на MNIST-е. Ну или просто есть спортивный интерес попробовать обучить глубокую сеть, ибо где-то слышали про затухание градиентов и хочется понять, что это.

Почему MNIST? Потому что он простой, картинки маленькие 1х28х28, легко помещается в память, можно учить на компе без GPU...
Почему не MNIST? И практики на нём часто не проявляются те сложности, которые потом вылезают на реальных датасетах.
В данном случае нас это не затронет. Нам главное хоть как-то взлететь...

Начнём с самого простого: много линейных слоёв и ReLU. Например нейросеть из 1к слоёв можно описать так:

class ManyLinears(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)
        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16),
                nn.ReLU()
            ) for _ in range(1000)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = F.relu(x)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

Легко заметить, что никакого обучения не происходит. Значение loss-функции на обучающей выборке никак не меняется почему-то. Немного поэкспериментировав можно заметить, что если снизить количество линейных слоёв примерно до 20, то обучение начинается.

Тут можно вспомнить про затахание градиентов, и что от него хорошо помогает батч-нормализация. Меняем нашу сеть:

class ManyLinearsWithBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)
        self.bn1 = nn.BatchNorm1d(16)

        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16),
                nn.BatchNorm1d(16),
                nn.ReLU()
            ) for _ in range(10000)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = self.bn1(x)
        x = F.relu(x)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

Запускаем обучение и происходит вообще странное. На первых итерациях обучения значение loss-функции становится nan.

Что с этим делать? Если это редкая история, то можно в случае nan в лосе просто пропускать такой "плохой" батч. Добавляем следующие строчки:

if not torch.isnan(loss):
    loss.backward()
    optimizer.step()
    print('Loss is not nan')

Запускаем и узнаём, что все итерации кроме самой первой дают nan в loss-е. Выходит, что эта первая итерация как-то портит нам всю сеть.

Наверное оптимизатор плохо работает на первых итерациях. Меняем AdamW на классику: SGD.

optimizer = optim.SGD(model.parameters(), lr=args.lr)

Не помогло.

Тогда попробуем скорость обучения уменьшить в 1000 раз. Нам ведь пока важно от nan избавиться, поэтому пока не важно, если будет медленно обучаться. Делаем 'lr=3e-7'.
Тоже не помогло.

Вернём оптимизатор и lr обратно и опять попробуем уменьшить нашу сеть. Экспериментально можно найти, что примерно 400 линейных слоёв дают nan в значении loss-функции не сразу, а ближе к середине первой эпохи. 300 линейных слоёв не имеют проблемы с nan в лосе, но и не обучаются. При 50 слоях вялое обучение начинается.

Самое время посмотреть, а что там в сети происходит, что не даёт ей обучаться. Давайте посмотрим, какие значения градиентов в начале, середине и конце сети во время обучения.

if not torch.isnan(loss):
    loss.backward()

    for name, parameter in model.named_parameters():
        if name in ['fc1.weight', 'fc1.bias', 'blocks.25.0.weight', 'blocks.25.0.bias', 'fc.weight', 'fc.bias']:
            print(name, parameter.grad.view(-1)[parameter.grad.view(-1)!=0].abs().log10().median().item())

    optimizer.step()
    #print('Loss is not nan')
else:
    print('Loss is nan')

(если есть идеи, как вывести значения градиентов в текстовом виде в более симпатичном виде, пишите варианты в комментах)
запускаем и получаем что-то вроде такого:

fc1.weight 1.8429147005081177
fc1.bias -4.214419841766357
blocks.25.0.weight 0.10246949642896652
blocks.25.0.bias -6.622660160064697
fc.weight -1.8109673261642456
fc.bias -1.6103878021240234

Они имеют плюс-минус один порядок значений. Т.е. не сказал бы, что они затухают. Значит дело в чём-то другом.

Выдвинем гипотезу, что ReLU, применённый много раз, почти ничего не оставляет от исходного сигнала. Каждый раз он зануляет все отрицательные значения. А после BatchNorm-а - это примерно половина. Да и если бы ReLU стоял после линейного слоя, то то тоже наверное была бы примерно половина входного сигнала занулена. На входе нашей нейросети у нас 28*28*256 бит информации. Мы 50 раз занулили половину из них. На выходе получили пшик и по нему пробовали обучаться. Попробуем заменить активацию на такую, которая не теряет информацию. Например, на LeakyReLU. И ожидаем, что она будет обучаться поживее. Пробуем вот такую сеть:

class ManyLinearsWithBNAndLeakyReLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)
        self.bn1 = nn.BatchNorm1d(16)

        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16),
                nn.BatchNorm1d(16),
                nn.LeakyReLU()
            ) for _ in range(50)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = self.bn1(x)
        x = F.leaky_relu(x)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

Обучилось до лучшего значения accuracy: 22% . Но всёравно маловато. Причём, если уменьшить глубину сети, то обучается до лучших значений метрики. Например, при 10 слоях: 95% .

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

Выдвинем ещё одну гипотезу, что дефолтное значение negative_slope у LeakyReLU, равное 0.01, применённое много раз к половине входного сигнала очень сильно снижает его силу. Количество информации в теории остаётся тем же, но на практике оно теряется при операциях с очень маленькими числами, попадая за границы точности float32. Пробуем negative_slope=2.0.

class ManyLinearsWithBNAndLeakyReLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)
        self.bn1 = nn.BatchNorm1d(16)

        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16),
                nn.BatchNorm1d(16),
                nn.LeakyReLU(negative_slope=2.0)
            ) for _ in range(50)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = self.bn1(x)
        x = F.leaky_relu(x, negative_slope=2.0)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

Обучается весьма шустро. И даёт 93% accuracy на 50 линейных слоях. Маленькая победа!
Пробуем увеличить количество слоёв и упираемся в число 120, когда вяленько, но учится.

P.S.
Выше мной написана глупость. Сила сигнала проходящего сигнала не снижается. Просто при negative_slope=2 начали меньше гаснуть градиенты.

Попробуем убрать BatchNorm и потюнить negative_slope у LeakyReLU так, чтобы в начале и конце нейросети градиенты были примерно одного порядка. Идея почертнута из выступлений Владислава Голощапова.

class ManyLinearsWithTunedLeakyReLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)

        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16),
                nn.LeakyReLU(negative_slope=2.33)
            ) for _ in range(120)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = F.leaky_relu(x, negative_slope=2.33)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

Сеть стала быстрее и легче, но обучению это не помогло. И сеть глубже так же не учится.

Если снова взглянуть на градиенты промежуточных слоёв, то можно заметить, что у bias и у weight они примерно одного порядка. Это приводит к тому, что на выходе слоя значения постоянно качает то в сторону увеличения, то в сторону уменьшения. И обучиться на таких входных данных следующему слою очень сложно. Это наблюдение также взято из выступлений Владислава Голощапова.

Для всех bias-ов снизим LR в сотню раз. Разделим bias и weight по разным группам параметров и в таком виде отдадим оптимизатору.

def param_groups(model, wieght_lr, bias_lr):
    wieght_params = []
    bias_params = []
    for name, param in model.named_parameters():
        if name.endswith(".bias"):
            #print("bias_params:", name, param.size())
            bias_params.append(param)
        else:
            #print("wieght_params:", name, param.size())
            wieght_params.append(param)

    return [{'params': wieght_params, 'lr': wieght_lr},
            {'params': bias_params, 'lr': bias_lr}]

...

    model_parameters = param_groups(model, args.lr, args.lr/100)
    optimizer = optim.AdamW(model_parameters, lr=args.lr)

Сразу начинает бодренько обучаться на 120 линейных слоях и доходит до 83% ассuracy. Возможно если дольше учить, настроить шудуллер LR и подобрать оптимальный LR получилось бы и выше accuracy. Но пока попробуем достичь большей глубины. При 250 линейных слоях при ещё большем уменьшении LR для bias-ов обучение ещё начинается, а больше - уже нет.

Возникает мысль вообще отказаться от bias-а в линейных слоях:

class ManyLinearsWithTunedLeakyReLUAndWithoutBias(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16, bias=False)

        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16, bias=False),
                nn.LeakyReLU(negative_slope=2.33)
            ) for _ in range(200)])
        self.fc = nn.Linear(16, 10, bias=False)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = F.leaky_relu(x, negative_slope=2.33)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

Но этот подход не даёт результат. Сеть и на 200 слоях не хочет обучаться. Хотя на 10 слоях учится шустро.

В поисках решения проблемы набредаем на Self-normalizing Neural Networks (SNNs) https://arxiv.org/pdf/1706.02515.pdf . И переделываем сеть так, чтобы активацией у нас была SELU, а веса были иницилизованы так, как это требует автор статьи. Нормализуем вход чуть иначе на всякий случай, хотя этого вроде не требуется. BatchNorm-ы такой сети тоже не нужны.

def init(m):
    if isinstance(m, nn.Linear):
        nn.init.constant_(m.bias, 0)
        #nn.init.kaiming_uniform_(m.weight, nonlinearity='linear')
        nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='linear')

...

class ManyLinearsWithSELU(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)

        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Linear(16, 16),
                nn.SELU()
            ) for _ in range(90)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = torchvision.transforms.functional.normalize(x, x.mean(), x.std())
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = F.selu(x)

        for block in self.blocks:
          x = block(x)

        x = self.fc(x)
        return x

...

    model = ManyLinearsWithSELU().to(device)

    model.apply(init)

К сожалению оказывается, что такая сеть тянет только 90 слоёв. Но вот что интересно: когда слоёв 100, она начинает обучаться, а потом перестаёт почему-то...

Причина, почему SSN начинает учиться, а потом перестаёт - это weight decay, который в AdamW по дефолту задан 0.01. Он очень большой и для начала обучения он не нужен совершенно. Если его отключить, и при это очень сильно ограничить величины градиентов:

torch.nn.utils.clip_grad_norm_(model.parameters(), 1e-6)

то SSN начинает учиться на 150 линейных слоях.
Правда на 200 начинает, а потом перестаёт почему-то.

Если вернуться к модельке с тюненым LeakyReLU и добавить адативный клипинг градиентов, то можно начать учиться на 400 линейных слоях:

# https://arxiv.org/abs/2102.06171
# https://github.com/rwightman/timm/blob/main/timm/utils/agc.py
def adaptive_clip_grad(parameters, clip_factor=0.01, eps=1e-3, norm_type=2.0):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    for p in parameters:
        if p.grad is None:
            continue
        p_data = p.detach()
        g_data = p.grad.detach()
        max_norm = unitwise_norm(p_data, norm_type=norm_type).clamp_(min=eps).mul_(clip_factor)
        grad_norm = unitwise_norm(g_data, norm_type=norm_type)
        clipped_grad = g_data * (max_norm / grad_norm.clamp(min=1e-6))
        new_grads = torch.where(grad_norm < max_norm, g_data, clipped_grad)
        p.grad.detach().copy_(new_grads)

def unitwise_norm(x, norm_type=2.0):
    if x.ndim <= 1:
        return x.norm(norm_type)
    else:
        # works for nn.ConvNd and nn,Linear where output dim is first in the kernel/weight tensor
        # might need special cases for other weights (possibly MHA) where this may not be true
        return x.norm(norm_type, dim=tuple(range(1, x.ndim)), keepdim=True)

...
            loss.backward()

            #torch.nn.utils.clip_grad_norm_(model.parameters(), 1e-6)

            adaptive_clip_grad(model.parameters(), clip_factor=0.0001)

            optimizer.step()

Accuracy выходит 84% , если немного поучить.
Если делать сеть глубже, то обучение не начинается.

700 линейных слоёв начало обучаться при LR, уменьшенном до 3e-5. Сеть достигла accuracy 74% и скорее всего было бы больше, если б дольше обучалась. Видимо обучение не идёт из-за того, что веса сильно шатает и часто не туда.

1000 линейных слоёв начало обучаться при LR = 1e-5 и acc 60%. Если оставить на подольше, то до 84% доходит . C LR=3e-6 уже не обучается. C LR = 9e-6 тоже не обучается. C LR = 2e-5 тоже не обучается.

Предварительный прогрев оптимизатора на около нулевом LR ничего не изменил.
Прогрев на LR=3e-6, на котором ещё не начинается обучение, приводит к тому, что обучение вообще после него не начинается. Хотя можно было бы предположить, что он сэкономит одну эпоху обучения в самом начале.

Увеличение размера батча вдвое до 128 приводит к тому, что перестаёт обучаться при LR=1e-5. Увеличивать LR при увеличении размера батча правильно было бы конечно.

Можно предположить, что LR, на котором происходит обучение находится в очень узких рамках. Что наводит на грустные мысли, ибо если увеличивать глубину сети ещё больше, то мы возможно вообще не сможем обучаться ни с каким LR.

Если подумать, то нам бы в начале по максимому избавиться от любых регуляризаций, любых источников шума, всего, что двигает веса сети не туда. Так на вскидку в голову приходят только две вещи, которые привносят лишний шум: разделение датасета на батчи (т.е. большая стохастичность при градиентом спуске) и разнообразие входных данных. Аугментации у меня пока никакой нет, но можно представить, что многообразие разных начертаний цифр и есть аугментированные цифры. Пробуем из всего датасета взять только один батч и на нём учить. А как обучение начнётся, можно и оставшиеся семплы в обучающую выборку добавлять.

    dataset_train_small = datasets.MNIST('../data', train=True, download=True, transform=transform)

    # повторяем первые 64 семпла вместо всего датасета, т.е. чтобы батч всегда был один и тот же
    dataset_train_small.data = dataset_train_small.data[0:args.batch_size].repeat(len(dataset_train_small.data)//args.batch_size,1,1)
    dataset_train_small.targets  = dataset_train_small.targets[0:args.batch_size].repeat(len(dataset_train_small.targets)//args.batch_size)

На 1000 слоях обучение начинается в самом начале второй эпохи. До этого начиналось в конце.

1500 слоёв. Более точно подобран: negative_slope=2.387 у LeakyReLU. Адапивный клипинг градиентов отключён, ибо вдвое замедлял обучение, и вместо него включил torch.nn.utils.clip_grad_norm_(model.parameters(), 1e-4) . Обучение идёт только на первом батче обучающей выборки. При LR=1e-5 - не обучается. C LR=9e-6 обучается с первой эпохи!

Если не мелочиться, а сразу сделать 10тыс. слоёв, то получается лосс = 14102130520384488865792.000000 и он растёт. Ситуация исправляется, если опять потюнить negative_slope у LeakyReLU в сторону уменьшения с 2.387 до 2.33. Но тогда до первых слоёв доезжают nan-ы вместо градиентов.

Нужно что-то, что не усиливало бы проходящий сигнал при прямом проходе и не гасило бы градиенты при обратном проходе....

Пока подобрал negative_slope=2.376, так, чтобы градиенты до первых слоёв доходили в районе +-1. Надежда на то, что они потекут, и оптимизатор потом не даст им уехать в ноль.

Как порекомендовал Владислав Голощапов, сделал функцию, которая усиливает градиенты на обратном проходе так, чтобы в начале и конце сети они были примерно одного порядка:

# усиливает проходящий через него градиент
class GradAmplifier(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        return input

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output*1.145178 # 10^(разница порядков ослабления/усиления градиентов / количество слоёв между которыми она наблюдается)

class ManyLinearsWithTunedLeakyReLUAndGradAmp(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 16)
        self.blocks = nn.ModuleList([nn.Linear(16, 16) for _ in range(10000)])
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        grad_amp = GradAmplifier.apply

        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = F.leaky_relu(x, negative_slope=2)
        x = grad_amp(x)

        for block in self.blocks:
            x = block(x)
            x = F.leaky_relu(x, negative_slope=2)
            x = grad_amp(x)

        x = self.fc(x)
        return x

Градиенты стали проходить нормально. Сетка пока не обучается почему-то, правда и LR пока не подбирался...

10000 линейных слоёв. Градиенты текут нормально. Заодно проверил mean и variance сигнала на прямом проходе. Там тоже всё хорошо.
Протестил разные LR: 1e-7, 2e-7 ... 9e-7 . По 10 эпох, по 937 итераций оптимизатора в каждой, с одним и тем же батчем размером 64 семпла. Норма градиентов клипается до 1e-4. У bias-ов LR в 1000 раз ниже, чем у weight. Обучение не началось ни с одним из перечисленных LR. Ещё в интервале 1-9e-6 запустил обучаться. Может и 1-9e-8 проверить надо.
Но вангую, что обучение так же не начнётся, ибо попасть на правильный LR будет большой удачей. Если такой LR вообще существует при 10000 слоях.