FastGraph: GPU-beschleunigte Graphkonstruktion für die Teilchenrekonstruktion

Forscher von ETP, CMU und UZH haben kürzlich ein neues GPU-basiertes Algorithmenpaket namens FastGraph vorgestellt, das entwickelt wurde, um einen der zeitaufwändigsten Schritte in modernen graphbasierten Rekonstruktionsmethoden zu beschleunigen. Wenn Teilchen einen Detektor durchqueren, hinterlassen sie eine große Anzahl einzelner Signale ("Hits") . Um die Eigenschaften der Teilchen zu rekonstruieren, müssen diese Hits kombiniert und gemeinsam interpretiert werden. Eine leistungsstarke Methode hierfür besteht darin, die Hits miteinander Informationen austauschen zu lassen: Wenn wir benachbarte Hits miteinander verbinden, bilden sie einen Graphen, auf dem graphische neuronale Netze (GNNs) arbeiten können, um die zugrunde liegende Teilchenstruktur wiederherzustellen.

Die Herausforderung besteht darin, dass die Erstellung dieses Graphen alles andere als trivial ist. Eine häufig verwendete Strategie besteht darin, jeden Hit mit seinen k nächsten Nachbarn in einem gelernten Raum zu verbinden. Die Suche nach den nächsten Nachbarn wird jedoch sehr schnell sehr aufwendig: Bei einem naiven Ansatz muss jeder Hit mit jedem anderen verglichen werden, was mit der Quadratzahl der Hits skaliert. Und die Zahlen sind groß – je nach Detektorregion oder Simulationsaufbau kann ein einzelnes Ereignis leicht Hunderttausende von Hits enthalten. Die viele Wiederholung dieses Prozesses während des Trainings des neuronalen Netzwerks kann die gesamte Laufzeit dominieren und den gesamten verfügbaren Speicher belegen.

FastGraph behebt genau diesen Engpass. Es führt die Nachbarsuche direkt auf der GPU durch und ist auf die in GNNs verwendeten niedrigdimensionalen Räume zugeschnitten. Anstatt jeden Hit mit allen anderen zu vergleichen, unterteilt FastGraph den Raum in kleine Bereiche und beschränkt die Suche auf nahegelegene Regionen. Zusammen mit einer sorgfältig optimierten GPU-Implementierung ermöglicht dies dem Algorithmus, die exakten nächsten Nachbarn äußerst effizient zu berechnen, wobei die vollständige Kompatibilität mit modernen Frameworks für maschinelles Lernen gewährleistet ist.
In dem für die detektorbezogene Graphenrekonstruktion relevantesten Bereich erreicht FastGraph im Vergleich zu den modernsten Bibliotheken Geschwindigkeitssteigerungen zwischen 8 und 300, während der Speicheraufwand im Wesentlichen vernachlässigbar bleibt. Das bedeutet, dass der Engpass der Graphenkonstruktion – oft der langsamste Teil graphenbasierter Rekonstruktionspipelines – um Größenordnungen schneller gemacht werden kann.

Für die graphbasierte Teilchenrekonstruktionsforschung eröffnet dies mehrere neue Möglichkeiten: größere und detailliertere Graphen, schnellere Trainingszyklen und die Möglichkeit, komplexere Netzwerkarchitekturen ohne unerschwingliche Rechenkosten zu erforschen. FastGraph ist als Open Source unter der MIT-Lizenz verfügbar und kann direkt in bestehende PyTorch-basierte Workflows integriert werden.

Kontakt: Prof. Dr. Jan Kieseler