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

49a64801c7d2f010667da796c36f0f5f

С прошлой статьи я внёс несколько изменений:
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)

© Habrahabr.ru