We are excited to announce that the Pruna AI team will present 9 papers at NeurIPS 2024, showcasing our contributions across diverse fields of AI and machine learning. These papers explore cutting-edge developments in efficiency, uncertainty, robustness, privacy, and molecular generation, highlighting the breadth of expertise within our team. Here's an overview of some key papers, particularly those focused on improving efficiency in AI:

Shaving Weights with Occam's Razor: Bayesian Sparsification for Neural Networks using the Marginal Likelihood

Authors: Rayen Dhahri, Alexander Immer, Bertrand Charpentier, Stephan Günnemann, Vincent Fortuin

Neural network sparsification is a promising avenue to save computational time and memory costs, especially in an age where many successful AI models are becoming too large to naively deploy on consumer hardware. While much work has focused on different weight pruning criteria, the overall sparsifiability of the network, i.e., its capacity to be pruned without quality loss, has often been overlooked. We present a sparsification framework that highlights the effectiveness of using the Bayesian marginal likelihood in conjunction with sparsity-inducing priors for making neural networks more sparsifiable. Our approach implements an automatic Occam's razor that selects the most sparsifiable model that still explains the data well, both for structured and unstructured sparsification.

Efficient Adversarial Training in LLMs with Continuous Attacks (Spotlight)

Authors: Sophie Xhonneux, Alessandro Sordoni, Stephan Günnemann, Gauthier Gidel, Leo Schwinn

Large language models (LLMs) are vulnerable to adversarial attacks that can bypass their safety guardrails. While adversarial training has shown promise in improving robustness, its application to LLMs has been limited by high computational costs. We introduce C-AdvUL, a fast adversarial training algorithm operating in the continuous embedding space of the LLM, and C-AdvIPO, an adversarial variant of IPO for robust alignment. Through evaluation across multiple model families and scales, we demonstrate that these methods significantly enhance LLM robustness against discrete attacks while maintaining utility. Our results demonstrate that robustness to continuous perturbations can extrapolate to discrete threat models. Thereby, we present a path toward scalable adversarial training algorithms for robustly aligning LLMs.

Expected Probabilistic Hierarchies

Authors: Marcel Kollovieh, Bertrand Charpentier, Daniel Zügner, Stephan Günnemann

Traditional hierarchical clustering methods often rely on discrete optimization, constrained by large search spaces or continuous relaxations, which may not align with optimal discrete outcomes. In this work, we introduce Expected Probabilistic Hierarchies (EPH), a probabilistic model that optimizes expected scores of the Dasgupta cost and Tree-Sampling Divergence (TSD). We show that the optima of these scores align with their discrete counterparts, unlike those of previous works. EPH employs differentiable hierarchy sampling and unbiased subgraph sampling to efficiently scale to large datasets to find the hierarchies. Our approach achieves state-of-the-art performance across synthetic and real-world graph and vector datasets.

Unified Guidance for Geometry-Conditioned Molecular Generation

Authors: Sirine Ayadi, Leon Hetzel, Johanna Sommer, Fabian Theis, Stephan Günnemann

Designing molecular geometries is essential to advancing drug discovery, a field that has gained significant attention with the rise of generative models, particularly diffusion models. Despite their success, current molecular diffusion models are typically tailored for specific downstream tasks and lack adaptability to diverse applications. We introduce UniGuide, a novel framework built upon a condition map that enables controlled geometric guidance of unconditional diffusion models. UniGuide introduces flexible conditioning during inference, eliminating the need for additional training or auxiliary networks. We showcase how UniGuide flexibly integrates applications such as structure-based, fragment-based, and ligand-based drug design, achieving on-par or superior performance to task-specific approaches. By offering a more versatile approach, UniGuide has the potential to accelerate the development of molecular generative models, contributing towards novel pharmaceutical innovation.

Energy-based Epistemic Uncertainty for Graph Neural Networks (Spotlight)

Authors: Dominik Fuchsgruber, Tom Wollschläger, Stephan Günnemann

