Analyzing Neural Network Information Flow Using Differential Geometry

TL;DR

This work reimagines NN information flow through differential geometry by introducing neural curvature, derived from Ollivier-Ricci curvature, to rank neural connections by their contribution to classification accuracy. By constructing a neural graph and embedding sample-dependent neighbor distributions via activations, the method yields a continuous, data-aware edge-importance score that highlights bottleneck edges as carriers of essential information. Extensive pruning experiments on MNIST, CIFAR-10, and CIFAR-100 show that negative-curvature edges are crucial for maintaining performance, while many positive-curvature edges can be removed with minimal impact, outperforming several established pruning baselines in several settings. The approach is scalable with GPU-accelerated computation and offers a framework for robust NN analysis, model repair, and alignment, with future work targeting more efficient Wasserstein approximations.

Abstract

This paper provides a fresh view of the neural network (NN) data flow problem, i.e., identifying the NN connections that are most important for the performance of the full model, through the lens of graph theory. Understanding the NN data flow provides a tool for symbolic NN analysis, e.g.,~robustness analysis or model repair. Unlike the standard approach to NN data flow analysis, which is based on information theory, we employ the notion of graph curvature, specifically Ollivier-Ricci curvature (ORC). The ORC has been successfully used to identify important graph edges in various domains such as road traffic analysis, biological and social networks. In particular, edges with negative ORC are considered bottlenecks and as such are critical to the graph's overall connectivity, whereas positive-ORC edges are not essential. We use this intuition for the case of NNs as well: we 1)~construct a graph induced by the NN structure and introduce the notion of neural curvature (NC) based on the ORC; 2)~calculate curvatures based on activation patterns for a set of input examples; 3)~aim to demonstrate that NC can indeed be used to rank edges according to their importance for the overall NN functionality. We evaluate our method through pruning experiments and show that removing negative-ORC edges quickly degrades the overall NN performance, whereas positive-ORC edges have little impact. The proposed method is evaluated on a variety of models trained on three image datasets, namely MNIST, CIFAR-10 and CIFAR-100. The results indicate that our method can identify a larger number of unimportant edges as compared to state-of-the-art pruning methods.

Figures (9)

• Figure 1: ORC illustration on an undirected graph with unit edge weights. Red-colored edges (left) have a positive curvature; blue edge (bridge) has a negative curvature; yellow edges (right) have a zero curvature.
• Figure 2: Approach overview. Given a trained NN $f_\theta$ and a calibration set $\mathcal{D}$, we aim to identify the importance of NN connections, respectively. Specifically, we 1) construct a neural graph $G_{f_\theta}$ induced by the NN architecture (including an the input layer); 2) for each edge $(u,v) \in G_{f_\theta}$ and each example $x \in \mathcal{D}$, we calculate the neural curvature $\kappa_N(u,v,x)$, as defined in Definition \ref{['def:nrc']}; 3) finally we order the whole edge set by positive-first-curvature edges.
• Figure 3: Edge removal evaluation on MNIST, CIFAR-10, and CIFAR-100. Each subfigure shows a comparison of our neural curvature algorithm with magnitude, SNIP, and SynFlow pruning methods.
• Figure 4: Edge removal evaluation on MNIST and CIFAR-10 using our neural curvature algorithm. Each subfigure shows a comparison of the impact of removing edges in order of positive-curvature-first versus negative-curvature-first. Numbers on the curves indicate the minimum curvature per removed edge over the validation set.
• Figure 5: Ablation experiment on per-layer edge removal. Each subfigure shows the edge removal analysis for each layer based on neural curvature value (top 9 figures) and magnitude-based pruning (bottom 9 figures) for the VGG9-lite, CE ReLU configuration. The number above the figure is the total number of parameters in that layer.

Theorems & Definitions (16)

• Remark 1
• Definition 1: Wasserstein Distance (Optimal Transport Distance)
• Definition 2: ORC ollivier2009ricci
• Example 1
• Definition 3: $\alpha$-ORC lin2011ricci
• Definition 4: Ricci Curvature on Graphs lin2011ricci
• Remark 2
• Definition 5: Neural Graph tan24xiao24
• Remark 3
• Definition 6: Neural Neighbor Distribution