-
Benchmarking GPU and TPU Performance with Graph Neural Networks
Authors:
xiangyang Ju,
Yunsong Wang,
Daniel Murnane,
Nicholas Choma,
Steven Farrell,
Paolo Calafiura
Abstract:
Many artificial intelligence (AI) devices have been developed to accelerate the training and inference of neural networks models. The most common ones are the Graphics Processing Unit (GPU) and Tensor Processing Unit (TPU). They are highly optimized for dense data representations. However, sparse representations such as graphs are prevalent in many domains, including science. It is therefore impor…
▽ More
Many artificial intelligence (AI) devices have been developed to accelerate the training and inference of neural networks models. The most common ones are the Graphics Processing Unit (GPU) and Tensor Processing Unit (TPU). They are highly optimized for dense data representations. However, sparse representations such as graphs are prevalent in many domains, including science. It is therefore important to characterize the performance of available AI accelerators on sparse data. This work analyzes and compares the GPU and TPU performance training a Graph Neural Network (GNN) developed to solve a real-life pattern recognition problem. Characterizing the new class of models acting on sparse data may prove helpful in optimizing the design of deep learning libraries and future AI accelerators.
△ Less
Submitted 21 October, 2022;
originally announced October 2022.
-
Spatial Graph Attention and Curiosity-driven Policy for Antiviral Drug Discovery
Authors:
Yulun Wu,
Mikaela Cashman,
Nicholas Choma,
Érica T. Prates,
Verónica G. Melesse Vergara,
Manesh Shah,
Andrew Chen,
Austin Clyde,
Thomas S. Brettin,
Wibe A. de Jong,
Neeraj Kumar,
Martha S. Head,
Rick L. Stevens,
Peter Nugent,
Daniel A. Jacobson,
James B. Brown
Abstract:
We developed Distilled Graph Attention Policy Network (DGAPN), a reinforcement learning model to generate novel graph-structured chemical representations that optimize user-defined objectives by efficiently navigating a physically constrained domain. The framework is examined on the task of generating molecules that are designed to bind, noncovalently, to functional sites of SARS-CoV-2 proteins. W…
▽ More
We developed Distilled Graph Attention Policy Network (DGAPN), a reinforcement learning model to generate novel graph-structured chemical representations that optimize user-defined objectives by efficiently navigating a physically constrained domain. The framework is examined on the task of generating molecules that are designed to bind, noncovalently, to functional sites of SARS-CoV-2 proteins. We present a spatial Graph Attention (sGAT) mechanism that leverages self-attention over both node and edge attributes as well as encoding the spatial structure -- this capability is of considerable interest in synthetic biology and drug discovery. An attentional policy network is introduced to learn the decision rules for a dynamic, fragment-based chemical environment, and state-of-the-art policy gradient techniques are employed to train the network with stability. Exploration is driven by the stochasticity of the action space design and the innovation reward bonuses learned and proposed by random network distillation. In experiments, our framework achieved outstanding results compared to state-of-the-art algorithms, while reducing the complexity of paths to chemical synthesis.
△ Less
Submitted 11 May, 2022; v1 submitted 3 June, 2021;
originally announced June 2021.
-
Performance of a Geometric Deep Learning Pipeline for HL-LHC Particle Tracking
Authors:
Xiangyang Ju,
Daniel Murnane,
Paolo Calafiura,
Nicholas Choma,
Sean Conlon,
Steve Farrell,
Yaoyuan Xu,
Maria Spiropulu,
Jean-Roch Vlimant,
Adam Aurisano,
V Hewes,
Giuseppe Cerati,
Lindsey Gray,
Thomas Klijnsma,
Jim Kowalkowski,
Markus Atkinson,
Mark Neubauer,
Gage DeZoort,
Savannah Thais,
Aditi Chauhan,
Alex Schuy,
Shih-Chieh Hsu,
Alex Ballow,
and Alina Lazar
Abstract:
The Exa.TrkX project has applied geometric learning concepts such as metric learning and graph neural networks to HEP particle tracking. Exa.TrkX's tracking pipeline groups detector measurements to form track candidates and filters them. The pipeline, originally developed using the TrackML dataset (a simulation of an LHC-inspired tracking detector), has been demonstrated on other detectors, includ…
▽ More
The Exa.TrkX project has applied geometric learning concepts such as metric learning and graph neural networks to HEP particle tracking. Exa.TrkX's tracking pipeline groups detector measurements to form track candidates and filters them. The pipeline, originally developed using the TrackML dataset (a simulation of an LHC-inspired tracking detector), has been demonstrated on other detectors, including DUNE Liquid Argon TPC and CMS High-Granularity Calorimeter. This paper documents new developments needed to study the physics and computing performance of the Exa.TrkX pipeline on the full TrackML dataset, a first step towards validating the pipeline using ATLAS and CMS data. The pipeline achieves tracking efficiency and purity similar to production tracking algorithms. Crucially for future HEP applications, the pipeline benefits significantly from GPU acceleration, and its computational requirements scale close to linearly with the number of particles in the event.
△ Less
Submitted 21 September, 2021; v1 submitted 11 March, 2021;
originally announced March 2021.
-
Track Seeding and Labelling with Embedded-space Graph Neural Networks
Authors:
Nicholas Choma,
Daniel Murnane,
Xiangyang Ju,
Paolo Calafiura,
Sean Conlon,
Steven Farrell,
Prabhat,
Giuseppe Cerati,
Lindsey Gray,
Thomas Klijnsma,
Jim Kowalkowski,
Panagiotis Spentzouris,
Jean-Roch Vlimant,
Maria Spiropulu,
Adam Aurisano,
V Hewes,
Aristeidis Tsaris,
Kazuhiro Terao,
Tracy Usher
Abstract:
To address the unprecedented scale of HL-LHC data, the Exa.TrkX project is investigating a variety of machine learning approaches to particle track reconstruction. The most promising of these solutions, graph neural networks (GNN), process the event as a graph that connects track measurements (detector hits corresponding to nodes) with candidate line segments between the hits (corresponding to edg…
▽ More
To address the unprecedented scale of HL-LHC data, the Exa.TrkX project is investigating a variety of machine learning approaches to particle track reconstruction. The most promising of these solutions, graph neural networks (GNN), process the event as a graph that connects track measurements (detector hits corresponding to nodes) with candidate line segments between the hits (corresponding to edges). Detector information can be associated with nodes and edges, enabling a GNN to propagate the embedded parameters around the graph and predict node-, edge- and graph-level observables. Previously, message-passing GNNs have shown success in predicting doublet likelihood, and we here report updates on the state-of-the-art architectures for this task. In addition, the Exa.TrkX project has investigated innovations in both graph construction, and embedded representations, in an effort to achieve fully learned end-to-end track finding. Hence, we present a suite of extensions to the original model, with encouraging results for hitgraph classification. In addition, we explore increased performance by constructing graphs from learned representations which contain non-linear metric structure, allowing for efficient clustering and neighborhood queries of data points. We demonstrate how this framework fits in with both traditional clustering pipelines, and GNN approaches. The embedded graphs feed into high-accuracy doublet and triplet classifiers, or can be used as an end-to-end track classifier by clustering in an embedded space. A set of post-processing methods improve performance with knowledge of the detector physics. Finally, we present numerical results on the TrackML particle tracking challenge dataset, where our framework shows favorable results in both seeding and track finding.
△ Less
Submitted 30 June, 2020;
originally announced July 2020.
-
Graph Neural Networks for IceCube Signal Classification
Authors:
Nicholas Choma,
Federico Monti,
Lisa Gerhardt,
Tomasz Palczewski,
Zahra Ronaghi,
Prabhat,
Wahid Bhimji,
Michael M. Bronstein,
Spencer R. Klein,
Joan Bruna
Abstract:
Tasks involving the analysis of geometric (graph- and manifold-structured) data have recently gained prominence in the machine learning community, giving birth to a rapidly developing field of geometric deep learning. In this work, we leverage graph neural networks to improve signal detection in the IceCube neutrino observatory. The IceCube detector array is modeled as a graph, where vertices are…
▽ More
Tasks involving the analysis of geometric (graph- and manifold-structured) data have recently gained prominence in the machine learning community, giving birth to a rapidly developing field of geometric deep learning. In this work, we leverage graph neural networks to improve signal detection in the IceCube neutrino observatory. The IceCube detector array is modeled as a graph, where vertices are sensors and edges are a learned function of the sensors' spatial coordinates. As only a subset of IceCube's sensors is active during a given observation, we note the adaptive nature of our GNN, wherein computation is restricted to the input signal support. We demonstrate the effectiveness of our GNN architecture on a task classifying IceCube events, where it outperforms both a traditional physics-based method as well as classical 3D convolution neural networks.
△ Less
Submitted 17 September, 2018;
originally announced September 2018.
