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

Пробую через self supervised learning (как это делают в DINO https://arxiv.org/pdf/2104.14294 ) учить несколько подряд идущих свёрток. Но не как в DINO - всю сеть целиком, а каждую по отдельности.
Оказалось, что учить всю сеть целиком на пару порядков быстрее, чем отдельно её кусочки.
Оказалось, что так же как и в целиковой сети - сначала учатся последние слои. Видимо входной сигнал сильно сложнее, чем он же, но прошедший через несколько случайно инициализированных свёрток. Изначально у меня было предположение, что будет наоборот или же все начнут учиться одновременно.
Оказалось, что оптимальный LR сильно отличается в каждой свёртке. Но общее в том, что при LR=3e-4 (рекомендованный Andrej Karpathy) обычно начинает обучаться, а уже потом оптимальный LR обычно падает. Оптимальный - этот тот, на котором loss на тренировочной выборке в среднем падает.
Подтвердилось, что градиенты у bias-ов на пару порядков больше, чем у weight и потому для bias-ов LR надо в 100 раз меньше делать. А более общем случае надо конечно разделять свёртки на два слоя: в одном свёртка без bias-а, во втором - только bias для каждого канала.

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

Росс пишет про эволюцию в обучении свёрточных сетей: https://huggingface.co/blog/rwightman/mobilenet-baselines . Как видно, заметную часть составляют изменения того, что никто никогда и не думал менять.

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

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

Всё это легко представить. Берём плоскость и на ней несколько точек, к которым нейросеть должна при обучении максимально приблизиться. И изначально так инициализируем веса, чтобы нейросеть на плоскости рисовала прямую линию (это не обязательно, но для визуализации полезно). И начинаем наше обучение и с очень низким LR. Мы увидим, как наша прямая линия начинает ломаться в нескольких точках и становится ломанной (если активация у нас ReLU). Далее становится очевидным, что чем больше переломов на этой линии, тем проще подогнать веса так, чтобы ломанная проходила через заданные точки на плоскости. Если отдельные наши точки сгруппированы в какой-то одной области, а там переломов недостаточно, то нейросети будет тяжело сжаться именно под эту группу точек. Она уже подогналась под другие и они не пускают её.

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