Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning

要約

この論文では、累積的な後悔を最小限に抑える最適に近いアルゴリズムを学習することを目標とする、マルチタスク構造化バンディット問題を研究します。
タスクは共通の構造を共有しており、アルゴリズムはその共有構造を利用して、目に見えないが関連するテスト タスクに対する累積的な後悔を最小限に抑えます。
この共有構造を学習してテスト タスクに一般化するための意思決定アルゴリズムとしてトランスフォーマーを使用します。
DPT のような事前トレーニング済み意思決定変換器の以前の作業では、トレーニング中に最適なアクションにアクセスする必要がありますが、いくつかのシナリオでは困難な場合があります。
これらの研究とは異なり、私たちの学習アルゴリズムは、トレーニング中にタスクごとの最適なアクションの知識を必要としませんが、さまざまなトレーニング タスクから観察されたオフライン データのみを使用して、各アクションの報酬ベクトルを予測します。
最後に、推論時間中に、目に見えないテスト タスクのコンテキスト内でさまざまな探索戦略を採用した報酬予測を使用してアクションを選択します。
私たちのモデルは、いくつかの構造化バンディット問題 (線形、双線形、潜在、非線形) に関する一連の実験において、DPT やアルゴリズム蒸留などの他の SOTA 手法よりも優れています。
興味深いことに、私たちのアルゴリズムは、根底にある問題構造を知らなくても、多様なタスクにわたる共有構造を活用することで、コンテキスト内で最適に近いポリシーを学習できることがわかりました。
私たちは、新しいアクションで目に見えないタスクを活用しながらも、根底にある潜在構造を学習して最適に近いポリシーを導き出すことができることを示すことで、事前トレーニング済み意思決定トランスフォーマーの分野をさらに拡張します。
これをいくつかの実験で検証し、私たちが提案するソリューションが非常に一般的であり、テスト時に潜在的に新たなオンラインおよびオフライン戦略に幅広く応用できることを示します。
最後に、アルゴリズムのパフォーマンスを理論的に分析し、コンテキスト内のマルチタスク学習設定における一般化限界を取得します。

要約(オリジナル)

In this paper, we study multi-task structured bandit problem where the goal is to learn a near-optimal algorithm that minimizes cumulative regret. The tasks share a common structure and the algorithm exploits the shared structure to minimize the cumulative regret for an unseen but related test task. We use a transformer as a decision-making algorithm to learn this shared structure so as to generalize to the test task. The prior work of pretrained decision transformers like DPT requires access to the optimal action during training which may be hard in several scenarios. Diverging from these works, our learning algorithm does not need the knowledge of optimal action per task during training but predicts a reward vector for each of the actions using only the observed offline data from the diverse training tasks. Finally, during inference time, it selects action using the reward predictions employing various exploration strategies in-context for an unseen test task. Our model outperforms other SOTA methods like DPT, and Algorithmic Distillation over a series of experiments on several structured bandit problems (linear, bilinear, latent, non-linear). Interestingly, we show that our algorithm, without the knowledge of the underlying problem structure, can learn a near-optimal policy in-context by leveraging the shared structure across diverse tasks. We further extend the field of pre-trained decision transformers by showing that they can leverage unseen tasks with new actions and still learn the underlying latent structure to derive a near-optimal policy. We validate this over several experiments to show that our proposed solution is very general and has wide applications to potentially emergent online and offline strategies at test time. Finally, we theoretically analyze the performance of our algorithm and obtain generalization bounds in the in-context multi-task learning setting.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google