UMA-LedgerQA-4B (Specialist)

UMA-LedgerQA-4B is the Specialist checkpoint of the Unified Memory Agent (UMA) introduced in Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning.

UMA is a tool-using memory agent that incrementally maintains a compact core summary and a structured key-value Memory Bank. The same policy performs memory construction and downstream question answering through explicit memory and retrieval operations.

Code: github.com/ictnlp/unified-memory-agent
Paper: arXiv:2602.18493
Generalist checkpoint: ICTNLP/UMA-4B

Checkpoint Variant

This repository contains the Specialist UMA checkpoint adapted to Ledger-QA, the paper's controlled benchmark for persistent state tracking under accumulated updates, overwrites, and aggregation.

Property	Value
Base model	`Qwen/Qwen3-4B-Instruct-2507`
Parameters	4B
Weight format	BF16 Safetensors
Training method	End-to-end reinforcement learning with Task-Stratified GRPO
Specialization data	Ledger-QA
Reported default context budget	16K

The Generalist and Specialist checkpoints share the same UMA architecture and tool interface. Use this checkpoint to reproduce or extend the Ledger-QA Specialist setting; use ICTNLP/UMA-4B for the broader Generalist setting.

Intended Use

This checkpoint is intended for research on:

persistent state tracking over long input streams;
updates, overwrites, contradiction handling, and aggregation;
proactive structured memory construction;
memory maintenance with explicit tool calls;
evaluation and extension of the UMA framework and Ledger-QA.

The checkpoint is designed to run inside the UMA two-phase agent loop. A plain text-generation call loads the language model, but does not by itself instantiate the Memory Bank, retrieval tools, prompts, or memory-to-QA workflow.

Loading the Weights

python

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "ICTNLP/UMA-LedgerQA-4B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

The model can also be served through an OpenAI-compatible inference server:

bash

vllm serve ICTNLP/UMA-LedgerQA-4B \
  --max-model-len 16384 \
  --gpu-memory-utilization 0.8

For full memory-agent inference, including the Memory Bank, memory tools, embedding retrieval, prompts, Ledger-QA data preparation, and evaluation, follow the official repository.

Limitations

This checkpoint is specialized for Ledger-QA and should not be presented as the paper's Generalist checkpoint.
Ledger-QA is a controlled diagnostic benchmark rather than comprehensive real-world validation.
The checkpoint is a research model and may generate incorrect answers or perform incorrect memory updates.
Agent behavior depends on the UMA prompt templates, tool implementations, retrieval backend, chunking policy, and inference configuration.
Persistent-memory applications can involve sensitive information. Deployments should provide appropriate privacy controls, retention policies, and user oversight.
This checkpoint should not be used as the sole basis for high-stakes decisions.

Citation

bibtex

@article{zhang2026learning,
  title   = {Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning},
  author  = {Zhang, Kehao and Gui, Shangtong and Yang, Sheng and Chen, Wei and Feng, Yang},
  journal = {arXiv preprint arXiv:2602.18493},
  year    = {2026}
}