Quantifying the epistemic uncertainty of Graph Neural Networks (GNNs) is difficult because uncertainty sources can occur at different scales, e.g. node-level, cluster-level, etc. Existing approaches only distinguish uncertainty in the presence or absence of graph structure. We propose a Graph Energy-Based-Model (GEBM) that computes energy-based uncertainty at different structural scales and aggregates them into a single high-quality measure. It builds on previous logit-based energy measures which we regularize to induce an integrable probability density in the data space provably. This links GEBM to evidential methods and enables more robust predictions. Our uncertainty estimate can be applied to any pre-trained GNN ad-hoc and is highly effective in practice: It is the only estimator that simultaneously detects different families of distribution shifts applied to the data.

Neural Pfaffians: Solving Many Many-Electron Schrödinger Equations (Oral)

Authors: Nicholas Gao, Stephan Günnemann

Neural wave functions accomplished unprecedented accuracies in many-electron systems, though at a high computational cost. Recent works proposed amortizing the cost by learning generalized wave functions across different structures and compounds. Enforcing the permutation antisymmetry of electrons in such generalized neural wave functions remained challenging as existing methods require discrete orbital selection via non-learnable hand-crafted algorithms. This work tackles the problem by defining overparametrized, fully learnable neural wave functions suitable for generalization across molecules. We achieve this by relying on Pfaffians. The Pfaffian allows us to enforce the antisymmetry on arbitrary electronic systems. Our empirical evaluation finds that a single neural Pfaffian calculates the ground state and ionization energies with chemical accuracy across various systems.

Soft Prompt Threats: Attacking Safety Alignment and Unlearning in Open-Source LLMs through the Embedding Space

Authors: Leo Schwinn, David Dobre, Sophie Xhonneux, Gauthier Gidel, Stephan Günnemann

As open-source Large Language Models (LLMs) advance in capability, ensuring their safety becomes increasingly imperative. Yet, attacks tailored to open-source LLMs that exploit full model access remain largely unexplored. We address this research gap and propose the embedding space attack, which directly attacks the continuous embedding representation of input tokens. We find that embedding space attacks circumvent model alignments and trigger harmful behaviors more efficiently than discrete attacks or model fine-tuning. Additionally, we demonstrate that models compromised by embedding attacks can be used to create discrete jailbreaks in natural language. Lastly, we present a novel threat model in the context of unlearning and training data extraction and show that embedding space attacks can extract supposedly deleted information and training data from LLMs across multiple datasets and models.

Spatio-Spectral Graph Neural Networks

Authors: Simon Geisler, Arthur Kosmala, Daniel Herbst, Stephan Günnemann

Due to recursive aggregations over neighborhoods, Message-Passing Graph Neural Networks (MPGNNs) have a limited "receptive field", and information exchange between distant nodes is moreover limited by over-squashing. Our Spatio-Spectral Graph Neural Networks (S²GNNs), a new modeling paradigm for Graph Neural Networks, overcome these limitations by synergistically combining spatially and spectrally parametrized graph filters. We show that S²GNNs vanquish over-squashing and yield strictly tighter approximation-theoretic error bounds than MPGNNs. Furthermore, rethinking graph convolutions at a fundamental level unlocks new design spaces, such as free positional encodings that make S²GNNs strictly more expressive than the 1-Weisfeiler-Lehman (WL) test. Moreover, to obtain general-purpose S²GNNs, we propose spectrally parametrized filters for directed graphs. S²GNNs outperform spatial MPGNNs, graph transformers, and graph rewirings at a variety of benchmarks, are competitive with state-of-the-art sequence modeling, and scale to millions of nodes.

Unified Mechanism-Specific Amplification by Subsampling and Group Privacy Amplification

Authors: Jan Schuchardt, Mihail Stoian, Arthur Kosmala, Stephan Günnemann

Amplification by subsampling is one of the main primitives in machine learning with differential privacy (DP): If we train a model on random batches, then each training iteration has a certain chance of not accessing the private information of an individual in a dataset, thus limiting privacy leakage. But what if we want to protect the private information of multiple individuals, e.g., a family that contributed to a dataset? Can we formally derive group privacy amplification guarantees? We answer this question by developing a general framework for analyzing the privacy of subsampled mechanisms in a principled manner while simultaneously providing a unified view on over a decade of prior work on privacy amplification. Applied to group privacy, our tight analysis yields significantly better bounds on privacy leakage compared to classic group privacy results. In practice, this allows us to train group-private machine learning models for a much larger number of iterations at any given privacy budget.