Publications

Discovering Object-Centric Generalized Value Functions From Pixels

Somjit Nath

Gopeshh Raaj Subbaraj

Deep Reinforcement Learning has shown significant progress in extracting useful representations from high-dimensional inputs albeit using ha… (see more)nd-crafted auxiliary tasks and pseudo rewards. Automatically learning such representations in an object-centric manner geared towards control and fast adaptation remains an open research problem. In this paper, we introduce a method that tries to discover meaningful features from objects, translating them to temporally coherent"question"functions and leveraging the subsequent learned general value functions for control. We compare our approach with state-of-the-art techniques alongside other ablations and show competitive performance in both stationary and non-stationary settings. Finally, we also investigate the discovered general value functions and through qualitative analysis show that the learned representations are not only interpretable but also, centered around objects that are invariant to changes across tasks facilitating fast adaptation.

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

Discrete Key-Value Bottleneck

Frederik Träuble

Anirudh Goyal

Nasim Rahaman

Michael Curtis Mozer

Kenji Kawaguchi

Yoshua Bengio

Bernhard Schölkopf

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

Equivariance with Learned Canonicalization Functions

Sékou-Oumar Kaba

Arnab Kumar Mondal

Yan Zhang

Yoshua Bengio

Siamak Ravanbakhsh

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

FAENet: Frame Averaging Equivariant GNN for Materials Modeling

Alexandre AGM Duval

Victor Schmidt

Alex Hernandez-Garcia

Santiago Miret

Fragkiskos D. Malliaros

Yoshua Bengio

David Rolnick

Applications of machine learning techniques for materials modeling typically involve functions known to be equivariant or invariant to speci… (see more)fic symmetries. While graph neural networks (GNNs) have proven successful in such tasks, they enforce symmetries via the model architecture, which often reduces their expressivity, scalability and comprehensibility. In this paper, we introduce (1) a flexible framework relying on stochastic frame-averaging (SFA) to make any model E(3)-equivariant or invariant through data transformations. (2) FAENet: a simple, fast and expressive GNN, optimized for SFA, that processes geometric information without any symmetrypreserving design constraints. We prove the validity of our method theoretically and empirically demonstrate its superior accuracy and computational scalability in materials modeling on the OC20 dataset (S2EF, IS2RE) as well as common molecular modeling tasks (QM9, QM7-X). A package implementation is available at https://faenet.readthedocs.io.

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

Flexible Phase Dynamics for Bio-Plausible Contrastive Learning

Ezekiel Williams

Colin Bredenberg

Guillaume Lajoie

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

FusionRetro: Molecule Representation Fusion via In-Context Learning for Retrosynthetic Planning

Songtao Liu

Zhengkai Tu

Minkai Xu

Zuobai Zhang

Lu Lin

Zhitao Ying

Rex Ying

Jian Tang

Peilin Zhao

Dinghao Wu

Retrosynthetic planning aims to devise a complete multi-step synthetic route from starting materials to a target molecule. Current strategie… (see more)s use a decoupled approach of single-step retrosynthesis models and search algorithms, taking only the product as the input to predict the reactants for each planning step and ignoring valuable context information along the synthetic route. In this work, we propose a novel framework that utilizes context information for improved retrosynthetic planning. We view synthetic routes as reaction graphs and propose to incorporate context through three principled steps: encode molecules into embeddings, aggregate information over routes, and readout to predict reactants. Our approach is the first attempt to utilize in-context learning for retrosynthesis prediction in retrosynthetic planning. The entire framework can be efficiently optimized in an end-to-end fashion and produce more practical and accurate predictions. Comprehensive experiments demonstrate that by fusing in the context information over routes, our model significantly improves the performance of retrosynthetic planning over baselines that are not context-aware, especially for long synthetic routes. Code is available at https://github.com/SongtaoLiu0823/FusionRetro.

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

proceedings.mlr.press

openreview.net

GFlowNet-EM for Learning Compositional Latent Variable Models

Edward J Hu

Nikolay Malkin

Moksh J. Jain

Katie E Everett

Alexandros Graikos

Yoshua Bengio

