Rubin Wei

I’m a first-year Ph.D student in Computer Science at the School of Artificial Intelligence, Shanghai Jiao Tong University (SJTU), advised by Prof. Zhouhan Lin in the Language Understanding and Machine Intelligence Algorithms (LUMIA) Group. Previously, I received my B.E. in Information Security from Wuhan University.

My research centers on memory, pretraining, and architectures for large language models and agents.

Long-Term Memory

Training memory modules to internalize and retain knowledge from large-scale training corpora.

Memory-Augmented Architectures

Designing memory-augmented models and pretraining paradigms for efficient and adaptive knowledge access.

Pretraining & Model Architectures

Advancing pretraining methods and model architectures for foundation models.

News

I am currently seeking research internship opportunities with foundation model teams, especially in pretraining.

Intern-S2-Preview technical report released.

Memory Decoder at Scale released.

MemSFT released.

MLP Memory accepted to ICLR 2026.

Memory Decoder accepted to NeurIPS 2025.

Selected Publications

All publications on Google Scholar.

2026

Preprint · 2026 Overview of the MemSFT architecture

MemSFT: Mitigating Alignment Tax with an External Parametric Memory

Jiarui Wang*, Xiang Shi*, Jiaqi Cao, Rubin Wei, Xiquan Wang, Hao Sun, Jingzhi Wang, Zhiqi Yang, Qipeng Guo, Bowen Zhou, Zhouhan Lin

Preprint, 2026

* Equal Contribution

ICLR 2026 Overview of the MLP Memory architecture

MLP Memory: A Retriever-Pretrained Memory for Large Language Models

Rubin Wei*, Jiaqi Cao*, Jiarui Wang, Jushi Kai, Qipeng Guo, Bowen Zhou, Zhouhan Lin

International Conference on Learning Representations (ICLR), 2026

* Equal Contribution

2025

NeurIPS 2025 Overview of the Memory Decoder architecture

Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models

Jiaqi Cao*, Jiarui Wang*, Rubin Wei, Qipeng Guo, Kai Chen, Bowen Zhou, Zhouhan Lin

Annual Conference on Neural Information Processing Systems (NeurIPS), 2025

* Equal Contribution