본문 바로가기
AI/Paper Review

[논문 리뷰] A Generalist Agent (Gato)

by hyeok1235 2024. 8. 1.

0. Abstract

LLM(Large-scale Language Modeling)에서 영감을 얻은 Gato는 동일한 네트워크와 동일한 weight로,

multi-modal (이미지, 텍스트, 음성과 같이 여러 형식의 데이터 처리 가능),

multi-task (다양한 과제 수행 가능),

multi-embodiment (로봇처럼 실제로 물리적으로 구현이 필요한 과제)

를 모두 수행 가능한 generalist agent이다. 

 

1. Introduction

모든 과제에 대해서 하나의 neural sequence 모델을 사용하는 것에는 많은 장점이 있다. 각 과제마다 직접 policy model을 설정하지 않아도 되고, flat sequence로 만들 수 있는 데이터라면 학습에 모두 넣을 수도 있다. 하나의 큰 transformer sequence 모델인 Gato는 학습의 폭을 넓혀 대화, 이미지 캡션, 실제 로봇 팔로 블록을 쌓는 등 여러 개의 과제를 수행할 수 있는 범용 agent이다. 

논문이 증명하고자 한 가설 = 비록 훈련 데이터를 벗어난 모든 과제를 완벽하게 수행할 수 있는 agent는 불가능하지만, 매우 많은 수의 과제를 수행할 수 있는 agent는 가능하다는 것

데이터 양을 늘리고, parameter를 조정하고, training 데이터의 분야를 늘리면 어떠한 과제/행동/구현을 할 수 있는 agent가 가능하다고 예측한다. 이때, 다양한 작업에 대해 가이드 역할로 자연어를 활용한다. 

 

2. Model

Gato의 설계 원칙은 학습 데이터로 최대한 넓은 분야의 관련 데이터들을 사용하는 것이다. Multi-modal을 위해서 데이터들을 모두 flat sequence of token으로 직렬화하게 된다. 배포 시에는 학습한 데이터들을 바탕으로 대화형 반응, 캡션, 버튼 누르기처럼, 주어진 맥락에 맞는 적절한 작업을 수행한다. 

Gato의 Training 과정

 

2.1 Tokenization

데이터를 토큰화하는 방식은 아주 많지만, 해당 논문에서는 Gato의 성능이 제일 높이는 방법들을 선택해서 설명한다.

  • 텍스트 : SentencePiece를 사용해 인코딩되어 raw input과 동일한 순서로 입력
  • 이미지 : ViT와 같이 16x16 패치로 변형한 뒤, 픽셀 단위로 normalize하여 raster order(좌측 상단부터 우측 하단까지 row 단위로 보는 순서)로 입력
  • Tensor나 일련의 데이터들은 모두 row-major 순서로 입력
  • timestep, episode처럼 시간과 관련된 데이터는 token 뒤에 separator를 붙이고 시간 순서대로 입력

2.2 Embedding Input tokens and setting output tokens

토큰화된 데이터들에 parameterized embedding 함수를 적용해서 최종으로 모델에 넣을 input을 구한다. 데이터가 어떤 modal이었는지에 따라서 embedding 함수는 다른 작업을 수행하게 된다.

> 텍스트, 이산형/연속형, 행동 데이터들은 lookup 테이블을 사용해서 학습된 벡터 embedding 공간으로 입력된다.

> 이미지 데이터들은 하나의 ResNet 블록을 사용해서 패치별로 하나의 벡터를 만들어낸다.

 

데이터는 autoregressive(자기 회귀)적으로 계속 들어오기 때문에 각 토큰은 이전 토큰의 타겟 label의 역할을 가지게 된다. 텍스트, 이산형/연속형, 행동 데이터들을 토큰화 이후 바로 타겟으로 설정할 수 있지만, 이미지나 텍스트가 아닌 데이터와 같은 경우에는 아직 Gato가 예측하지는 못한다. 

 

2.3 Training

