[논문 리뷰] Think Before You Act: Decision Transformers with Working Memory

Jikun Kang, Romain Laroche외 4명|arXiv (Cornell University)|2023. 05. 24.

Explainable Artificial Intelligence (XAI)인용 수 5

한 줄 요약

이 논문은 내용 기반 가용 메모리 모듈을 사용해 작업별 경험을 명시적으로 저장하고 검색하는 메모리 증강 강화 학습 아키텍처인 Decision Transformers with Memory(DT-Mem)을 제안한다. 저랭크 미세조정(LoRA)을 통합함으로써 DT-Mem은 오직 모델 파라미터의 10%만을 사용하여 아케이드 게임과 Meta-World 작업에서 최신 기술 수준의 일반화 성능을 달성하며, 샘플 효율성과 적응 가능성 면에서 뚜렷한 향상을 이룬다.

ABSTRACT

Decision Transformer-based decision-making agents have shown the ability to generalize across multiple tasks. However, their performance relies on massive data and computation. We argue that this inefficiency stems from the forgetting phenomenon, in which a model memorizes its behaviors in parameters throughout training. As a result, training on a new task may deteriorate the model's performance on previous tasks. In contrast to LLMs' implicit memory mechanism, the human brain utilizes distributed memory storage, which helps manage and organize multiple skills efficiently, mitigating the forgetting phenomenon. Inspired by this, we propose a working memory module to store, blend, and retrieve information for different downstream tasks. Evaluation results show that the proposed method improves training efficiency and generalization in Atari games and Meta-World object manipulation tasks. Moreover, we demonstrate that memory fine-tuning further enhances the adaptability of the proposed architecture.

연구 동기 및 목표

• 암묵적인 파라미터 기반 메모리에 의존하고 치명적인 잊음 현상에 시달리는 대규모 언어 모델 기반 에이전트의 비효율성을 해결하기 위해.
• 인간의 인지 과정을 영감으로 삼아 명시적인 작업 메모리 모델링을 통해 학습 효율성과 다중 작업 간 일반화 능력을 향상시키기 위해.
• 메모리 모듈에 저랭크 미세조정(LoRA)을 적용하여 최소한의 미세조정 데이터로도 새로운 작업에 효율적으로 적응할 수 있도록 하기 위해.
• 기존 방법들에 비해 계산 및 데이터 요구량을 줄이면서도 성능을 유지하거나 향상시키기 위해.
• 결과적으로 파라미터 기반 메모리에 대한 더 해석 가능하고 모듈화된 대안을 제공하기 위해.

제안 방법

• 내부 작업 메모리 모듈을 콘텐츠 기반 주소화 매트릭스로 구현하여 작업별 경험을 저장한다.
• 메모리 업데이트는 입력 시퀀스와 메모리 키 간의 상관관계를 계산하는 어텐션 기반 메커니즘을 통해 수행되며, 주목된 가중치를 사용해 값을 업데이트한다.
• 메모리 검색은 쿼리가 저장된 키와 유사도 기반으로 관련 정보를 위치하고 검색한다.
• 메모리 모듈은 저랭크 미세조정(LoRA)을 통해 미세조정되며, 이는 새로운 작업에 대해 메모리 출력을 조정하는 데 소량의 학습 가능한 파라미터를 적용한다.
• 전체 아키텍처는 일반 지식을 위한 미리 훈련된 Decision Transformer를 활용하고, 미세조정 시에는 오직 메모리 모듈만을 적응시켜 파라미터 효율성을 유지한다.
• 이 시스템은 이중 메커니즘을 사용한다: 새로운 경험을 저장하는 메모리 업데이트와 과거의 작업 관련 정보를 기반으로 결정보정을 수행하는 메모리 검색.

실험 결과

연구 질문

• RQ1 명시적인 내부 작업 메모리가 다중 작업 강화 학습에서 일반화 능력과 샘플 효율성을 향상시키는가?
• RQ2 내용 기반 주소화 메모리 검색이 치명적인 잊음 현상과 성능 측면에서 암묵적인 파라미터 기반 기억보다 어떻게 비교되는가?
• RQ3 저랭크 미세조정(LoRA)이 최소한의 데이터로 새로운 작업에 대해 메모리 모듈을 효율적으로 미세조정할 수 있는가?
• RQ4 제안된 메모리 증강 아키텍처가 MDT, HDT, PDT와 같은 기존 방법보다 아케이드 게임과 Meta-World 벤치마크에서 모두 우월한 성능을 보이는가?
• RQ5 메모리 크기와 미세조정 데이터 양이 모델의 성능 및 적응 가능성에 어떤 영향을 미치는가?

주요 결과

• DT-Mem은 Meta-World ML45 벤치마크에서 HDT, PDT, MDT를 모두 능가하는 최신 기술 수준의 성능을 달성했으며, 미세조정 후 테스트 점수는 0.95±0.10을 기록했다.
• Meta-World에서 DT-Mem은 미세조정 없이도 HDT 대비 8% 향상된 테스트 성능을 기록했고, 미세조정 후에는 3% 향상된 성능를 달성했다.
• 아케이드 게임에서 DT-Mem은 모델 파라미터의 10%만을 사용해 MDT-Top3와 비슷하거나 뛰어난 일반화 성능를 달성했다.
• Pong 데이터셋의 20%를 사용해 200만 스텝 동안 미세조정했을 때, DT-Mem Top3는 MDT-Top3 대비 1.154% 향상된 성능를 기록하여 효과적인 적응 능력을 입증했다.
• 미세조정 데이터 양을 늘림에 따라 성능 저하가 0.978%에서 1.154% 향상으로 전환됨을 확인하여, 더 많은 미세조정 데이터가 일반화 능력을 향상시킨다는 점을 입증했다.
• LoRA 적응을 적용한 메모리 모듈은 단지 147K 파라미터만을 사용했으며, HDT의 147K 적응 파라미터와 230만 개의 하이퍼네트워크 파라미터보다 훨씬 적어 계산 오버헤드를 감소시켰다.