ultra-mem

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}
}

Metadata

Release files for ultra-mem 0.1.2

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for ultra-mem 0.1.2
File	Size	Uploaded
ultra_mem-0.1.2.tar.gz	36.8 MB	Nov 4, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for ultra-mem 0.1.2
File	Interpreter	ABI	Platform	Reset
ultra_mem-0.1.2-py3-none-any.whl 7.0 kB Nov 4, 2025	Python 3	none	any	Details

Total release size: 36.8 MB

Release files / ultra_mem-0.1.2.tar.gz

Download URL	ultra_mem-0.1.2.tar.gz
Size	36.8 MB
Tags	Source
SHA-256 checksum How to use checksums	`b3e8b06d289a1c24cc0c983b9704544eec6c1d2a7d43b2d7d9adb91600fe4b47`
BLAKE2b-256 checksum How to use checksums	`c08eb9db4669f413ac8da7ad7d1a8c6d039e4b95ee460f6589a5b501011af25a`
Upload date	Nov 4, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.9.25`

Release files / ultra_mem-0.1.2-py3-none-any.whl

Download URL	ultra_mem-0.1.2-py3-none-any.whl
Size	7.0 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`a0914e032b78d32ff09adb9e9856c970ad0adf40d69c507dfd5f42e069352841`
BLAKE2b-256 checksum How to use checksums	`8355add0886d9c51395ba343e3825b085a0fb2b52fde7024d61ba50d5ec74ca5`
Upload date	Nov 4, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.9.25`