PSSA, un modelo de lenguaje no-transformer escrito en Rust desde cero, aprende más rápido y genera 12 veces más rápido que un transformer
PSSA: A non-transformer language model written from scratch in Rust
PSSA es un modelo de lenguaje recurrente de espacio de estados con memoria episódica y pesos plásticos, implementado en Rust sin PyTorch ni TensorFlow. Con 1.5M de parámetros y entrenado en 12.7M de tokens de WikiText-103, alcanza una entropía cruzada de 3.98 frente a 4.43 de un transformer con los mismos parámetros, y genera 200 tokens en 226 ms frente a 2,735 ms en la misma CPU. El autor busca acceso a GPU para escalar el prototipo.
PSSA terminó en 3.98 de entropía cruzada de entrenamiento, el transformer en 4.43. Eso es una brecha de 0.45 nats, perplejidad 53.7 contra 83.7.
- JPLeRouzic
He leído en algún sitio que la arquitectura Transformer tiene un coste cuadrático [0] (lo que explica los altos costes asociados a los LLM y la dificultad de mejorar de forma constante sin recurrir a trucos de tamaño de estado).
Por lo que entiendo, PSSA pertenece a una línea de investigación sobre LLM con arquitectura escalable, porque no necesitas cargar todo el KV en memoria para generar un solo token:
[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_
Por lo visto, ninguno de los que se quejan del uso de Rust leyó más allá del título, porque es la segunda sección del readme y es imposible pasarla por alto.
- janalsncm
OP, no deberías haber escrito esto en Rust. Debería estar en PyTorch, que es con diferencia lo más popular. No podemos saber si esta arquitectura es buena o si hay un problema en tu implementación.
Puedes probar todo gratis en una GPU con Google Colab. Prueba tanto el transformer como tu nueva arquitectura en un conjunto de datos más grande. Algo que ponga la GPU al máximo durante una hora en cada ejecución.
Además, el readme menciona mantener el mismo schedule del optimizador, lo cual suena bien al principio, pero son arquitecturas completamente distintas. La loss es alta en el transformer, ¿probaste a subirle el learning rate?
En general, me interesan las arquitecturas eficientes en parámetros. No creo que los transformers sean óptimos, y de hecho se han hecho muchas mejoras a los transformers vanilla. Pero si tienes una idea para algo mejor, tienes que demostrarlo.
- mllev15
Se nota enseguida cuándo la idea en sí fue generada
- a19486
El colega simplemente tenía tokens que quemar a fin de mes y lo de "en Rust" es puro clickbait para HN.