PSSA на Rust учится быстрее трансформера и генерирует текст в 12 раз быстрее
PSSA: A non-transformer language model written from scratch in Rust
PSSA — это языковая модель, написанная с нуля на Rust без PyTorch или TensorFlow. В ней нет трансформерной архитектуры: вместо внимания используется рекуррентный state-space слой, банк эпизодической памяти на 512 слотов и пластичные веса, которые обновляются во время работы. На 12,7 млн токенов WikiText-103 при 1,5 млн параметров PSSA достигает кросс-энтропии 3,98 против 4,43 у трансформера, а на невиданном тексте разрыв сохраняется. Генерация 200 токенов занимает 226 мс против 2735 мс на том же CPU.
PSSA закончила с кросс-энтропией 3.98, трансформер — с 4.43. Это разрыв в 0.45 нат, перплексия 53.7 против 83.7. Трансформер потратил весь свой бюджет в 12.7M токенов, чтобы достичь потери, которую PSSA уже прошла примерно на 2M токенов.
- JPLeRouzic
Я где-то читал, что архитектура Transformer имеет квадратичную стоимость [0] (что объясняет высокие затраты, связанные с LLM, и сложность постоянного улучшения без карманов состояния).
Насколько я понимаю, PSSA относится к линии исследований LLM с масштабируемой архитектурой, потому что для генерации одного токена не нужно загружать весь KV в память:
[0] https://aclanthology.org/2023.findings-emnlp.936/
https://arxiv.org/abs/2503.00392
https://papers.nips.cc/paper_files/paper/2023/hash/6ceefa7b1...
- prospero_
Похоже, никто из тех, кто жалуется на rust, не читал дальше заголовка, потому что это второй раздел readme и его невозможно пропустить.
- janalsncm
OP, не стоило писать это на Rust. Надо было на PyTorch, который на сегодняшний день самый популярный. Мы не можем понять, хороша ли эта архитектура или в вашей реализации есть проблема.
Вы можете бесплатно протестировать всё это на GPU с помощью Google Colab. Протестируйте и transformer, и вашу новую архитектуру на большем датасете. На таком, который загружает GPU на час при каждом запуске.
Кроме того, в readme упоминается, что вы сохранили тот же график оптимизатора — на первый взгляд звучит хорошо, но это совершенно разные архитектуры. У transformer высокий loss, вы пробовали повысить learning rate для него?
В целом мне интересны параметроэффективные архитектуры. Я не считаю transformers оптимальными, и действительно, в vanilla transformers было внесено много улучшений. Но если у вас есть идея получше, нужно это показать.
- mllev15
Просто сразу видно, когда идея сама себя сгенерировала
- a19486
Чувак просто решил сжечь немного токенов в конце месяца, а «на Rust» — это чистый кликбейт для HN.