Причина, почему 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 слоях.

Чуда не случилось. Проверил LR так же 1e-8, 2e-8 ... 9e-8 и 1e-6, 2e-6 ... 9e-6 . Обучение не началось.

Пока в голову приходит обучать по 100 слоёв, добавляя их по мере того, как прошлые 100 обучились хотя бы до 50% accuracy. Или, как как-то писал Владислав Голощапов - прокинуть скипконекшны до каждого 100-ого слоя, а потом их постепенно убрать. Но это всё те хаки, которые пока хочется избежать...

Скрипт для обучения нейросети из 1 млн. линейных слоёв. Итог того, что было начато тут: https://t.me/chivinya/11

Некоторое время назад разбирался с Self Supervised Learning (SSL). И постепенно наткнулся на DINO (https://github.com/facebookresearch/dino , https://arxiv.org/pdf/2104.14294.pdf ) . Там описан оригинальный механизм обучения любой сети без разметки и не требующий дополнительной сети, как в BYOL, например. Пока мне кажется DINO самым простым и верным подходом. И ещё он в топе https://paperswithcode.com/sota/self-supervised-image-classification-on . DINO v2 - это DINO с несколькими твиками не меняющих идею.

Есть две сети одинаковой архитектуры: студент и учитель. Учитель дистилирует в студента через softmax + CE. Студент обучается градиентным спуском. На вход подаются разные кусочки изображения, аугментированные и не очень. Выходы сетей трансформируются так, что студенту сложно приблизиться к тому, что выдаёт учитель. Из-за этого их выходы не колапсируют в константу и обучение идёт. При этом веса учителя обновляются через Exponential Moving Average (EMA) от весов студента. Т.е. сети очень похожи, но из-за EMA учитель чуточку лучше студента. И этого "чуточку" достаточного для очень медленного выучивания скрытых представлений данных, подаваемых на вход.

Вот ещё описание того, как работает DINO: https://t.me/gonzo_ML/688 , https://deepschool-pro.notion.site/DINO-Self-distilation-with-no-labels-18e894250fae4e5d87e99a3312c1effd .
Там же есть упрощённый код алгоритма обучения DINO, который снимает основные вопросы по тому, как оно всё работает.

Я попробовал воспроизвести алгоритм и после нескольких попыток оно заработало. И, как оказалось, многие описанные в работе трюки не сильно нужны и без них модели обучаются. Может хуже, но тем не менее. Так, например, аугментация и кропы не нужны. Это сразу позволяет учиться на любых данных, а не только картинках. Вместо аугментации можно зашумлять входные данные и это немного улучшает итоговую модель. Размер выходного вектора вроде как не очень важен. Очень важны гиперпараметры.

Похоже, что есть ещё какой-то смысл в этих трансформациях (sofmax + centering + sharpening + CE), который заставляет студента выучивать репрезентации того, что поступает на вход, а не шатать свои веса в хаотичном порядке. Предположу, что всё похоже на классификацию, и тут учитель сильнее выпячивает какой-то правильный с его точки "класс", а студент пытается сделать тоже самое. Наверное можно придумать не приводящие к колапсу трансформации и для других задач, а не только для классификации.

В случае SSL много времени требует валидация. Мы хотим понять, насколько хороши выученные репрезентации. Поэтому надо взять сеть (студента или учителя), прицепить к неё линейный слой, обучить на трейне только его, а потом посмотреть значение метрики на проверочной выборке. На простых датасетах, типа MNIST-а, работает хорошо. Попробовал ещё на SVHN ( http://ufldl.stanford.edu/housenumbers/ ). Это тоже цифры, но цветные 32х32 пикселя, с огромным количеством шума, но зато реальные данные и их довольно много: 600 тыс.. Тоже обучается, но как-то о-о-очень медленно. И accuracy значительно ниже, что ожидаемо.

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

Я к тому, что баланс между оптимизатором и WD чаще всего не наступает в обученных сетях. И WD скорее всего можно шатать в очень больших пределах. В начале обучения это излишне, а как обучение выходит на плато - возможно временное увеличение WD сильно поможет.

На праздниках запускал обучение сетки из 1 млн. линейных слоёв. Прошло 3 эпохи и результаты такие:

Test set: Average loss: 1.5966, Accuracy: 5306/10000 (53%)
Test set: Average loss: 1.5637, Accuracy: 5415/10000 (54%)
Test set: Average loss: 1.5002, Accuracy: 5640/10000 (56%)

LR был 1e-7 . Наверное надо было взять его побольше...

В https://habr.com/ru/companies/airi/articles/816125/ мельком описан интересный лосс: он стремится так изменить веса, чтобы вход и выход совпадали. Это регуляризация ИМХО позволяет убрать шум от случайной инициализации, оставляя только те веса, которые реально нужны для обучения. При этом может применяться к любым слоям/блокам, ничего не зная об их структуре.

Запускал обучение сетки из 1 млн. линейных слоёв. Взял LR побольше: 3e-7. результаты сильно хуже:

Test set: Average loss: 1.8179, Accuracy: 4412/10000 (44%)
Test set: Average loss: 1.7524, Accuracy: 4307/10000 (43%)
Test set: Average loss: 1.7991, Accuracy: 3730/10000 (37%)
Test set: Average loss: 1.7636, Accuracy: 3763/10000 (38%)
Test set: Average loss: 1.8692, Accuracy: 2959/10000 (30%)

Решил воспользоваться идеей @kraidiky по периодическому поиску оптимальной скорости обучения. Только искать не один learning rate, но ещё и weight decay. Причём отдельно для свёрток и отдельно для MLP, а также отдельно для весов и отдельно для сдвигов. Получилось 4 пары LR и WD. В пределе можно было бы искать их для каждого слоя, но у меня сетка примитивная. Прототипом в очередной раз послужил код из примеров Торча: https://github.com/pytorch/examples/tree/main/mnist . Если брать сети вроде ResNet-ов и EfficientNet-ов, то там очень большое пространство подбора и надо подумать эффективно его делать. Возможно внутри блока свёрток будут разные LR и WD, и возможно они будут разными у каждой группы блоков, работающих с одним разрешением.

Исходный скрипт имеет оптимизатор Adadelta, который судя по его описанию сам подбирает оптимальный LR для каждого веса. С другими учится сильно хуже. Проблема данной архитектуры сети в том, что для MLP нужен LR раз в 10-20 больше, чем для свёрток. Так происходит потому, что архитектура сети нестандартная, в ней нет пулинга полностью схлопывающего ширину и высоту:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, kernel_size=2)
        x = self.dropout1(x)
        x = torch.flatten(x, start_dim=1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return x

Перед каждой эпохой я добавил поиск оптимальных learning rate и weight decay. Спомощью Optuna я обучал 30 сетей, клонированных из текущей, и подбирал 4 пары LR и WD, которые потом использовал для обучения одной эпохи. Всего было 100 эпох. На MNIST-е отработало быстро.

Исходный код давал accuracy 99.25% на 14 эпохах. Если учить дольше, то получалось иногда 99.3%.
С подбором оптимальных LR и WD после 66-ой эпохе получилось достичь 99.43%, что не мало. И само обучение шло в начале быстрее.

Вот таблица с данными, полученными при обучении 100 эпох. Там же есть разные графики. https://docs.google.com/spreadsheets/d/1PrkH42b10jFQ8ExgRhFixocoS3n8hALLLVwOU7NOLzE/edit?usp=sharing .

Ещё один важный момент: сеть не переобучается. Т.е. фактически переобучение происходит тогда, когда мы используем неподходящие LR и WD. Возможно в этом утверждении я не прав, но на графиках переобучения не видно.

У подхода есть и минусы. Возможно существуют такие LR и WD, которые дают на одной эпохе плохое значение метрики, но на нескольких эпохах - значительно лучшее. И это никак не узнать.