LSWM 2.0: нейросеть для создания 3D моделей обходит LSTM с меньшим числом параметров

Автор LSWM опубликовал первую версию сети несколько дней назад, но при повторном тестировании обнаружил критические проблемы: высокую чувствительность к начальному значению генератора случайных чисел (сиду), избыточные линейные слои и отдельный LWM-слой, который вредил стабильности. Версия 2.0 убирает эти узкие места и перестраивает внутреннюю логику ячейки. Исходный код выложен автором в открытый доступ.

Что изменилось в архитектуре?

Главное нововведение: входной сигнал проходит через три линейных слоя (запрос, ключ, значение), как в трансформере или xLSTM. Но, в отличие от трансформера, каждый из этих компонентов поэлементно умножается на предыдущее скрытое состояние. Результат: запрос, ключ и значение становятся зависимыми от прошлой истории последовательности.

Кроме того, автор:

• Заменил два линейных слоя одним, что сократило число параметров при росте качества
• Полностью убрал LWM-слой, который делал сеть чувствительной к сиду
• Перешёл на суммирование вместо конкатенирования (объединения данных в одну длинную строку)
• Использовал softsign (функция активации, плавно сжимающая числа в диапазон от минус единицы до единицы) и scaled softsign для нормализации

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

Что понадобится

• Python 3.x с установленными библиотеками PyTorch, NumPy
• GPU с поддержкой CUDA (автор запускает модель на видеокарте)
• Собственный датасет или тестовый набор автора (описан в статье)
• 30 минут на запуск первого эксперимента с размером скрытого слоя 128

Пошаговая инструкция

1. Скопируйте код модели. Автор выложил полный исходный код класса LSWM. Ключевой фрагмент (определение модели):

import torch
import torch.nn as nn
import torch.nn.functional as F

device = torch.device("cuda")

class LSWM(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.d = d_model
        self.embedding = nn.Embedding(vocab_size, d_model).to(device)
        self.W_f = nn.Linear(d_model, d_model).to(device)
        self.W_i = nn.Linear(d_model, d_model).to(device)
        self.W_o = nn.Linear(d_model, d_model).to(device)
        with torch.no_grad():
            self.W_f.bias.fill_(3.0)
            self.W_i.bias.fill_(0.0)
            self.W_o.bias.fill_(0.0)
        self.W_q = nn.Linear(d_model, d_model).to(device)
        self.W_k = nn.Linear(d_model, d_model).to(device)
        self.W_v = nn.Linear(d_model, d_model).to(device)
        self.norm = nn.LayerNorm(d_model)

1. Задайте размер скрытого слоя. Автор проверял четыре конфигурации:

2. hidden dim = 128, 400 эпох

3. hidden dim = 2048, 1200 эпох
4. hidden dim = 4096, 1500 эпох
5. hidden dim = 128, 20000 эпох

Для первого запуска начните с hidden dim = 128 и 400 эпох.

1. Обучите модель на своём датасете. Формат обучения совпадает с описанным автором в предыдущей статье: стандартный цикл PyTorch с оптимизатором и функцией потерь.

2. Запустите инференс (инференс, применение обученной модели к новым данным) на длинной последовательности. Автор тестировал на цепочке из 1000 элементов с «иголкой» на позиции 500:

model.eval()
with torch.no_grad():
    tokens_pool = [a, b, c]
    random_noise = []
    for _ in range(1000):
        if _ == 500:
            random_noise.extend([a])
        random_noise.extend([random.choice(tokens_pool), TOKEN_ARROW])
    chain1 = random_noise + [c, TOKEN_Q_1]
    chain2 = random_noise + [c, TOKEN_Q_2]
    test1 = torch.tensor([chain1]).to(device)
    test2 = torch.tensor([chain2]).to(device)
    pred_live = torch.argmax(model(test1), dim=1).item()
    pred_work = torch.argmax(model(test2), dim=1).item()
    print(f"need: 3 answer: {pred_live}")
    print(f"need: 12 answer: {pred_work}")

1. Сравните результаты с LSTM и GRU. Автор запускал каждую сеть 10 раз с одинаковым числом параметров и считал долю верных ответов. По его результатам, LSWM 2.0 дала больше верных ответов, чем LSTM (второе место) и GRU (худший результат). Точность у всех моделей составляла от 93% до 100% на обучающей выборке.

Что это даёт на практике, по ролям?

Для разработчиков и ML-энтузиастов. Нейросеть для создания 3D моделей и других задач с длинными последовательностями, работающая на рекуррентной архитектуре вместо трансформера. Меньше параметров, значит, можно экспериментировать на менее мощном железе. Код открыт, можно форкнуть и дообучить (дообучение, обучение модели на своих примерах под узкую задачу) под свою задачу.

Для авторов Дзена и контент-практиков. Если вы пишете про нейросети, это пример того, как одиночный исследователь выкладывает архитектуру, тестирует её публично и честно фиксирует провалы. Это рабочий формат для технического контента: не «я создал прорыв», а «вот что сломалось и вот как починил».

Для предпринимателей РФ и СНГ. Прямой коммерческой применимости у LSWM 2.0 пока нет: это исследовательская работа на раннем этапе. Но сам подход, рекуррентная сеть с Q/K/V вместо полного трансформера, интересен тем, что требует меньше вычислительных ресурсов. Если работаете с последовательными данными (временные ряды, логи, диалоги), стоит отслеживать, как развивается эта архитектура.

Идея засунуть Q, K, V в рекуррентную ячейку и сделать каждый шаг зависимым от полной предыдущей истории оказалась рабочей. Вопрос в том, сможет ли один исследователь довести её до уровня, где это будет конкурировать не только с GRU на синтетических тестах, но и на реальных задачах. Код открыт, и это главное: любой может проверить.