Продолжение исследования RNN

С прошлой статьи я внёс несколько изменений:
1. Планировщик был сломан и не изменял скорость. Починил.
2. Остаточное соединение через умножение.
3. WindowedDense для выходной проекции.
4. Добавил clipnorm 1, cutoff_rate 0.4

Как обычно это всё добавляет стабильности и 1% точности.

WindowedDense по неизвестной мне причине добавляет SMR стабильность.

class SMR(layers.Layer):
  def __init__(self, units):
    super().__init__()
    self.state_size = units
    self.s_l = layers.Dense(units, use_bias=False)

  def get_in_proj(self):
    return WindowedDense(self.state_size, 16)

  def call(self, i, states):
    s = states[0]
    s = self.s_l(s)
    o = i * (s + 0.1)
    return o, [o]

Обновлённые тесты для "crime-and-punishment-2554" (1128 шагов):

Bigram: ~0.27

LSTM
0.4951, 0.5609, 0.5795, 0.5886

GRU
0.5036, 0.5654, 0.5834, 0.5932

SMR
0.5335, 0.5843, 0.6002, 0.6103

SLR(emb=3072, lr=0.001)
0.5267, 0.5799, 0.6021, 0.6151

В дальнейшем эксперименты будут на "Принцип неопределённости" (4356 шагов).
Русский фанфик по Звёздным Войнам. Намного более сложный текст.
Размер моделей будет ~500т.

Bigram: 0.2449

LSTM
0.5047, 0.5530, 0.5642, 0.5673

GRU
0.5006, 0.5458, 0.5611, 0.5652

SMR
0.5156, 0.5588, 0.5716, 0.5750

SLR(emb=3072, lr=0.001)
0.4849, 0.5279, 0.5398, 0.5423

LSTM и GRU опять "обделались". Посла множества экспериментов у меня появился список правил построения эффективных RNN:

1. Для стека ячеек необходима входная матрица.

2. Матрица для входа и/или состояния повышает стабильность.

3. Остаточное соединение через умножение лучше сложения.
   Объясняю это тем что так легче масштабировать и инвертировать значения.

4. Минимальная глубина вычислений ячейки. Для стабилизации градиентов.

5. Использование активаторов-зажимов(TanH, Sigmoid, Softmax) приводит к угасанию градиентов. TanH наименее вредный.

6. Внутренние нормализации - плохая практика.

7. Отдельное состояние не обязательно.

8. Нелинейности нужны только если задача требует.

9. Языковое моделирование не требует активаторов и биасов.

10. Проблемы градиентов первичны.

11. Если состояние используется только в остаточном соединении то RNN можно распараллелить. Вот так: h = x + h-1;

12. Количество обучаемых параметров и размер состояния должны быть сбалансированными.

13. Обучение с сохранением состояния даёт мощную регуляризацию.

14. Ценность эмбеддинга зависит от архитектуры: Есть матрица для состояния - минимальная. Есть матрица для входа - средняя. Матриц нет - максимальная.

15. Нормализацию лучше делать перед слоем и остаточным соединением: x = norm(x); y = RNN(x); x = x * y;

16. Ворота в разы повышают эффективность состояния. Но проигрывают из за ухудшения градиентов и баланса параметров/состояния. Возможно это нужно рассматривать как смесь экспертов.

17. Чем больше опора на вход тем больше нужда в стеке или многослойности.

18. Ортогональная инициализация не обязательна.

Список "ессесвенно" не исчерпывающий. Общий его посыл в том что SMR лучший.
Но я таки придумал архитектуру ещё лучше:

class MSMR(Layer):
  def __init__(self, units, cells=3):
    super().__init__()
    self.units = units
    self.state_size = [units, cells * units]
    self.mem_shape = [-1, cells, units]
    self.k_l = Dense(cells, use_bias=False)
    self.d_l = Dense(units, use_bias=False)

  def get_in_proj(self):
    return Dense(self.units, use_bias=False)

  def call(self, i, states):
    s, m = states
    m = tf.reshape(m, self.mem_shape)

    k = tf.nn.softmax(self.k_l(i * s), axis=-1)
    k = tf.expand_dims(k, -1)
    d = tf.reduce_sum(m * k, axis=-2)

    o = i * (self.d_l(d) + 0.1)

    k = tf.tanh(self.k_l(o))
    k = tf.expand_dims(k, -1)
    w = tf.expand_dims(o, -2)
    m = tf.tanh(m * k + w * (1 - k))

    m = tf.reshape(m, [-1, self.state_size[1]])
    return o, [o, m]

  0.5187, 0.5705, 0.5860, 0.5902

По сути это универсальная обёртка для расширения памяти RNN.

ЗЫ:
Я пытался сравнивать с другими RNN и трансформерами. GPT2, Mamba, RWKV, Gemma2, ...
Все они показали сомнительные результаты. С ними вообще сложно сравнивать. Это принципиально другие архитектуры. Похоже я близок к пределу точность/шаги.
За исключением семейств SSM и RWKV все RNN вертятся вокруг ворот LSTM/GRU и не предлагают ничего нового.

ЗЫЫ:
В моих экспериментах с трансформерами линейное кодирование позиций значительно превзошло синусоидальное.

p_emb = tf.cast(tf.range(0, 1, 1 / seq_len), t_emb.dtype)
p_emb = tf.expand_dims(tf.expand_dims(p_emb, 0), -1)
p_emb = tf.tile(p_emb, [batch_size, 1, 1])
x = tf.concat([t_emb, p_emb], -1)