Авторы добавляют одномерную depthwise свёртку с ядром k=2 (так понимаю, текущий токен и предыдущий токен) поверх фич, уже прошедших гейт: H_conv = σ(W_dwconv * H_ffn) К теме про такую активацию они, как я понимаю, пришли после изучения абляций, показавших, что последовательное убирание нелинейности из функции активации монотонно уменьшает перформанс на ARC-AGI-1. Что, мне кажется, в целом согласуется с ранжированием упомянутых там функций активации: SwiGLU → SiLU → ReLU, тут вроде ничего нового нет, что SiLU/swish лучше ReLU, а функция с гейтингом ещё лучше (https://t.me/gonzo_ML/4070). Поэтому решили ещё нелинейности подбавить. На картинке, кстати, у них ещё один SiLU заявлен, уже после свёртки, его нет в формуле статьи, но в коде он есть. Провели эксперименты со свёртками разных размеров, для ядра размером 2 оказался лучший результат. Попробовали попереставлять свёртки в разные места трансформера. Внутри механизма внимания позиция мало на что влияет и порой даже ухудшает всё. Лучший результат (на ARC-AGI) если ставить после MLP expansion. Но это вроде не очень соответствует положению на картинке — оно конечно после expansion, но оно уже и после гейта. Может они просто неточто это всё описали. TBPTL нужен для ограничения глубины рекурсии, он считает только градиенты поздних циклов. Внутри TRM и HRM тоже была аналогичная логика: HRM использовал градиенты только с последнего цикла (финальное состояние H модуля и финальное состояние L-модуля), а TRM при deep recursion прогонял внутренний цикл без отслеживания градиентов для всех раз кроме последнего. Да ещё и при самом высокоуровневом deep supervision выходные значения отсоединялись от графа вычислений и передавались на следующий шаг улучшения просто как входные данные. TBPTL делает примерно аналогичное. Если взять модель с D слоями и применять её итеративно в течение M итераций, то новые репрезентации h_t^d слоя d ∈ {1, . . . , D} на итерации t ∈ {1, . . . , M} будут вычисляться как функция от h_t^{d-1} (предыдущий слой той же итерации) и h_{t-1}^d (тот же слой предыдущей итерации). Тут я кстати тоже не уверен, что они это верно написали, эта вот тема с тем же слоем предыдущей итерации какая-то сомнительная имхо. Я это воспринимал как вложенные циклы. Здесь вместо полного бэкпропа через все M итераций мы задаём индекс отсечения N<M, так что для всех шагов от 1 до N бэкпроп не делается, а для N+1 .. M -- делается. Идейно абсолютно та же логика, в лоссе учитываем только последние вычисления. Например, для модели c D=4 слоя и M=8 внутренних циклов (что по идее эквивалентно 32 слоям) при выборе N=2 только 6 последних циклов (t=3..8) повлияют на градиент. Такая конфигурация с 6 из 8 шагов и была выбрана по результатам перебора всех вариантов на ARC-AGI (правда это делалось на двуслойной модели без свёрток, а не на четырёхслойной со свёртками). Эксперименты Авторы взяли те же датасеты и аугментации, что у TRM/HRM (респект авторам оригинальной HRM за то, что дали референсный код, на котором смогли строить все остальные). В TRM (но не HRM) использовалась EMA (модель обучается и обновляет свои параметры, но параллельно этому мы держим другую модель, которая является экспоненциальным скользящим средним от весов обновляемой модели, и на этой модели и делается оценка). Обучали с AdamAtan2 как в оригинальной работе. Weight decay также как в предыдущих работах. Использовалась модель с 4 слоями размерности 512 и с 8 головами. Итого, весь процессинг включает 4 слоя на внутреннем уровне, 8 итераций (из которых только 6 последних участвуют в бэкпропе) и внешний цикл с ACT и максимум 16 шагами. То есть, если я правильно всё понял, как бы 4*8*16=512-слойная модель. Между ACT шагами, как я понимаю, градиенты не передаются, но вот эта часть в статье не описана, надо по коду перепроверять.