Yiyi Lu
Author directory2026
Mem-LM: An LM-style Model for Reinforced Memory Selection
Yanhan Huang | Yihao Zhong | Wenwen Ouyang | Xing Tian | Zining Yin | Yiyi Lu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing: Industry Track
Yanhan Huang | Yihao Zhong | Wenwen Ouyang | Xing Tian | Zining Yin | Yiyi Lu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing: Industry Track
External memory banks can improve test-time reasoning in frozen large language models without updating backbone parameters. Yet similarity-based retrievers and pointwise selectors score memories independently, ignoring interactions and presentation order. We introduce , an LM-style model for reinforced memory selection that casts context construction as autoregressive index decoding. From a densely retrieved candidate pool, a lightweight Transformer encoder–decoder generates an ordered index sequence [i1,…,ik], modeling dependencies among memories while leaving the backbone frozen. This formulation supports standard language-model training. For sparse binary verifiable rewards, we further propose , a confidence-aware extension of group-relative policy optimization that promotes correct, low-confidence rollouts, penalizes confidently incorrect ones, and retains learning signals in reward-homogeneous groups. Across three mathematical reasoning and two reasoning-intensive question-answering benchmarks with three Qwen backbones, achieves an average score of 84.9 on Qwen-plus, exceeding the strongest memory-augmented baseline by 2.3 points.