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 for Large Language Models and Agents, particularly in:

• Long-term Memory: Training a memory module to retain the entire training dataset.
• Memory-augmented Architectures: Designing novel model architectures and pre-training paradigms for efficient and adaptive knowledge access.
• Decoupling common and factual knowledge within LLMs

I’m always open to discussions and collaborations — feel free to reach out if you’d like to chat!

🔥 News

• 2026.01 🎉🎉 MLP Memory accepted to ICLR 2026
• 2025.10 🚀🚀 MLP Memory released — model weights and code are now open-sourced.
• 2025.09 🎉🎉 Memory Decoder accepted to NeurIPS 2025
• 2025.06 🏆🏆 My undergraduate thesis was selected as an Outstanding Thesis
• 2024.12 🎉🎉 One paper accepted to IEEE TIFS 2024
• 2024.06 🎉🎉 Beowulf accepted to ACM CCS 2024

📚 Publications

• MLP Memory: A Retriever-Pretrained Memory for Large Language Models
  Rubin Wei*, Jiaqi Cao*, Jiarui Wang, Jushi Kai, Qipeng Guo, Bowen Zhou, Zhouhan Lin
  ICLR 2026
  🔗 arXiv · GitHub · HuggingFace · * Equal Contribution

• 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
  NeurIPS 2025 Poster
  🔗 arXiv · GitHub · HuggingFace · * Equal Contribution

• Beowulf: Mitigating Model Extraction Attacks via Reshaping Decision Regions
  Xueluan Gong, Rubin Wei, Ziyao Wang, Yuchen Sun, Jiawen Peng, Yanjiao Chen, Qian Wang
  ACM CCS 2024
  🔗 ACM CCS

• Augmenting Model Extraction Attacks Against Disruption-Based Defenses
  Xueluan Gong, Shuaike Li, Yanjiao Chen, Mingzhe Li, Rubin Wei, Qian Wang, Kwok-Yan Lam
  IEEE Transactions on Information Forensics & Security (TIFS) 2024
  🔗 IEEE TIFS

🎓 Education

• 2025.09–Present, Ph.D. Student, School of Artificial Intelligence, Shanghai Jiao Tong University
• 2021.09–2025.06, B.E. in Information Security, School of Cyber Science and Engineering, Wuhan University

💼 Internships

• 2024.11–2025.02, TeleAI, Shanghai, China
• 2024.04–2024.06, Shanghai AI Lab, Shanghai, China

🏅 Honors and Awards

• National Scholarship, 2023
• Outstanding Undergraduate Thesis Award, 2025