Latent variable models (LVMs) with discrete compositional latents are an important but challenging setting due to a combinatorially large nu… (see more)mber of possible configurations of the latents. A key tradeoff in modeling the posteriors over latents is between expressivity and tractable optimization. For algorithms based on expectation-maximization (EM), the E-step is often intractable without restrictive approximations to the posterior. We propose the use of GFlowNets, algorithms for sampling from an unnormalized density by learning a stochastic policy for sequential construction of samples, for this intractable E-step. By training GFlowNets to sample from the posterior over latents, we take advantage of their strengths as amortized variational inference algorithms for complex distributions over discrete structures. Our approach, GFlowNet-EM, enables the training of expressive LVMs with discrete compositional latents, as shown by experiments on non-context-free grammar induction and on images using discrete variational autoencoders (VAEs) without conditional independence enforced in the encoder.

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining

Shengchao Liu

weitao Du

Zhi-Ming Ma

Hongyu Guo

Jian Tang

Molecule pretraining has quickly become the go-to schema to boost the performance of AI-based drug discovery. Naturally, molecules can be re… (see more)presented as 2D topological graphs or 3D geometric point clouds. Although most existing pertaining methods focus on merely the single modality, recent research has shown that maximizing the mutual information (MI) between such two modalities enhances the molecule representation ability. Meanwhile, existing molecule multi-modal pretraining approaches approximate MI based on the representation space encoded from the topology and geometry, thus resulting in the loss of critical structural information of molecules. To address this issue, we propose MoleculeSDE. MoleculeSDE leverages group symmetric (e.g., SE(3)-equivariant and reflection-antisymmetric) stochastic differential equation models to generate the 3D geometries from 2D topologies, and vice versa, directly in the input space. It not only obtains tighter MI bound but also enables prosperous downstream tasks than the previous work. By comparing with 17 pretraining baselines, we empirically verify that MoleculeSDE can learn an expressive representation with state-of-the-art performance on 26 out of 32 downstream tasks.

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

Guided-topic modelling of single-cell transcriptomes enables sub-cell-type and disease-subtype deconvolution of bulk transcriptomes

Lakshmipuram Seshadri Swapna

Michael Huang

Yue Li

Cell-type composition is an important indicator of health. We present Guided Topic Model for deconvolution (GTM-decon) to automatically infe… (see more)r cell-type-specific gene topic distributions from single-cell RNA-seq data for deconvolving bulk transcriptomes. GTM-decon performs competitively on deconvolving simulated and real bulk data compared with the state-of-the-art methods. Moreover, as demonstrated in deconvolving disease transcriptomes, GTM-decon can infer multiple cell-type-specific gene topic distributions per cell type, which captures sub-cell-type variations. GTM-decon can also use phenotype labels from single-cell or bulk data as a guide to infer phenotype-specific gene distributions. In a nested-guided design, GTM-decon identified cell-type-specific differentially expressed genes from bulk breast cancer transcriptomes.

2023-07-03

bioRxiv (preprint)

doi.org

Hidden Symmetries of ReLU Networks

J. Grigsby

Elisenda Grigsby

Kathryn Lindsey

David Rolnick

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

High-Probability Bounds for Stochastic Optimization and Variational Inequalities: the Case of Unbounded Variance

Abdurakhmon Sadiev

Marina Danilova

Eduard Gorbunov

Samuel Horváth

Gauthier Gidel

Pavel Dvurechensky

Alexander Gasnikov

Peter Richtárik

During the recent years the interest of optimization and machine learning communities in high-probability convergence of stochastic optimiza… (see more)tion methods has been growing. One of the main reasons for this is that high-probability complexity bounds are more accurate and less studied than in-expectation ones. However, SOTA high-probability non-asymptotic convergence results are derived under strong assumptions such as boundedness of the gradient noise variance or of the objective’s gradient itself. In this paper, we propose several algorithms with high-probability convergence results under less restrictive assumptions. In particular, we derive new high-probability convergence results under the assumption that the gradient/operator noise has bounded central

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net

Maximal Initial Learning Rates in Deep ReLU Networks

Gaurav Iyer

Boris Hanin

David Rolnick

2023-07-03

Proceedings of the 40th International Conference on Machine Learning (published)

doi.org

openreview.net