O Pré-treinamento Supervisionado Pode Aprender Aprendizado por Reforço In-Contexto

Supervised Pretraining Can Learn In-Context Reinforcement Learning

June 26, 2023

Autores: Jonathan N. Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, Emma Brunskill

cs.AI

Resumo

Grandes modelos de transformadores treinados em conjuntos de dados diversos têm demonstrado uma habilidade notável de aprender em contexto, alcançando alto desempenho em poucos exemplos (few-shot) em tarefas para as quais não foram explicitamente treinados. Neste artigo, estudamos as capacidades de aprendizado em contexto de transformadores em problemas de tomada de decisão, ou seja, aprendizado por reforço (RL) para bandidos e processos de decisão de Markov. Para isso, introduzimos e estudamos o Transformador Pré-treinado para Decisão (Decision-Pretrained Transformer - DPT), um método de pré-treinamento supervisionado no qual o transformador prevê uma ação ótima dado um estado de consulta e um conjunto de dados de interações em contexto, em uma variedade de tarefas. Esse procedimento, embora simples, produz um modelo com várias capacidades surpreendentes. Descobrimos que o transformador pré-treinado pode ser usado para resolver uma gama de problemas de RL em contexto, exibindo tanto exploração online quanto conservadorismo offline, apesar de não ter sido explicitamente treinado para isso. O modelo também generaliza além da distribuição de pré-treinamento para novas tarefas e adapta automaticamente suas estratégias de tomada de decisão a estruturas desconhecidas. Teoricamente, mostramos que o DPT pode ser visto como uma implementação eficiente de amostragem posterior bayesiana, um algoritmo de RL comprovadamente eficiente em termos de amostras. Aproveitamos ainda essa conexão para fornecer garantias sobre o arrependimento (regret) do algoritmo em contexto gerado pelo DPT, e provamos que ele pode aprender mais rápido do que os algoritmos usados para gerar os dados de pré-treinamento. Esses resultados sugerem um caminho promissor e simples para instilar fortes habilidades de tomada de decisão em contexto em transformadores.

English

Large transformer models trained on diverse datasets have shown a remarkable ability to learn in-context, achieving high few-shot performance on tasks they were not explicitly trained to solve. In this paper, we study the in-context learning capabilities of transformers in decision-making problems, i.e., reinforcement learning (RL) for bandits and Markov decision processes. To do so, we introduce and study Decision-Pretrained Transformer (DPT), a supervised pretraining method where the transformer predicts an optimal action given a query state and an in-context dataset of interactions, across a diverse set of tasks. This procedure, while simple, produces a model with several surprising capabilities. We find that the pretrained transformer can be used to solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline, despite not being explicitly trained to do so. The model also generalizes beyond the pretraining distribution to new tasks and automatically adapts its decision-making strategies to unknown structure. Theoretically, we show DPT can be viewed as an efficient implementation of Bayesian posterior sampling, a provably sample-efficient RL algorithm. We further leverage this connection to provide guarantees on the regret of the in-context algorithm yielded by DPT, and prove that it can learn faster than algorithms used to generate the pretraining data. These results suggest a promising yet simple path towards instilling strong in-context decision-making abilities in transformers.