Gato의 학습 과정에서 핵심적인 아이디어는 앞서 언급되었듯이 크게 두가지이다. 1) 토큰을 토큰 embedding으로 변환해주는 parameterized embedding 함수와 2) 다음 순서의 토큰을 만들어내는 sequence 모델이다. 다른 수행 과제여도 같은 도메인에 속해 있다면 embodiment나 관찰 형식, 동작 방식과 같은 것을 공유할 수 있다. 해당 논문에서는 수행 과제끼리 구별하기 위해서 one-hot task identifier를 제공하기보다는 프롬프트 컨디셔닝을 사용했다.

 

2.4 Deployment

Gato를 배포할 때는 연속적인 행동 모델로 작동하게 된다. 먼저 시범 결과에 대한 내용이 프롬프트로 Gato에 입력이 되고, 환경에 대해서는 입력 정보를 토큰화해서 토큰 하나씩  autoregressive하게  관찰한다. 만약 관찰이 끝났을 때는 수행할 작업에 대한 토큰을 디코딩(토큰화 과정의 역순)한다. 디코딩 된 작업을 실제 환경에서 수행한 뒤에, 변화하는 환경을 다시 관찰하게 되면서 과정을 계속 반복한다. 

 

3. Datasets

Gato는 많고 다양한 데이터셋에 대해서 학습이 된다. 실제 환경과 가상 환경에 대한 정보를 포함하며, 텍스트와 이미지에 대한 데이터셋에 대해서도 학습한다. 

3.1 Simulated control tasks

제어 과제에서는 SOTA나 SOTA에 가까운 강화학습 agent들을 사용해서 데이터셋을 구축했다. 학습 과정에서는 각 환경에서마다 agent가 만들어내는 경험의 집합들(상태, 동작, 보상)을 기록했다.

3.2 Vision and language

텍스트 데이터는 웹페이지, 기사, 책, 코드를 모아둔 MassiveText (Rae et al., 2021)  데이터셋을 사용해서 학습했다. vision-language 데이터셋으로는 ALIGN (Jia et al., 2021) , LTIP (Alayrac et al., 2022), COCO (Chen et al., 2015) 등을 사용했다. 

3.3 Robotics - RGB Stacking Benchmark (real and sim)

Gato는 블록 쌓기 환경으로 로보틱스와 관련된 과제를 평가한다. 이때 가능한 관찰 환경으로는 카메라 이미지, 로봇 팔과 잡는 부분의 각도이다. 평가 기준으로는 1) Skill mastery : 제공된 데이터를 테스트 했을 때의 성능 2) Skill Generalization : 새로운 케이스에 대해서 테스트 했을 때의 성능이 있다. 

 

4. Capabilities of the generalist agent 

Atari 게임에서는 대부분의 경우 인간의 평균 점수보다 높게 나왔다. BabyAI (Chevalier-Boisvert et al., 2018)에서는 거의 모든 레벨에서 80%의 전문가 점수를 받았고, Meta-World (Yu et al., 2020)에서도 대부분의 과제에 대해서 좋은 점수를 받았다. 로봇과 관련된 평가는 RGB 블로 쌓기로 했을 때, 좋은 점수를 받았다. 텍스트 데이터를 바탕으로 챗봇으로 사용했을 때는, 답변들이 대부분 관련이 있는 내용이었지만 너무 피상적이거나 틀린 정보를 제공할 때도 있었다. 이런 오류는 추가적인 Scaling과 수정이 필요하다는 것을 뜻한다.

 

5.Analysis

5.1 Scaling Laws Analysis

모델의 크기가 성능에 미치는 영향을 알아보기 위해서 파라미터의 수를 조절해서 모델을 만들었다. 이때, 파라미터의 수가 늘어날수록(모델이 커질수록) 성능이 더 좋아지는 것을 확인할 수 있었다.

모델 사이즈 scaling laws 결과

5.2 Out of distribution tasks

처음 보는 과제에 대한 성능을 확인하기 위해서 pre-training 데이터셋에서 4개의 과제에 대한 데이터를 제외시킨 뒤에, 해당 과제들에 대해서 성능을 평가하였다. 제어가 가능한 다양한 데이터들에 대해서 학습하며 fine-tuning을 진행했을 때는 좋은 결과를 얻었지만, 그렇지 않은 경우(no control data)에는 특정 환경에 대해서 전이(transfer)가 잘 이루어지지 않았다.

