bingreeky/MemGen

📚 MemGen이란?

MemGen은 ICLR-2026 논문에서 소개된 MemGen 모델을 구현한 연구용 코드베이스입니다. 핵심 아이디어는 외부 데이터베이스가 아닌 모델 자체의 추론 스트림 내에 존재하는 잠재 메모리를 대규모 언어 모델에 제공하는 것입니다. 두 구성 요소가 함께 작동합니다:

메모리 트리거 – 생성 중에 모델이 잠재 메모리에서 가져와야 하는지 결정합니다.
메모리 위버 – 과거 경험을 요약하고 현재 컨텍스트에 통합하는 컴팩트한 "메모리 토큰"을 생성합니다. 이들은 함께 에이전트가 기본 모델 가중치를 변경하지 않고도 상호 작용을 통해 행동을 진화시킬 수 있게 합니다.

⚙️ 시작하기

환경 생성 – requirements.txt를 pip로 설치하거나 제공된 memgen.yml conda 파일을 사용합니다.
(선택 사항) 검색기 설정 – 저장소는 외부 검색을 위해 Search‑R1 환경에 연결할 수 있지만 핵심 MemGen 실험에는 필요하지 않습니다.
모델 및 데이터 세트 선택 – README에는 KodCode, GSM8K, TriviaQA용 즉시 사용 가능한 체크포인트(Qwen2.5‑1.5B‑Instruct 및 SmolLM3‑3B)가 나열되어 있습니다.
스크립트 실행 – 훈련 및 평가는 scripts/weaver_sft.sh(Weaver 모듈용) 또는 scripts/trigger_train.sh(Trigger 모듈용)와 같은 간단한 bash 파일로 래핑됩니다. scripts/eval.sh의 LOAD_MODEL_PATH를 테스트하려는 체크포인트를 가리키도록 조정합니다.

🚀 무엇을 할 수 있나요?

논문 결과 재현 단일 턴 및 다중 턴 추론 작업에서.
잠재 메모리 실험 Weaver 또는 Trigger 구성 요소를 교체하거나 configs/latent_memory/*.yaml 파일을 편집합니다.
바닐라 기준선과 비교 – 저장소에는 표준 generate 함수의 주석 처리된 버전이 포함되어 있습니다. 이를 전환하면 잠재 메모리가 성능을 얼마나 향상시키는지 확인할 수 있습니다.

📦 내부 구성

폴더 / 파일	용도
`memgen/model/modeling_memgen.py`	메모리 토큰을 처리하는 사용자 지정 `generate`가 있는 핵심 모델 클래스.
`scripts/`	훈련(`_sft.sh`, `_grpo.sh`, `trigger_train.sh`) 및 평가(`*_eval.sh`)용 Bash 래퍼.
`configs/latent_memory/`	각 데이터 세트에 대한 YAML 구성. 주의 구현 및 빠른 SFT 모드와 같은 옵션을 제어합니다.
`requirements.txt` / `memgen.yml`	종속성 목록(PyTorch, FlashAttention‑2, HuggingFace Transformers 등).
`assets/memgen.png`	Trigger‑Weaver 파이프라인 다이어그램.

📊 평가 및 체크포인트

README의 표에는 각 모델‑데이터 세트 쌍에 대한 HuggingFace 링크가 나열되어 있습니다. scripts/eval/qwen2_5_gsm8k_sft_trigger.sh와 같은 스크립트는 체크포인트를 다운로드하고 테스트 분할에서 추론을 실행하며 표준 메트릭(예: GSM8K의 정확도, TriviaQA의 정확한 일치)을 보고합니다.

🛠️ 알려진 제한 사항

멀티 GPU 훈련은 분산 데이터 병렬(DDP)에서는 작동하지만 완전 샤딩 데이터 병렬(FSDP)에서는 작동하지 않습니다. 저자는 FAQ에서 이 점을 언급합니다.
다중 턴 GRPO 스크립트와 체크포인트는 아직 보류 중입니다. 현재는 단일 턴 SFT/GRPO만 출시되었습니다.
작은 형식 문제(예: 프롬프트가 \boxed{} 뒤에 마침표 또는 줄 바꿈으로 끝나는지 여부)는 1.5B 모델의 성능에 영향을 줄 수 있습니다 – 최신 코드는 일관된 ChatML 템플릿을 적용하여 이를 수정합니다.

📖 인용

연구에서 MemGen을 사용하는 경우 ICLR‑2026 논문을 인용하세요:

@article{zhang2025memgen,
  title={MemGen: Weaving Generative Latent Memory for Self-Evolving Agents},
  author={Zhang, Guibin and Fu, Muxin and Yan, Shuicheng},
  journal={arXiv preprint arXiv:2509.24704},
  year={2025}
}

🙏 감사의 말

이 프로젝트는 Search‑R1 검색 환경, 이전 잠재 추론 연구(LatentSeek, SoftCoT, R3Mem) 및 LAVIS 코드 프레임워크를 기반으로 합니다.

TL;DR: MemGen은 Trigger‑Weaver 메커니즘을 통해 LLM에 잠재적이고 내부적으로 생성된 메모리를 추가하는 연구 코드베이스로, 에이전트가 가중치 업데이트 없이 시간이 지남에 따라 개선할 수 있게 합니다. 여러 데이터 세트와 소형 및 중형 모델을 위한 즉시 실행 가능한 훈련/평가 스크립트가 포함되어 있습니다.