서론

현재, 과거 상태 및 행동을 고려하여 실시간으로 적절한 행동을 선택하는 인공지능은 다양한 응용 분야에서 중요한 역할을 한다. 이러한 인공지능을 만드는 방법 중 하나로 Online Decision Transformer(ODT)가 주목받고 있다. 그러나 ODT와 같은 Transformer 기반 모델은 긴 기억을 요구하는 작업에 적합하지 않은데, 이는 추론 비용이 $O(context\_length^2)$으로 높기 때문이다. 하지만 Retentive Networks(RetNet)는 $O(context\_length)$에 불과한 추론 비용과 낮은 메모리 사용량을 보인다. 따라서 필자는 이러한 단점을 극복하기 위해 ODT를 Retentive Networks(RetNet)로 대체하는 Online Decision Retentive Networks(ODRN)를 고려했지만, 1. 소수의 벡터로 관측치를 표현할 수 있을 것이라는 판단과 2. context length를 늘리는 것만으로 좋은 agent를 만들 수 없기에 ODRN을 시도하지 않고 표현 학습에 집중하기로 결정했다.

배경

Decision Transformer

Decision Transformer(DT)는 미래의 Return To Go와 과거 및 현재의 상태 그리고 행동을 고려하여 의사결정 과정을 수행한다. 기존 강화학습 알고리즘과는 다르게 sequence modeling을 통해 문제를 해결하며, 모델 크기를 키우기가 용이하고 안정적인 학습이 가능한 장점이 있다. 그러나 긴 기억을 필요로 하는 작업에는 트랜스포머의 높은 추론 비용이 문제가 되며 behavioral cloning 대비 큰 모델 capacity를 요구한다는 단점이 있다.

Online Decision Transformer

Online Decision Transformer(ODT)는 오프라인 데이터셋만으로 policy를 학습한 DT와 다르게, 환경과의 상호작용을 통해 policy를 미세조정 할 수 있도록 개선한 모델이다. 비용 함수에 policy의 엔트로피를 추가함으로써 exploration을 유도하여 온라인 미세조정 성능을 높인다. 또한 entropy term이 더해진 maximum likelihood가 loss function이기에 온라인 미세 조정 또한 지도학습에서와 같이 안정적이라는 장점이 있다.

Retentive Networks

트랜스포머와 비슷하거나 나은 성능을 보이면서 수열의 길이에 대해 낮은 추론 비용을 갖는 네트워크는 많이 제시되고 있다( Linformer, Infinity Memory Transformer, S3, S4, HIPPO, RetNet). 그중에서 RetNet은 트랜스포머의 attention을 retention으로 바꿈으로써 수열을 병렬적으로 처리할 수 있고 RNN처럼 recurrent하게 처리할 수도 있다. Recurrent하게 추론할 때, context length와 무관한 계산량을 갖는다는 강점이 있다. 다만 "기억"의 역할을 하는 상태 매트릭스의 차원이 임베딩 차원의 제곱이기 때문에, 더 많은 정보를 기억하기 원한다면 임베딩 차원을 늘려야하고 이에 따라 계산량이 늘어난다는 문제를 갖고 있다. 또한 모델의 크기가 작을 때 트랜스포머보다 언어 모델링 성능이 떨어지는 것으로 알려져 있다.

Scene Memory Transformer

Scene Memory Transformer는 partially observed long horizon tasks를 수행하는 agent가, 긴 기억을 가질 수 있도록 하는 모델이다. 긴 기억을 가지기 위해 긴 context length가 있어야 하는데, Vision Transformer와 같이 한 개의 이미지를 여러 개의 벡터로 표현하면 필요한 context length 또한 같이 커지는 문제가 있다. 그렇기에, 한 개의 vector로 하나의 observation을 표현함으로써 짧은 context length로 긴 메모리를 구현하는 방식을 선택한 모델이다.

왜 ODRN이 기대되지 않는가?

다수의 벡터로 한 개의 observation을 표현할 필요는 없다

처음에는 ViT와 같이 한 개의 image를 다수의 벡터로 표현해야 한다고 생각했다. 256개의 벡터로 이미지를 표현하고 1초에 4개의 이미지만 받는다고 해도 1분의 기억을 갖기위해 context length는 $256 * 4* 60 = 61440$이 돼야한다. 이러한 context length를 갖기 위해서, inference cost가 context length와 무관한 네트워크가 필요하리라 생각했다. 하지만 생각을 하면 할수록, 이것이 틀릴 수도 있다는 사실을 인정해야만 했다.

많은 경우 Scene Memory Transformer와 같이 하나의 observation을 하나의 벡터로 표현한다면 트랜스포머로도 충분히 긴 메모리를 가질 수 있으리라 생각한다. 예시로 새로운 집에서 물건을 재배치하는 로봇을 생각해 보자. 해당 로봇은 집을 돌아다니면서 카메라를 통해 얻은 이미지와 본인의 위치를 기억하여 집의 구조와 물건의 위치 등 새로운 환경에 대한 지식을 구성할 필요가 있다. 이를 위해, 긴 메모리를 갖는 것은 도움이 될 것이다. 단순히 직관에 의존해 하나의 관측치(이미지)로 부터 어떤 정보를 기억해야할지 고민해보자. 물건의 종류, 크기, 모양 그리고 위치 정도에 대한 정보를 기억해야할 것이다. 이처럼 하나의 관측치로 부터 기억해야 할 정보가 512차원 혹은 1024차원의 벡터에 담기지 않을 정도로 많지 않을거라 생각한다.

Context Length만 늘린다고 성능이 좋아지지 않는다

배경에서 설명했듯이 RetNet의 경우 "기억"의 역할을 하는 상태 행렬의 크기는 임베딩 차원의 제곱이다. 그렇기에 메모리의 크기를 키우기 위해 임베딩 차원을 늘려야하고 이는 RetNet의 추정 비용을 늘린다. 트랜스포머의 경우, self attention 헤드 혹은 레이어를 늘리지 않고 단순히 context length를 늘리는 것의 한계가 이론적으로 알려져있다.

결론

단순히 sequence model의 context length를 늘리는 것만으로 크고 긴 기억을 갖는 agent를 만들 수 있을지 모르겠다. 또한 좋은 표현을 학습할 수 있다면 짧은 context length로도 긴 기억을 가질 수 있다. 그렇기에 작업을 수행하는 데 필요한 내용을 간결하게 담는 표현을 학습할 방안에 대한 고민할 것이다.