5.3 Fine-tuning on Robotic Stacking Tasks

5.2와 유사하게 fine-tuning이 가능할 때 robotics와 관련된 성능을 확인하고자 했다. 모델 크기가 커질수록 training 과정에서 학습한 데이터를 더 잘 활용하는 것을 확인할 수 있었다. 또한 ,특정 물체에 대한 정보가 없었음에도 불구하고 stacking 과제를 잘 수행한 것으로부터 skill generalization(일반화 능력) 또한 높은 것을 알 수 있다.

5.4 Robotics: Skill Mastery

in-distribution 과제에 대해서는 평균적으로 BC-IMP (Lee ET AL., 2021) agent보다 높은 성능을 보여줬다.

5.5 Specialist single-domain multi-task agents

Gato와 다르게 특정 과제에 대한 데이터로만 학습시킨 agent, generalist가 아닌 specialit agent의 성능 또한 확인한다. Meta-World와 ALE Atari에 대해서도 인간의 평균 점수보다 높게 나오면서, generalist로 학습하지 않고 Gato의 모델 구조 만으로도 특정 과제에 대해서는 SOTA 모델에 도달할 수 있음을 시사한다. 

5.6 Attention Analysis

transformer attention weight을 실제 관찰하는 이미지 위에 렌더링함으로써 각 단계에서 Gato가 어떤 것에 집중하는지를 시각화하였다. 이 부분에서 Gato가 과제와 관련된 영역과 오브젝트에 집중하는 것을 확인한다.

Attention Map

5.7 Embedding Visualization

다양한 과제들을 위해서 embedding, 샘플링, 토큰화하는 과정에 대해서 설명한다. PCA와 T-SNE 같은 기법들을 활용해서 2차원 상에 표시할 수 있도록 dimensionality reduction을 수행한다. 

 

6. Related Work

Gato와 관련된 연구에 대해서 설명하는데, 크게 다음과 같은 분야에 해당된다.

  • Decision Transformer와 Trajectory Transformer
  • Multi-Modal, Multi-Embodiment
  • Generalist 모델
  • Neuroscience 

 

7. Broader Impact

Gato가 배포됨에 따라 영향을 미칠 수 있는 상황에 대해서 설명합니다. Gato와 같은 generalist agent를 악용하는 사례나 방법에 대한 방지책은 비교적으로 연구가 되어 있지 않기 때문에, 실제로 배포되기 전에 이러한 연구가 필요하다고 설명하였다. 또한, 연구자/윤리학자/대중들이 generalist agent에 대해 토론하며 장점과 위험성을 인지하는 것과 Gato에 대해서 지속적인 평가와 모니터링이 필요하다고 이야기한다. 

 

8. Limitations and Future Work

현재 Gato의 context lengthsms 1024개의 토큰으로 제한되어 있기 때문에, 더 효율적인 아키텍쳐를 기반으로 성능 향상을 기대할 수 있다. Gato의 설명 가능성과 성능을 높이기 위해서는 시행하는 실험들을 더 증가시켜야 하고, Counterfactual teaching과 같은 기법을 활용할 수 있다. 오프라인 강화학습을 바탕으로 Gato가 관찰만 가능한 데이터들로부터 다양한 분야에서 성능을 높일 수 있다. 

 

9. Conclusion

Transformer sequence model을 통해 multi-task, multi-embodiment(물리적 구현), 다양한 데이터에 대해서 효과적으로 과제를 수행하는 generalist agent를 만들 수 있었다. 파라미터와 데이터를 확장시킴으로써 일반화 능력이 증가한 Gato의 접근 방식을 바탕으로, 모델과 데이터셋의 크기를 늘려서 더 유용한 general-purpose agent이 만들어질 수 있다는 것을 시사한다. 

 

논문 링크 : 
A Generalist Agent
728x90
반응형