ultra-mem 0.1.2

UltraMem

UltraMem

Implementation of UltraMem, improved Product Key Memory design, from Bytedance AI labs

Install

$ pip install ultra-mem

Usage

import torch
from ultra_mem import UltraMem

ultra_mem = UltraMem(
    dim = 512,
    core_heads = 2,
    topk = 32,
)

tokens = torch.randn(1, 1024, 512)

out, mem_indices, aux_loss = ultra_mem(tokens) # (1, 1024, 512), (2, 1, 1024, 32), ()

Char-level LM

$ uv run train.py

Citations

@misc{huang2025ultrasparsememorynetwork,
    title   = {Ultra-Sparse Memory Network}, 
    author  = {Zihao Huang and Qiyang Min and Hongzhi Huang and Defa Zhu and Yutao Zeng and Ran Guo and Xun Zhou},
    year    = {2025},
    eprint  = {2411.12364},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2411.12364}, 
}

@inproceedings{Lin2025ContinualLV,
    title   = {Continual Learning via Sparse Memory Finetuning},
    author  = {Jessy Lin and Luke S. Zettlemoyer and Gargi Ghosh and Wen-tau Yih and Aram H. Markosyan and Vincent-Pierre Berges and Barlas Ouguz},
    year    = {2025},
    url     = {https://api.semanticscholar.org/CorpusID:282203348}
}