Not All Latents Are Equal: Importance-Weighted Latent Graph Generation

Graph generation has shown tremendous effectiveness in modeling complex scientific tasks ranging from molecule generation to optimized circuit design. Most recent advancements have developed models that learn to denoise noisy graphs using either discrete diffusion or flow matching. However, discrete graph diffusion scales quadratically with the number of nodes, making it difficult to use such models on larger graphs. On the other hand, latent graph generation instead seeks to denoise noisy graph latents, which scale better to more nodes due to their lower dimension. Latent generation works by first learning a set of latents via an autoencoder trained on graph reconstruction. The learned latents are then denoised in the continuous space, where the final denoised latent is then mapped back to the discrete graph. However, in practice, latent graph generation models lag behind their discrete counterparts in performance. We study this problem, finding that it is due to the fact that the importance of specific latent units (i.e. dimensions) in a graph latent vary tremendously in their impact on graph reconstruction. That is, the vast majority of impact is contained in only a smaller number of latent units. However, existing work places equal emphasis on reconstructing the entire latent in the loss function, regardless of how much influence each unit has on reconstructing the graph. We propose to tackle this by simply weighting the latent units differently as per their importance to final graph reconstruction. The importance is calculated as the relative influence of a single latent unit on graph reconstruction, which is represented via the Jacobian of the decoder relative to the latent. Our method is model-agnostic, and we show that it can both boost performance and lead to faster training convergence when applied to two different autoencoders on various benchmark datasets.