Если не мелочиться, а сразу сделать 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, которые дают на одной эпохе плохое значение метрики, но на нескольких эпохах - значительно лучшее. И это никак не узнать.

Вот код скрипта.

Повторил вчерашний эксперимент и добавил данные с графиками туда же на отдельный лист: https://docs.google.com/spreadsheets/d/1PrkH42b10jFQ8ExgRhFixocoS3n8hALLLVwOU7NOLzE/edit?usp=sharing

Обучилось хуже. И ИМХО видны причины такого обучения: это высокий LR у bias-ов свёрток.

Сейчас интервал поиска нового параметра меняется от 1/3 до 3 от старого значения. Возможно стоит этот интервал увеличить, чтобы не происходило подобного залипания на высоком LR, а могло скатиться к маленькому. Хотя тогда графики станут более дёрганными.

Переобучения, кстати, опять не видно. Может, конечно, сетка очень маленькая...

Во время криво настроенных экспериментов выяснилось, что сетка вполне себе может обучаться с learning rate 1e+16 и weight decay 5e+23! Конечно она пришла к таким огромным величинам не сразу, но пришла. И при них вполне себе удерживала accuracy 99.33% на MNIST-е .

Из этого можно сделать вывод, что неверный LR можно компенсировать WD.

Провёл ещё 3 эксперимента с подбором оптимального LR и WD каждую эпоху. Данные и графики добавил на отдельные листы https://docs.google.com/spreadsheets/d/1PrkH42b10jFQ8ExgRhFixocoS3n8hALLLVwOU7NOLzE/edit?usp=sharing .

После каждой эпохи сеть обучалась 100 раз на одну эпоху и выбирались оптимальные значения LR и WD. Т.е. в последние два эксперимента было обучено по 100 тыс. эпох. Эффективнее всего оказалось не смотреть на предыдущие значения LR и WD, как я делал до этого, а выбирать их каждый заново из заранее определённого интервала. Я использовал интервал от 1e-8 до 1e-2. Графики выходят ужасными и пока не понял, как их сгладить в гугл-докс, но зато сеть быстро достигает наилучших значений метрики.

И если посмотреть на графики всех пяти экспериментов, то снова видно полное отсутствие переобучения. Для MNIST-а 1000 эпох вполне достаточно, чтобы оно проявилось хоть как-то.

Было время пообучать сетку в миллион линейных слоёв. Добавил данные по LR=3e-8:

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

lr=1e-7:
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=3e-8:
Test set: Average loss: 1.5202, Accuracy: 5493/10000 (55%)
Test set: Average loss: 1.4754, Accuracy: 5743/10000 (57%)
Test set: Average loss: 1.4550, Accuracy: 5795/10000 (58%)
Test set: Average loss: 1.4516, Accuracy: 5940/10000 (59%)
Test set: Average loss: 1.4456, Accuracy: 5930/10000 (59%)
Test set: Average loss: 1.4228, Accuracy: 6075/10000 (61%)
Test set: Average loss: 1.4007, Accuracy: 6107/10000 (61%)

Подумалось...
L1-регуляризация обычно делается через лос-функцию примерно так:

l1_loss = sum(p.abs().sum() for p in model.parameters())

Потом этот лосс прибавляют с каким-то коэффициентом к тому лосу, который используется для обучения, считают значения градиентов и оптимизатором меняют веса.

Но ведь L1-лос можно взять отдельно и считать его тоже отдельно, например, после обновления весов. Очевидно, что значения градиентов от такого лоса будут равны единице. И можно вообще обойтись без расчёта градиентов, а просто пройтись по всем весам и сделать так:

w -= 1e-6

Точнее с учётом того, что значения весов могут быть отрицательными и меньше 1e-6, то вот так:

w -= w.sign()*min(w.abs(), 1e-6)

где 1e-6 - это аналог описанного выше коэффициента у L1-лоса, помноженного на LR, который использовался бы в оптимизаторе.

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