Politiques Transformer à diffusion efficaces avec mélange de dénoiseurs experts pour l'apprentissage multitâches

Les politiques de diffusion sont devenues largement utilisées en apprentissage par imitation, offrant plusieurs propriétés attractives, telles que la génération de comportements multimodaux et discontinus. À mesure que les modèles deviennent plus grands afin de capturer des capacités plus complexes, leurs exigences computationnelles augmentent, comme le montrent les lois de mise à l’échelle récentes. Par conséquent, la poursuite des architectures actuelles entraînera un goulot d’étranglement computationnel. Pour combler cet écart, nous proposons Mixture-of-Denoising Experts (MoDE), une nouvelle politique pour l’apprentissage par imitation. MoDE dépasse les politiques de diffusion basées sur les Transformers actuelles tout en permettant une mise à l’échelle efficace en termes de paramètres grâce à des experts épars et un routage conditionné au bruit, réduisant à la fois le nombre de paramètres actifs de 40 % et les coûts d’inférence de 90 % grâce au cache d’experts. Notre architecture combine cette mise à l’échelle efficace avec un mécanisme d’attention auto-conditionnée au bruit, permettant un débruitage plus efficace à travers différents niveaux de bruit. MoDE atteint des performances de pointe sur 134 tâches dans quatre benchmarks établis d’apprentissage par imitation (CALVIN et LIBERO). Notamment, en pré-entraînant MoDE sur des données robotiques diversifiées, nous obtenons un score de 4,01 sur CALVIN ABC et 0,95 sur LIBERO-90. Elle dépasse à la fois les politiques de diffusion basées sur les CNN et celles basées sur les Transformers d’une moyenne de 57 % sur quatre benchmarks, tout en utilisant 90 % moins de FLOPs et un nombre réduit de paramètres actifs par rapport aux architectures de diffusion Transformer par défaut. En outre, nous menons des ablations approfondies sur les composants de MoDE, offrant des perspectives pour concevoir des architectures Transformer efficaces et évolutives pour les politiques de diffusion. Le code et des démonstrations sont disponibles à l’adresse suivante : https://mbreuss.github.io/MoDE_Diffusion_Policy/.