Инструмент для оперативной диагностики памяти нейросетевых архитектур языковых моделей

Большие языковые модели (Large Language Models, LLM) прошли путь от простых N-граммных систем до современных универсальных архитектур, однако ключевым ограничением остается квадратичная сложность механизма самовнимания по длине входной последовательности. Это существенно увеличивает потребление памяти и вычислительных ресурсов, а с появлением задач, требующих рекордно длинных контекстов, создает необходимость разработки новых архитектурных решений. Поскольку для исследования предлагаемой архитектуры требуется длительное и дорогостоящее обучение полновесной сети, необходимо разработать инструмент, который позволял бы быстро дать предварительную оценку архитектуре с точки зрения внутренней памяти.


В настоящей работе предложен метод количественной оценки внутренней памяти нейросетевых архитектур на основе синтетических тестов, не требующих больших корпусов данных. Под внутренней памятью понимается объем информации, который модель способна воспроизвести без обращения к исходным входам.


Для верификации подхода разработан программный комплекс, апробированный на архитектурах GPT-2 и Mamba. Использованы задачи копирования, инверсии и извлечения значения по ключу. Проведенное сравнение по точности предсказаний, распределению ошибок и вычислительным затратам позволяет оперативно оценивать эффективность и перспективность архитектур LLM.

большие языковые модели, архитектура нейросетей, внутренняя память, долговременное хранение информации, обработка последовательностей, измерение функциональной памяти, сравнение архитектур

1. Kaplan J., McCandlish S., Henighan T., et al. Scaling Laws for Neural Language Models // arXiv preprint arXiv:2001.08361. 2020. https://doi.org/10.48550/arXiv.2001.08361

2. Brown T., Mann B., Ryder N., et al. Language Models are Few‑Shot Learners // Advances in Neural Information Processing Systems. 2020. Vol. 33. P. 1877-1901. https://doi.org/10.5555/3495724.3495883

3. Beltagy I., Peters M. E., Cohan A. Longformer: The Long‑Document Transformer // arXiv preprint arXiv:2004.05150. 2020. https://doi.org/10.48550/arXiv.2004.05150

4. Radford A., Wu J., Child R., Luan D., Amodei D., Sutskever I. Language Models are Unsupervised Multitask Learners // OpenAI. 2019.

5. Common Crawl Foundation. Common Crawl dataset. https://commoncrawl.org

6. Gu A., Goel K., Ré C. Efficiently Modeling Long Sequences with Structured State Spaces // International Conference on Learning Representations (ICLR). 2022.

7. Gao L., Biderman S., Black S., et al. The Pile: An 800 GB Dataset of Diverse Text for Language Modeling // arXiv preprint arXiv:2101.00027. 2020. https://doi.org/10.48550/arXiv.2101.00027

8. Eldan R., Li Y. TinyStories: How Small Can Language Models Be and Still Speak Coherent English? // arXiv preprint arXiv:2305.07759. 2023. https://doi.org/10.48550/arXiv.2305.07759

9. Dao T. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness // Advances in Neural Information Processing Systems (NeurIPS). 2022. Vol. 35. P. 16344-16359. https://doi.org/10.48550/arXiv.2205.14135

10. Gu A., Goel K., Dao T., et al. Mamba: Linear-Time Sequence Modeling with Selective State Spaces // International Conference on Learning Representations (ICLR). 2024.

11. Kwon W., Lee S., Li S., Zaharia M., Zhang H., Stoica I., Sheng Y., Crichton W., Xie S., Gonzalez J. Efficient Memory Management for Large Language Model Inference with KV-Caching // arXiv preprint arXiv:2309.06180. 2023. https://doi.org/10.48550/arXiv.2309.06180

12. Vaswani A., Shazeer N., Parmar N., et al. Attention Is All You Need // Advances in Neural Information Processing Systems (NIPS). 2017. Vol. 30. P. 5998–6008. https://doi.org/10.5555/3295222.3295349

13. Tay Y., Bahri D., Metzler D., et al. Long Range Arena: A Benchmark for Efficient Transformers // arXiv preprint arXiv:2011.04006. 2020. https://doi.org/10.48550/arXiv.2011.04006

14. Bulatov A., Kuratov Y., Burtsev M. Recurrent Memory Transformer // Advances in Neural Information Processing Systems. 2022. Vol. 35. P. 11079-11091. https://doi.org/10.48550/arXiv.2207.06881