PUM-Net: Dual-Memory Retrieval for Long-Context State Space Models

Long-context language modeling requires access to details whose relevance may emerge thousands of tokens after they appear. Selective state space models make long-sequence computation efficient, but recurrent compression offers no explicit route to revisit information discarded from the state. The challenge is to retain access to input-specific evidence while enabling its use alongside related corpus knowledge. We propose PUM-Net, a dual-memory extension of Mamba that sepa- rates recurrent computation, causal context retention, and offline corpus reuse. In- ternal retrieval accesses completed input chunks, while external retrieval accesses corpus representations encoded offline. Bidirectional attention updates each re- trieved set using candidates from the other source before pooling and gated inte- gration. This design preserves Mamba’s parallel scan and incorporates retrieved evidence without appending passages to the input. After fine-tuning with 4k-token windows, the corpus-enabled configuration achieves the lowest perplexity in all twelve evaluated domain–length settings through 32k. At 130M, internal retrieval recovers passkeys in 18 of 20 single-trial settings beyond the training length, com- pared with 10 for Mamba. The results support explicitly retained memory as a complement to recurrent compression for prediction and recall over long contexts.