DuoMem: Towards Capable On-Device Memory Agents via Dual-Space Distillation

Large Language Model (LLM)-based agents can solve complex procedural tasks by interacting with environments over multiple turns, but this ability typically depends on large models, long contexts, and repeated inference calls. This makes advanced memory-augmented agents difficult to deploy on resource-constrained devices. We introduce DuoMem, a dual-space distillation framework that transfers procedural problem-solving ability from a large teacher model to compact student models. DuoMem distils in two complementary spaces: (1)context-space distillation, which replaces student-generated memories with higher-quality teacher-generated procedural memories prepended to the student's input, and (2)parameter-space distillation, which fine-tunes lightweight LoRA adapters on successful teacher trajectories. Evaluated on ALFWorld, a challenging embodied decision-making benchmark, DuoMem boosts a 4B-parameter model from 4.3% to 77.9% task success rate, closing most of the gap to a 72B teacher model (87.1%), while adding fewer than 10M trainable parameters and only a few megabytes of pre-computed teacher memories. Moreover, the DuoMem-enhanced 4B model completes tasks over 3x faster than the 72B teacher in wall-clock time, making it viable for real-time edge deployment, which would be challenging for the teacher.Extensive ablations across eight models spanning 2B-72B parameters reveal that both distillation axes contribute complementary

한 줄 요약

DuoMem은 4B 모델의 ALFWorld 성능을 4.3%에서 77.9%로 향상시키며, 72B 교사 모델과 유사한 성능을 3배 빠르게 달성하는 이중 공간 지도 학습 프레임워크이다.

핵심 기여도

**Context-Space Distillation**: 학습된 4B 모델의 입력에 교사 모델이 생성한 고품질 절차적 메모리를 추가하여 성능 향상.
**Parameter-Space Distillation (LoRA)**: 성공적인 교사 트래젝토리 기반 LoRA 모듈로 학습 파라미터 10M 미만 추가.
**ALFWorld 성능**: 4B 모델의 성공률 4.3% → 77.9% (72B 교사 모델 87.1%에 근접).
**속도 향상**: 4B 모델이 72B 교사 모델보다 3배 빠르게 작업 완료.

핵심 아이디어

DuoMem은 대규모 교사 모델의 절차적 기억 능력을 소규모 학습 모델로 이전하기 위해 **이중 공간 지도 학습**을 제안한다. 첫 번째는 **Context-Space Distillation**로, 학습 모델이 생성하는 메모리를 교사 모델이 생성한 메모리로 대체하여 입력 컨텍스트를 풍부하게 한다. 두 번째는 **Parameter-Space Distillation**로, 교사 트래젝토리에서 추출한 성공 패턴을 기반으로 LoRA 모듈을 미세 조정하여 학습 모델의 파라미터 공간을 최적화한다. 이는 기존의 단일 공간 지도 학습 방식보다 훨씬 효과적인 절차적 추론 능력 전달을 가능하게 한다.

기술적 접근법

**Context-Space Distillation**: 교사 모델이 생성한 메모리가 학습 모델의 입력에 **prepend**되어 제공.
**Parameter-Space Distillation**: **LoRA (Low-Rank Adaptation)** 모듈을 사용해 교사 트래젝토리 기반으로 학습.
**모델 크기**: 4B 학습 모델에 추가 파라미터는 10M 미만.
**메모리 사용량**: pre-computed teacher memories는 수메가바이트 수준.
**데이터셋**: ALFWorld (가정용 태스크 기반, multi-step planning 필요).

주요 결과

**ALFWorld 성능**: 4B 모델의 성공률 4.3% → 77.9% (72B 교사 모델 87.1%에 근접, +73.6% 개선).
**속도**: 4B 모델이 72B 교사 모델보다 3배 빠르게 작업 완료.
**모델 범위**: 2B–72B 파라미터 모델 8개에 대한 실험 수행.
**성능 향상**: 두 지도 학습 방식의 조합이 개별 방식보다 훨씬 높은 성능을 보임.

의의 및 한계

DuoMem은 대규모 모델의 절차적 추론 능력을 소규모 모델로 이전함으로써, **엣지 기기**에서의 실시간 실행 가능성을 열어준다. 특히, 4B 모델이 72B 모델과 유사한 성능을 3배 빠르게 달성한 점은 실용적 가치가 크다. 그러나, 교사 모델이 생성한 메모리에 의존하기 때문에, **동적 환경**에서는 제한이 있을 수 있다. 또한, **다중 모달 환경**에서의 확장성은 아직 검증되지 않았다.

실용적 활용

DuoMem은 스마트폰, 로봇, 임베디드 시스템 등 **자원 제약이 있는 기기**에 적합한 절차적 추론 모델 개발에 활용 가능하다. 특히, **ALFWorld와 유사한 가정용 태스크**나 **다단계 계획이 필요한 산업 자동화 시스템**에서 유용하게 사용될 수 있다.