FlowPrune: Accelerating Attention Flow Calculation by Pruning Flow Network
Authors: Shuo Xu
Yu Chen
Shuxia Lin
Xin Geng
Xu Yang
Venue: NeurIPS 2025 poster ()
Keywords: Transformer
Attention Flow
Graph Pruning
Max-Flow Min-Cut
Interpretability
Primary area: deep_learning
TL;DR: We accelerates attention flow computation by pruning attention graphs, achieving high efficiency while preserving interpretability.

Abstract: The Transformer architecture serves as the foundation of modern AI systems, powering recent advances in Large Language Models (LLMs) and Large Multimodal Models (LMMs). Central to these models, attention mechanisms capture contextual dependencies via token interactions. Beyond inference, attention has been widely adopted for interpretability, offering insights into model behavior. Among interpretability techniques, attention flow --- which traces global information transfer across layers --- provides a more comprehensive perspective than single-layer attention maps. However, computing attention flow is computationally intensive due to the high complexity of max-flow algorithms. To address this challenge, we introduce FlowPrune, a novel framework that accelerates attention flow analysis by pruning the attention graph before applying max-flow computations. FlowPrune uses the Max-Flow Min-Cut Theorem and two structural properties of Transformer to identify and eliminate non-critical graph regions. It comprises two components: Edge Pruning, which removes insignificant attention edges, and Layer Compression, which discards layers with minimal contributions to the flow. We conduct extensive experiments on LLaMA and LLaVA to evaluate the robustness and effectiveness of FlowPrune. Our results show that FlowPrune achieves high agreement with the original attention flow in both absolute and relative error metrics, as well as in identifying influential input tokens. Finally, case studies in both NLP and vision domains demonstrate that FlowPrune produces consistent interpretability outcomes as the original Attention Flow, validating its practical utility. The code for this paper is publicly available.

FlowPrune: Accelerating Attention Flow Calculation by Pruning Flow Network
Shuo Xu Southeast University stmxsp01@gmall com
Yu Chen Southeast University yu_chen@seu. edu. cn
Shuxia Lin
Southeast University Thuxialin@senl edu. cn
Xin Geng
Xu Yang* Southeast University xuyang_palmm@aeu edu. en
Southeast University xgengEseu. edu. en
Abstract
The Transformer architecture serves as the foundation of modern Al systems, pow- ering recent advances in Large Language Models (LLMM) and Large Multimodal Models (LMMs). Central to these models, attention mechanisms capture contex- tual dependencies via token interactions. Beyond inference, attention has been widely adopted for interpretability. offering insights into model behavior. Among interpretability techniques, attention flow which traces global information trans- fer across layers provides more comprehensive perspective than single-layer attention maps. However, computing attention flow is computationally intensive due to the high complexity of max-flow algorithms. To address this challenge, we introduce FlowPrune, novel framework that accelerates attention flow analysis by pruning the attention graph before applying max-flow computations. FlowPrune uses the Max-Floo Min-Cut Theorem and two structural properties of Transformer identify and eliminate non-critical graph regions. It comprises two components: Edge Pruning, which removes insignificant attention edges, and Layer Compres- sion, which discards layers with minimal contributions to the flow. We conduct extensive experiments on LLaMA and LLaVA VA to evaluate the robustness and effec- tiveness of FlowPrune. Our results show that FlowPrune achieves high agreement with the original attention flow in both absolute and relative error metrics, as well in identifying influential input tokens. Finally, case studies in both NLP and vision domains demonstrate that FlowPrune produces consistent interpretability outcomes the original Attention Flow, validating its practical utility. The code for this paper is publicly available.
Introduction
The Transformer architecture 1 has become the foundational backbone of modern Al systems, driving major advances in Large Language Models (LLMs) 24 and Large Multimodal Models (LMMs) [5-7]] Although various architectural modifications have since been proposed g 9 attention remains the core computational primitive. enabling contextual information flow via token interactions. Because attention weights reflect how information is propagated across tokens, they are frequently used as proxies for interpretability. Despite ongoing debate about their faithfulness as explanations [10 I.. attention weights remain valuable tools for interpreting, diagnosing, and analyzing Transformer behavior across diverse applications [2 33
"Corresponding Author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Compared to single-layer attention maps that capture only local, layer-specific token interactions, attention flow offers global perspective by tracing how information propagates from input to output across all layers. This makes it more effective diagnostic tool for understanding model behavior. To quantify attention flow, [4 models the entire network as a Directed Acyelic Graph (DAG), where nodes represent token embeddings and edges denote attention weights. Within this framework, attention flow between input and output tokens is formulated as a maximum flow problem computed by the Highest-Level Preflow Push algorithm (15]. Although this approach captures cross-layer information that single-layer maps overlook, its computational complexity of O(L2 (3)) - for layers and N tokens- makes it impractical for modern large-scale models.
To reduce computational complexity, [14] proposes Attention Rollout, which approximates attention flow by recursively multiplying layerwwise attention matrices, reducing complexity to ((LN2) but introduces significant limitations. Specifically, it neglects global capacity constraints and flow conser- vation principles inherent to max-flow optimization, potentially distorting information propagation paths [4. As a result, risks producing misleading explanations that contradict ground-truth flows an issue critical to safety- applications like medical diagnosis or legal reasoning [16] 17
To address this, we propose FlowPrune, a novel method that accelerates attention flow computation without heavily distorting the analyses results. FlowPrune prunes the attention graph before applying max-flow algorithms where the pruning operation is designed based on the Max-Flow Min-Cur Theorem [18), stating that the maximum flow between a source and a sink node equals the total capacity of the minimum cut separating them. This implies that graph regions not contributing to the minimum cut can be safely pruned without affecting the final flow. We leverage two key properties of Transformer-base based networks to identify such regions: the attention weights are calculated using the softmax function, and the attention graph is strictly layered.
FlowPrume consists of two components: Edge Pruning and Layer Compression, both designed to reduce the attention graph while preserving its critical flow structure. First, due to the softmax operation, many attention weights are close to zero [19], contributing little information propagation. These weights can be safely removed without significantly affecting the overall max-flow value. Second, because the attention graph strictly layered. we localize minimum cut edges to specific layers, allowing for substantial graph reduction. We use a random sampling heuristic to identify layers with fewer minimum cut edges, which are then removed, and new edges are created between adjacent layers to maintain connectivity. These two procedures significantly compress the attention graph, enabling more efficient max- -low computation.
We validate FlowPrune on two sets of experiments, First, we evaluate the fidelity and efficiency of our graph compression by applying to two widely used models: the unimodal Llama and the multimodal LLaVA 200 21]. We compare the max-flow computation on the original and compressed graphs in terms of runtime and absolute/relative error. Since attention flow often used to assess which input tokens contribute more to model predictions (14). we also measure the ranking discrepancy of input tokens by comparing their relative max- -fowwllow orderings. Second, we conduct case studies on real-world scenarios where attention flow has been applied for model analysis. Specifically, we test the practicality of FlowPrune in two real-world scenarios: analyzing Paraphrasing Verification [22 and constructing ViT Heatmaps, where results prove that FlowPrune is more similar with the original Attention Flow than Attention Rollout, validating its practicality. All our code is publicly available, and you can find it at
2 Related Work
Attention Flow and Interpretability. The interpretability of attention mechanisms has attracted significant interest, particularly in NLP and vision models. While some studies suggest that attention weights offer insight into model reasoning
23 24 others argue that they may not faithfully Despite this debate, attention remains a widely used tool for attributing model predictions to input elements. For instance, attention- based visualizations have been used to identify important areas in image classification 22 27), diagnose linguistic errors
12
reflect decision- -aaking processes 100 11 25
translation 28 29,, and trace causal relationships in sentence classification 10 30). Beyond direct visualization, recent work has formalized attention flow as global explanation mechanism [14], tracing information propagation across layers via graph-based analysis. However, such methods are computationally expensive and may become infeasible for large models. Several approximations,

---

including Attention Rollout [14, have been proposed to reduce complexity, though they often sacrifice theoretical properties such as flow conservation. Meanwhile, the observed sparsity of attention matrices [9] motivates efficient pruning-based approaches that preserve interpretive fidelity while improving scalability.
Maximum Flow Algorithms. The maximum flow problem is a classical topic in graph theory, which seeks to determine the maximum amount of flow that can be sent from a source sink under capacity constraints. Common algorithms include Edmonds-Karg [3.. which uses breadth-fins search find augmenting paths, Dinic's algorithm [32. which improves throughput via level graphs, and Bussh-Relabel method [33], which maintain a preflow and iteratively update local excess. Among these, the Highest-Leve Preflow-Puss (HLPP) algorithm [[5] has demonstrated strong practical efficiency, especially on layered or DAG- graphs, making it suitable for attention flow computation [14]
Building on this foundation, our method, FlowPrune, accelerates attention flow computation by applying HLPP to structurally simplified attention graph. We leverage the Max-Flow Min-Cut Theorem [8] to guide pruning decisions, safely removing low-capacity edges and compressible layers that do not contribute to the minimum cut. This design reduces computational complexity while preserving critical attribution paths, enabling scalable analysis of large Transformer models.
3 FlowPrune
To address the scalability limitations of attention flow analysis in large-scale Transformer networks, we propose an efficient approximation method named FlowPrune. In this section, we first review the standard Attention Flow formulation and its computational bottlenecks. We then introduce two key approximation strategies, namely Edge Pruning and Layer Compression, which jointly reduce time complexity while preserving essential attention flow structures to calculate maximum flow.
3.1 Calculating Maximum Flow
To calculate the maximum flow that quantifies how attention propagates across layers in a L-layer Transformer, 14] models the attention mechanism as a layered directed graph G = (V, E). Specif- ically, V (Vi..... V. where each V4 represents the set of nodes corresponding to token rep- resentations at the i-th Transformer layer. The edge set is defined as = (E EL-1), }. where & denotes directed edges capturing attention from tokens in layer i to tokens in layer + with normalized attention weights used as edge capacities.
Given this graph G constructed from an I--layer Transformer, the classical Highest-Level Preflow Push (HLPP) algorithm can be applied to compute the maximum flow, with a time complexity of O(V2VE) Since the number of nodes and edges scale as V LN and x N respectively
where N is the number of tokens per layer and L is the number of layers - the total computational complexity becomes ((L2.5 N*). This high cost makes the method impractical for large-scale Transformer models.
3.2 Pruning the Attention Graph based on Max-Flow Min-Cut Theorem
Our objective is to accelerate the computation of maximum flow in large-scale transformer archi- tectures without largely disturbing the analyses results. To achieve this, we brought ideas from the classic Max-Flow Min-Cut Theorem [18] theory, which states that: In directed graph with non-negative edge capacities, the maximum amount of flow that can be sent from a source node to a sink node equals the total capacity of the minimum cut that separates the source from the sink. This theory provides a useful intuition: if certain regions of the graph do not have the minimum cut, they can be safely removed without significantly changing the overall flow, e.g., by pruning low-capacity edges or compressing intermediate layers without minimum cut.
To apply the above-mentiones insight to prune the attention graph, we also need to exploit two key structural properties of Transformer. First, the computation of attention weights involves softmax operation and second, Traasformer exhibits strictly layered structure. We next show how to integrate these two properties with Max-Flow Min-Cut Theorem to design Edge Pruning and Layer Compression strategies for reducing the graph scale.
3

---

Expedtt inveshol Capecty a Tiveshel
Mani
New Eage
(a) Edge Pruning.
(b) Find Potential Min-cut Edge.
(c) Layer Compression.
Figure 1: Pipeline of FlowPrune. (a) The edges with low flow capacity are removed. (b) Three paths are sampled (blue, purple, and brown ones) to find the layers which has less mini-cut edge, where the third layer does contain min-cut edges. (c) The third layer is removed and the connectivity between the second and fourth layers are re-built by multiplying two binary adjacency matrices.
Edge Pruning. In Transformers, the softmax operation used in attention computation often produces many near-zero attention weights, which contribute minimally to the propagation of information across layers. These negligible weights correspond to edges in the attention graph that carry insignifi- cant flow and can therefore be safely pruned to reduce computational overhead. To eliminate such uninformative edges, as shown in Figure Ta we set pruning threshold and remove all edges whose capacities fall below this value, effectively treating them as having capacity.
Layer Compression. As previously diseussed, the attention graph of a Transformer is strictly layered structure, ensuring all flow paths follow layer-wise transitions. This property offers two key advantages for compression. First, allows to localize minimum cut edges to specific layers, allowing to treat each layer as a fundamental unit for compression. By removing layers with less critical contributions to the overall attention flow, we can achieve substantial graph reduction while not heavily damaging the maximum flow approximation. Second, the layered structure permits efficient removal of entire intermediate layers and reconstruction of cross-layer connections via matrix multiplication. Next, we show how to identify compressible layers and establish new cross-layer connections.
Motivated by the Max-Flov Min-Cut Theorem, we seek compress layers containing edges that minimally contribute to the network's 's minimum cut, as these are less critical for information flow. Since exact minimum cut computation computationally expensive - equivalent to solving the full max-flow problem [4] we introduce an efficient sampling-bass based heuristic, which is shown Figure Tb] After constructing the attention graph, we randomly sample source-to-sink paths in this graph and identify the lowest-capacity edge in each path potential min-cut edge. Then we compress layers containing relatively few such edges. Following Transformer architecture principles [3), we preserve the first and last layers maintain essential signal initiation and termination.
It should be noted that identifying the lowest-capacity edge along sampled path as a potential min-cut edge is a heuristic rather than guaranteed criterion. The intuition is as follows: in the case where source-to-sink path does not overlap with others, the smallest-capacity edge this path must belong to the minimum cut by definition 18]. However, more realistic settings where multiple paths often share edges, this assumption becomes less precise. Nevertheless, attention graphs derived from Transformers are densely connected in a strictly layered structure, where each token in layer typically attends to all tokens in layer + 1. This cear-uniform connectivity implies that sampled paths are equally likely to overlap on any edge. As a result, the edge with the lowest capacity on random path remains statistically more likely lie on the global min-cut, justifying its use as proxy in our sampling- based approximation.
After identifying layers for compression, we reconnect the remaining layers to maintain attention flow connectivity. When compressing layer i, we add direct edges from nodes layer layer i 2 whenever path exists through layer i + 1. In practice, this reconnection can be efficiently implemented by multiplying binary adjacency matrices representing layer-wise connectivity, as illustrated in Figure Ic The new edge capacities set to 1, which preserves the maximum flow

---

since: (1) all original attention weights are normalized to [0, 11, and (2) flow network theory guarantees that increasing edge capacities beyond actual flow values maintains the maximum flow 34.
Combining edge pruning and layer compression significantly reduces the attention graph size. We then compute the maximum flow using Highest-Level Preflow Push algorithm on this compressed graph, achieving substantial complexity reduction while maintaining flow accuracy.
4 Experiments
We assess the effectiveness of FlowPrune based computational efficiency, fidelity the original attention flow, and practical utility in real-world tasks, We evaluate FlowPrune in two cases: (I) generation cases on LLMs and LMMs, and (2) application cases on language and vision tasks. All our experiments were conducted on a single server equipped with 2 AMD EPYC 7453 28-Core Processors and 4 NVIDIA RTX A6000 GPUs.
4.1 Fundamental Evaluation
4.1.1 Experimental Settings and Evaluation Metries
We evaluate LJama 20 and Qwen 355, representing language-only transformers, on the other hand LLaVA (1 and VL 36) representing vision-language transformers, respectively. Each model has 32 layers. We experiment with seven compression configurations: the full model (32 layers), an extreme compression setting (3 layers), and five intermediate retain rates of 80%, 75%. 50%, 30%, and 25%, which correspond to 26, 24, 16, 10, and 8 layers, respectively. The edge pruning threshold set to 1 10-t throughout all experiments.
The datasets used by these models to generate attention maps are GSM8K 37 and OKVQA [8], respectively. For each model, provided 1,000 inputs and randomly selected 100 token pairs from each generated attention map (including those compressed using FlowPrune) to calculate the corresponding results.
We propose four fundamental metrics to assess the efficiency and approximation quality. (1) Com- putational Cost (Time). This metric measures the wall-clock time for computing attention flow using FlowPrune, compared to the original implementation, serving as an indicator of computa- tional efficiency. Lower values indicate faster inference.
(2) Approximation Error. We compute both the absolute and relative errors between the approximated and original attention flow matrices, averaged across all examples. Smaller values reflect higher accuracy in the approximation. (3) Top-K Token Retention. For each input, we extract the top-K highest- -foow tokens from both the original and approximated attention flows, and compute their Jaccard overlap [39]. Higher values indicate better retention of key tokens. Specifically, IOU |AnB! where A/ respectively repre- sents the top-K K token pairs of the attention flows calculated by original/compressed graphs, where K 1, 5, 10, 15, 20, 25, 50 in the experiments. (4) First Rank Divergence. We identify the first position in the sorted top- K token rankings where the approximation diverges from the original output, normalizing this position by sequence length. Higher values reflect better consistency with the original ranking.
AUU
4.1.2 Results Analysis
Computational Cost (Time). Figure presents the speed-up ratios of Llama and LLaVV under varying retain rates. We use 1,000 samples from the GSMSK dataset for Llama and 1,000 samples from OKVQA for LLaVA. Each sample corresponds to unique attention graph, and for each graph, attention flows are computed between 1,000 token pairs. The total computation time is used to derive the speed-up ratios.
As previously discussed, for a Transformer with L layers and N tokens per layer, the computational complexity of attention flow computation is ( With Layer Compression, retaining only L/R layers reduces the complexity to O(I2 N3 / /22-5, achieving theoretical speed-up of R2.5
The empirical results in Figures 2 and generally align with this analysis. For instance, with 50% retain rate, Llama's computation time is halved-down from 0.55 seconds per case. When compressed to just three layers, FlowPrune requires only 5.31% of the original attention flow time. On LLaVA.
5

---

FlowPrune achieves even greater gains, reducing computation time to approximately 35% of the original a 50% retain rate.
Interestingly, the two models exhibit different speed-up patterns. LL's performance closely matches the the- oretical expectation, whereas Llama demonstrates a more linear reduction average computation time. This dis- crepancy may be attributed to differences in the attention weight distributions across the two models. Llama may re- tain more edges during Layer Compression, which causes its time cost to decrease more slowly than the theoretical estimate.
0.8
0.6
0.4
Lama
Owen
OwerrV
0.6
0.8
1.0
Retain Rate
While compressing the attention graph introduces some overhead, this step is performed only once per graph. All subsequent token-pair computations benefit from the re- duced graph size. Figure 3 shows the average per-pair computation time- including compression overhead as the number of token pairs per graph increases. Due to longer input sequences (approximately 500 tokens), LLaVA has a higher average time cost than Llama (typ-
Figure Speed-up Ratio. The horizon- tal axis shows the retain rate (fraction of layers retained), and the vertical axis shows the average time cost normalized by the original Attention Flow runtime, measured over 100 token pairs.
ically around 100 tokens). Nonetheless, the compression overhead is amortized the number of token pairs increases. For a 50% retain rate, the overhead becomes negligible after computing more than cases for Llama and cases for LLaVA.
wore Time
2.0
1.5.
1.0-
15 layes
11.2% 38 layan
52.0% 11 layes
6hh 19yer
120
50
Oriinnll
56.0%
10.24
25,06, 1yyes
10.2% layers
2.5
2.0
1.5
3.5
3.0
50 ph. 14 layara
61 2%, tayers
byeers
layara
layer
100
90
40
geant Everear
30.2%, 26 heyar
75.2% 24 bayer
$1.24, 26 loyas
$3,24, layars
9.6%
Oriinal Avarage Time Cont
0.
0.5
0.0
20
0.0
10
15
20
2
10 1.0
Number Cases
Number Caose
Number of Cases
Number of Cases
(a) LLaMA
(b) LLaVA
(c) Qwen
(d) QwenVL
Figure 3: Comparison of Average Overhead. The horizontal axis indicates the number of token pairs computed by per attention map. The vertical axis shows the average time cost per token pair, including the overheads. Each curve corresponds to specific compression setting. The black dashed line denotes the average computational time of the original (uncompressed) Attention Flow.
Approximation Error. Figures 4 provide detailed analyses of how absolute and relative errors vary with the true attention flow values. Since relative eгror fluctuates with both retain rates and original attention flow values, we present results for all retain rates in Figures 4b and 4cl and separately report results for retain rates of 30% and above in Figures 4c and 4f
We observe that both absolute and relative errors decrease the original attention flow approaches This is due to two properties: (1) Since the sum of all edge capacities from a source node in an attention graph is 1, the maximum flow cannot exceed 1. Edge Pruning only reduces edges originating from a node, so this property remains valid for the compressed graph. (2) In Layer Compression, newly added edge capacities set to 1, which can be regarded positive infinity in the attention graph, while connectivity between remaining nodes stays unchanged. This ensures that the maximum flow from the compressed graph increases, not decreases.
Both properties constrain the errors when the original attention flow is close to 1, leading to reduced approximation errors. This property is generally favorable, as we are primarily concerned with the top source source-destinnge destinatiin pairs with the highest attention flow values. These pairs indicate the highest attention weights between token pairs, which the focus of conventional attention visualization. The small error in this range, as observed with FlowPrune, enhances confidence in its results.
Top-K K Token Retention. Many studies use attention flow methods focus on the top-K token pairs with the highest attention flow. To address this, we use the Top- Token Retention metrie to calculate the Jaccard overlap (IOU) of these token pairs. Figure 5 shows the results. In Figure 5 the horizontal

---

0.30
R0-25.0%, syens
M*d11.% ayer
6-
RR-01 26 lyyens
AR-75 byers
0.6
B5-01 25
RS- 24
0.08
0%, ayers
5
RR-50 byers
0.5
RM-30.3 15
AR-11 ayers
RR-31 28, 10 layer
Rbw31.2% 25.2% 12 uyen
0.66
09,25.0%, 8 ltyers
RM-25 0le layos
Aivd 4%, 3 syors
AR-9 44% 3 lbyer
0.3
0.44
0.2
0.02
0.50
6.0
0.
0.0 Original Attn Flow Value
0.2
0.
1.0
Original Attn Flow Value
Original Attn
(a) Llama Absolute Error
(b) Llama Relative Error
(c) Relative Error (RR 2 30%))
175
0.330
byers
layera
MR-11 layera
MR-15 layera
0.35
Rt.51.2% 12.22..22 RS: 75.0', 24 layan RS: 50.2's 12layan Hiw 21.25, 10 loyes
0.125
Ri-32 layer
Mi-00 layes
AB-31 loyers
MRw3l layers
0.100
28-22 layarss
0.25
0.75
4%, layera
6-
byers
0.20
oose
0.19
0.25
2-
0.05
0.000
5.0
0% Original Attn Flow Value
0.4
0.0
Original Attn Flow Value
0.2
0.4
2-6
0.0
1.0
Original Attn Flow
(d) LLaVA Absolute Error
(e) LLaVA Relative Error
(f) Relative Error (RR > 30%))
0.330
0.755
nR-50 0s 16 leyers
R5 0% 24 loyers
RR--1 2% 25 layera
17.5
keos
RRw 2%, 26 ayers
Rl.55.0% 15 kyas B8-33.2% ayars
75,05,000 24 ayes
RB-31 2%, 18 leyors
RBI-21 0% laaeers
12.5
nn-25 0%. leyes
0.075
AR-9 4%, byyers
10.0
AMw2.4%, lopers
5.0
0.25
0000
0.0
Original
0.2
0.0
Original Attn Flow Value
0.2
2.4
0.6
0.3
1.0
0.2 Original
1.0
(g) Qwen Absolute Error
(h) Qwen Relative Error
(i) Relative Error (RR 2 30%))
030
RM-75
BR
5R-81.2% 26 layers
88-15 28, byyer
Rh-dl: 28, 26 ayer
#5- pe, 24 aver
AA-50.0%
8Ru50 28, lyyers
Rk-50 ayers
0.00
RR 28, loyers
Rh-31 10 apens
0.66
ayers
Baaeers
layers
0.3-
0.44
4-
0.2
0.3
0.0
0.2
0.4
0.6
0.8
10
02
1.0
Original Attn Flow Value
Original Attn Flow Value
Original Value
(i) QwenVL Absolute Error
(k) QwenVL Relative Error
(I) Relative Error (RR 2 30%))
Figure 4: Approximation Error vs. True Attention Flow Value. Each subplot shows the approxi- mation error (either absolute or relative) between FlowPrune and the original Attention Flow. The horizontal axis denotes the original attention flow values, and the vertical axis shows the average error magnitude for all token pairs falling within each value. Subplots (a)-(c) show the absolute and relative errors for Llama, and subplots (f)-(f, (g)-())(i)- show the same for LLaVA, Qwen and Qwen VL. RR represents retain rates.

---

1.00
1.00
1,000
1.00
0.99
0.98
0.96
14.09
0.98
0.97
005
1athh
0.930
161
1823
0.96
0.00
0.96
1922
0.950
1833
0.90
0p20
02
0.6
0.6
0.5
1.9
0.5
10
3.6
1.5
Retain
Retain
(a) Llama
(b) LLaVA
(c) Qwen
(d) QwenVL
Figure 5: Top-K K Token Retention. We evaluate the overlap between the top-K tokens selected by FlowPrune and those selected by the original Attention Flow method for different values of K.
Retention is measured using the Jaccard index (IoU), defined as loU better preservation of salient attention targets.
[AnB]
Higher values indicate
AUB
axis represents the retain rate, defined as the ratio of layers the compressed attention map to those in the original map. A lower retain rate corresponds fewer remaining layers. The vertical axis shows Jaccard overlap values. The leftmost point represents the case where the retain rate is minimal, corresponding compressing the attention map three layers. Even in this case, the Top-K Token Retention remains around 0.90, demonstrating that our approximation algorithm maintains high fidelity for the top token pairs. This result aligns with findings from the Approximation Error experiment, which examined eгror variation with original attention flow values.
First Rank Divergence. Figure 6 shows the results of this metric for two models. The curves depict the average First Rank Divergence (FRD) across all attention map samples under varying retain rates. When the retain rate restricted 50%, the average FRD of both models exceeds 0.4. Notably, even at the highest compression level (retaining only 3 layers), the average FRD of Llama is close to 0.2, while the average FRD of LLaVA exceeds 0.3. This result fully sufficient for practical applications with lower precision requirements, While the results at a retain rate 50% are adequate for tasks with higher precision demands.
1.0
First Rank Divergence
0.8
0.6
Liama
ULVA
OwenVL
0.2
0.4
0.6
0.8
1.0
Retain Rate
Figure 6: First Rank Divergence. This metric records the first position in the sorted top- K token list where the approx- imation diverges from the original out- put, normalized by the sequence length.
4.2 Application Scenarios
We use two application scenarios in both NLP and Vi- sion, which are Paraphrasing Verification and Vision Transformer (VTT) Heatmaps, show the practicality
of FlowPrune. Paraphrase Verification (22 involves determining whether two sentences express the same similar meaning. Attention Flow offers a visual analytics framework for examining self-attention pattems, illustrating how fine-tuning sharpens attention toward task-relevant details. In ViT, attention heatmaps depict the spatial distribution of attention weights, highlighting the model's focus regions during inference.
4.2.1 Paraphrasing Verification
In this task, we need to trace the attention flow from the classification token in the final layer backward through the network, across layers and attention heads. By comparing attention patterns between pre-trained and fine-tuned models, it reveals how fine-tuning sharpens attention on task-relevant distinctions, such as phrases that differentiate sentence meaning. This visual comparison offers valuable insights into the model's reasoning and decision- decision-making process. Here, we replace the original Attention Flow algorithm with FlowPrune and compare its results against those of Attention Flow and Attention Rollout, demonstrating the practical efficiency and effectiveness of FlowPrune.
In this case, we implement the experiments on MRPC (Microsoft Research Paraphrase Corpus) 40 dataset with a 12-layer BERT model, which asks the model to determine whether two sentences have the same meaning. We compare the results of FlowPrune and Attention Rollout with the original

---

Table I: Max-n Overlap for Paraphrasing Verification. FP represenns Flowprune and AR represents
Attention (Mean + SE)
Max-n
FP Pretrained
AR Pretrained
FP Fine-tuned
AR Fine-tuned
98.50 + 0.86%
63.40 + 2.58%
95.00 + 1.50%
84.50 + 2.32%
n = 5
97.93 + 1.19%
84.71 + 1.83%
92.50 = 2.26%
76.62 = 1.93%
n 10
97.69 1.33%
66.38 + 1.58%
91.45 2.58%
70.85 = 1.56%
770
25
(a) Flow Prune (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
10
47
(d) FlowPrunc (Fine-tuned)
(c) Attention Flow (Fine-tuned) (f) Attention Rollout (Fine-tuned)
Figure 7: Paraphrasing Verification Example. Visualization of attention flow results across three methods FFlows-nee original Attention Flow, and Attention Rollout-on sentence-pair classification task. The two input sentences are: (I) "The pound also made progress against the dollar, reached fresh three-year highs $1.6789." and (2) "The British pound flexed its muscle against the dollar, last up 1 percent at $1.6672."
Attention Flow in both pre-trained and fine-tuned models. In this task, we primarily focus on the attention metrics from tokens in the initial layer pointing to the /cls] token in the final layer. The tokens with the highest computed attention metrics are considered the most important.
We calculate the Max-n Overlap between the most important tokens obtained using Flow- Prune/Attention Rollout and the original Attention Flow. For FlowPrune, we use a retain rate of 67%, reducing the original 12-layer attention map to 8 layers. Similar to the Top-K Token Re- tention, we also use the Jaccard overlap (IOU) as the corresponding metric. We take n = 3, 5, 10 and calculate the corresponding results.
the results, where we can see that FlowPrune has higher Max-n Overlap with original Attention Flow than Attention Rollout, proving FlowPrune is a better approximation than Attention Rollout in analyzing the model behaviour. We also show one example from MPRC in Figure 7 We find that the analysis results from FlowPrune are almost identical to those from the original Attention Flow, and both are capable of identifying the focal points of attention, effectively distinguishing between pretrained and fine-tuned models. In contrast, the Attention Rollout method is limited by its inherent flaws, resulting in inferior outcomes compared to the other two analysis methods.
4.2.2 Vision Transformer (ViT) Heatmaps
When conducting an attention interpretability analysis for the VIT model, we may be interested in understanding the attention relevance of specific token to all other tokens to construct an attention heatmap. To test whether FlowPrune is useful here, we the Deit-Small (1) model and implement the experiments on ILSVRC2012 [42), an image classification dataset. Similarly, we primarily focus on the attention metrics from tokens the initial layer pointing to the /cls] token in the final layer, and compare the relevant analysis results of Attention Flow, FlowPrune and Attention Rollout. For Phone, we use retain rate of 25%, reducing the 12-layer attention map to layers.

---

(a) Original Image
(h) Attention Flow
(c) FlowPrune
(d) Attention Rollout
Figure 8: VIT Heapmap Example. We visualize the attention attribution maps produced by three methods - FlowPrune, the original Attention Flow and Attention Rollout heatmaps correspond to the [CLS] token and show which regions of the image contribute most to the model's 's final prediction.
ViT model. All
Table 2: Max-n Overlap for VIT Heatmap. We
Like Paraphrasing Verification, we calculate the Max-n Overlap between Original Attention Rollout. Con-
report the top-n region overlap between FlowPrune (FP), Attention Rollout (AR), and the reference attention map computed via full Attention Flow. (Mean SE)
Flow and sidering that VIT has more input tokens, We take n 10, 20, 30 and get results in Table We can see Flowerune has much higher Max- n Overlap with the original Attention Flow than Attention Rollout, indicating FlowPrune is bet- ter approximation than Attention Rollout in con- structing the heatmaps for ViT. Figure 8 presents an example of a ViT heatmap, which was gener-
Max-n
FP
AR
10
100.0
0.00%
28.18 + 3.06%
20
99.53 + 0.12%
50.84 + 3.74%
=
99.59 0.11%
78.52 + 3.43%
ated using three different methods. Deit-Small model utilizes the [CLS] token for image classification tasks. Therefore, we focus on the attention metrics from all tokens to the [CLS] token to construct heatmaps. The analysis results of FlowPrune once again align with the original Attention Flow, while the Attention Rollout method, due to its inherent flaws, fails to effectively analyze the propagation of attention.
5 Conclusion and Limitation
In this work, we propose FlowPrune, novel framework to accelerate attention flow computation, aiming to facilitate the analysis of Transformer- based networks. Motivated by the Max-Flow -low Min-Cut Theorem, FlowPrune leverages two structural properties of Transformers: the use of the softmax operation in attention computation and the strictly layered architecture of attention graphs. FlowPrune introduces two key components: Edge Pruning, which removes insignificant attention edges, and Layer Compression, which discards layers with minimal contributions to the flow. Together, these techniques effectively reduce the scale of the attention graph while preserving the fidelity of attention flow analysis. Our method provides a practical and efficient solution for scalable interpretability in large Transformer models, and holds promise for guiding further analysis and application of Large Language Models
45).
While FlowPrune effectively accelerates attention flow computation with minimal loss in accuracy for most cases, its relative error tends to correlate with the magnitude of the original attention flow values. In particular, when the original flow value is very small, the relative error introduced by FlowPrune can become excessively large, which may render the results unreliable for interpretability in those regions. This suggests that caution should be exercised when analyzing low-magnitude flows, as the pruning process may disproportionately distort their values.
Acknowledgment
This work is supported by the National Science Foundation of China (62576091). This research work also supported by the Big Data Computing Center of Southeast University. This work was also partially supported by the Southeast University Kunpeng & Ascend Center of Cultivation and the Big Data Computing Center of Southeast University.
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, 2017.
Long Beach, CA, USA, pages
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subhiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[3] Simeng Han, Frank Palma Gomez, Tu Vu, Zefei Li, Daniel Cer, Hansi Zeng. Chris Tar, Arman Cohan, and Gustavo Hernandez Abrego. Ateb: Evaluating and improving advanced nip tasks for text embedding models. arXiv preprint arXiv:2502. 16766, 2025.
[4] Jack W Rae. Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann. Francis Song. John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112. 1446, 2021.
[5] Baptist Alayrac, Jeff Donahue, Pauline Luc. Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch. Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in neural information processing systems, 15:23716-237366 2022.
[6] Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. 21ip-2: Bootstrapping language image pre- training with frozen image encoders and large language models. In International conference on machine learning. pages 19730-19742. PMLR, 2023.
[7] Deyao Zhu, Jun Chen, Xiaogian Shen, Xiang Li, and Mohamed Elhoseiny. Minigpt-4: Enhancing vision-language understanding with advanced large language models. arXiv preprint 2.V:::3304.1099 2023.
[8] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint vrXiv:2001.08360 2020.
[9] Yutao Sun. Li Dong, Shaohan Huang, Shuming Ma. Yuging Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint 2077023307.08621, 2023.
[10] Sarthak Jain and Byron C. Wallace. Attention is not explanation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 1344-6556, 2019.
[11] Sarah Wiegreffee and Yuval Pinter. Attention is not not explamation. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNL) pages 4090-4100, 2019.
[12] Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. arXiv preprint arXiv: aexivv/9966.000 1906,04284,, 2019.
[13] Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in bertology: What we know about how bert works. Transactions of the association for computational linguistics, 8:842-866, 2021.
[14] Samira Abnar and Willem Zuidema. Quantifying attention flow in transformers. arXiv preprint urXiv:2005.0092 2000220000000,,,222 2020.
[15] Joseph Cheriyan and Kurt Mehlhom, An analysis of the highest-level selection rule in the preflow-push max- -lx-flow algorithm. Information Processing Letters, 69(5):239-244 1999.
[16] Rich Caruana, Yin Lou, Johannes Gehrke, Paul Koch, Marc Sturm, and Noemie Elhadad. Intelligible models for healthcare: Predicting pneumonia risk and hospital 30-day readmission. In Proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining, pages 1721-1730 2015.
[17] Finale Doshi and Been Kim. Towards a rigorous science of interpretable machine learning. arXiv preprint 1702:v:/002.88000 2017.
[18] Lester R Ford Jr and Delbert R Fulkerson. Maximal flow through network. Canadian journal of mathematics, 8(3):399 404, 1956.

---

[19] Liu Liu, Zheng Qu, Zhaodong Chen, Yufei Ding, and Yuan Xie. Transformer acceleration with dynamic sparse attention, 2021.
[20] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie- Anne Lachaux. Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume lample. Llama: Open and efficient foundation language models, 2023.
[21] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jac Lec. Visual instruction tuning. Advances in neural information processing systems, 36:34892-34916, 2023.
[22] Joseph F DeRose, Jiayao Wang, and Matthew Berger. Attention flows: Analyzing and comparing attention mechanisms in language models, 2020.
[23] Sofia Serrano and Noah A Smith. Is attention interpretable? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), pages 2931-2951, 2019.
[24] Jasmijn Bastings and Katja Filippova. The elephant in the interpretability room: Why use attention as explanation when we have saliency methods? arXiv preprint arXi::2010.0500, 2020.
[25] Danish Pruthi, Mansi Gupta, Bhuwan Dhingra, Graham Ncubig, and Zachary C Lipton. Learning to deceive with attention- based explanations. arXiv preprint 190.:::999.071333 2019.
[26] Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability heyond attention visualization. In Proceedings of the IEEF/CVF conference computer vision and pattern recognition, pages 782-791, 2021.
[27] Qiu-Feng Wang, Xin Geng, Shu-Xia Lin, Shi- Yu Xia, Lei Qi, and Ning Xu. Learngene: From open-world to your learning task. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 1557-8565. 2022.
[28] Dzmitry Bahdanau. Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint 1XX.::400904473 2014.
[29] Hamidreza Ghader and Christof Monz. What does attention in neural machine translation pay attention to? arXiv preprint arxiv:/710.03:: 1710.03348, 2017.
[30] Simeng Han, Tianyu Liu, Chuhan Li, Xuyuan Xiong, and Arman Cohan. Hybridmind: Meta selection of natural language and symbolic language for enhanced I1m reasoning. arXiv e-prints, pages arXiv-2409, 2024.
[31] Jack Edmonds and Richard M Karp. Theoretical improvements in algorithmic efficiency for network flow problems. Journal of the ACM (JACM), 19(2):248-22) 264, 1972.
[32] Etim A Dinic. Algorithu for solution of a problem of maximum flow in networks with power estimation. In Soviet Math. Doklady, volume 11, pages 1277-1280, 1970.
[33] Andrew V Goldberg and Robert Tarjan. A new approach the maximum-floe flow problem. Journal of the ACM (JACM), 35(4):921 940, 1988
[34] Thomas H Cormen, Charles E Leiserson, Ronald L. Rivest, and Clifford Stein. Introduction to algorithmss MIT press, 2022.
[35] An Yang. Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang. Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma. Jianxin Yang. Jin Xu. Jingren Zhou, Jinze Bai, Jinzheng He. Junyang Lin, Kai Dang, Keming Lu, Keqin Chen. Kexin Yang. Mei Li. Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng. Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfci Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, Zhifang Guo, and Zhihao Fan. Qwen2 technical report, 2024.
[36] Peng Wang. Shuai Bai, Sinan Tan, Shijie Wang. Zhihao Fan, Jinze Bai, Keqin Chen, Xuejing I Liu, Jialin Wang, Wenbin Ge, Yang Fan, Kai Dang, Mengfei Du, Xuancheng Ren, Rui Men, Dayiheng Liu, Chang Zhou, Jingren Zhou, and Junyang Lin. Qwen2-vl: Enhancing vision-language model's 's perception of the world at any resolution, 2024.
[37] Karl Cobbe, Vineet Kosaraju. Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jenry Tworek. Jacob Hilton, Reichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint 10.1680, 2021.
12

---

[38] Kenneth Marino, Mohammad Rastegari. Ali Farhadi, and Roozbeh Mottaghi. Ok-vqa: A visual question answering benchmark requiring external knowledge. In Proceedings of the IEEE/eyf conference computer vision and pattern recognition, pages 3195-3204, 2019.
[39] Paul Jaccard. Étude comparative de la distribution florale dans une portion des alpes et des jura. Bull Soc Vaudoise Sci Nat, 37:547-579, 1901.
[40] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman, Glue; multi-task benchmark and analysis platform for natural language understanding, 2019.
[41] Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention, 2021.
[42] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: large-scale hicrarchical image database. pages 248-255, 2009.
[43] Yang, Hanwang Zhang, and Jianfei Cai. Deconfounded image captioning: A causal retrospect. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(1 (5(11))129)) 19:11996-13001 2021.
[44] Xu Yang, Yongliang Wu, Mingzhuo Yang. Haokun Chen, and Xin Geng. Exploring diverse context configurations for image captioning. Advances in Newral Information Processing Systems, 36:40944-4093, 2023.
[45] Fu Feng, Jing Wang, Xu Yang, and Xin Geng. Learngene: Inheritable "genes" in intelligent agents. Artificial Intelligence, page 104421, 2025.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The Flowprune algorithm introduced in the abstract and the introduction is presented Section The experiments are described in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
In Section 5 introduce the limitations of the Flowprune algorithm.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification: The time complexity of the attention flow using the HLPP algorithm is analyzed Section The acceleration in time due to the compression ratio is analyzed in Section 4.1.2 under Computational Cost (Time) part.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: In Section 4 we have provided all the information needed to reproduce the experiments.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We have provided open access to our code. the code link is provided in
abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
public/guider httgggudddssccoo httas//ooSSsmbbsssooooooooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
pttps.cc//publ haapollccggudesseeee http://ooeeoooossssooooeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We specify all the training and test details necessary to understand the results in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The models used are relatively large and the datasets are extensive. How- ever, the random seed settings for all experiments are identical to ensure the fairness of the comparisons.
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All the information required to reproduce the experiments is mentioned the beginning of Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We have read and complied with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The creators original owners of the assets we used have been properly credited. These can be verified in Section
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

A Supplementary Materials
To provide additional qualitative evidence for the effectiveness of the proposed FlowPrune framework, this appendix presents five supplementary examples for both Paraphrasing Verification and Vision Transformer (ViT) Heatmaps Analysis in Section 4 We also included an application scenario experiment on Question Answer Verification [22), which is also based on BERT. In addition, we provide a detailed description of the dynamic programming algorithm used to determine which attention layers to compress in FlowPrune.
A.1 Paraphrasing Verification
To further validate the effectiveness of FlowPrune in interpretability tasks, we conduct additional experiments on the MRPC dataset for paraphrasing verification. This task involves determining whether two input sentences convey the same meaning. For example, the sentence pair [CLS] It affected earnings per share by a penny. [SEP] The company said this impacted earnings by a penny a share. [SEP] is labeled as a paraphrase (positive example), since both sentences express semantically equivalent content.
In this setting, compute attention flow values between token pairs from both sentences using FlowPrune, and compare the results with those obtained from the original Attention Flow and Attention Rollout algorithms. The analysis is performed using both the pre-trained and fine-tuned BERT models. FlowPrune is applied to reduce the computational complexity while preserving essential interpretability signals.
Figures 9 through 13 show five representative examples of the resulting attention heatmaps. As demonstrated, FlowPrune consistently aligns with the original Attention Flow in highlighting key semantic alignments and distinguishing features between the sentences. Notably, the differences between pre-trained and fine-tuned models are also clearly reflected in the FlowPrune heatmaps, showcasing its ability to retain meaningful interpretability insights under compression.
(a) FlowPrune (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
(d) FlowPrunc (Fine-tuned)
(e) Attention Flow (Fine-tuned) (f) Attention Rollout (Fine-tuned)
Figure 9: Paraphrasing Verification Example 1. The two input sentences are: (1) "The driver, Eugene Rogers, helped to remove children from the bus, Wood said " and (2) "At the accident scene the driver was covered in blood" but helped to remove children, Wood said ."
21

---

319
(a) FlowPrune (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
K02
148
329
315
205
440
420
(d) FlowPrune (Fine-tuned)
(c) Attention Flow (Fine-tuned) (f) Attention Rollout (Fine-tuned)
The two input sentences are: (1) "Cooley said he expects Muhanmad will similarly be called as wimees at a pretrial hearing for Malvo "" and (2) Lee Boyd Malvo will be called as a witness Wednesday in a pretrial hearing for fellow sniper suspect John Allen Muhanmad !
Figure 10: Paraphrasing Verification Example 2.
42
29
(a) Plowerune&o (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
Anwrtin Ieloct
02
(d) FlowPrune (Fine-tuned)
(e) Attention Flow (Fine-tuned) (f) Attention Rollout (Fine-tuned)
Figure 11: Paraphrasing Verification Example 3. The two input sentences are: (1) "" If we don 't march into Tehran, think we will be in pretty good shape, he said " and (2) "" As long as we don 't march on Tehran, I think we are going to be in pretty good shape, he said."
22

---

(a) Flow Prune (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
(d) Prune (Fine-tuned)
(c) Attention Flow (Fine-tuned) (f) Attention Rollout (Fine-tuned)
Figure 12: Paraphrasing Verification Example 4. The two input sentences are: (I) "Sales for the quarter beat expectations, rising 37 percent year-on-year to 1.76 billion euros." ," and (2) "Sales rose 37 per cent year-on-year to 1.76bn, beating expectations."
Congrvee Almbe Teclnd
Minksate Stat
Stmder Bed Boat
(a) Flow (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
erc. Comgreuee Almbee Poetoud
Mlrrastio Bext
#inber Miled feat
900
(d) FlowPrune (Fine-tuned)
(e) Attention Flow (Fine-tuned) (f) Attention Rollout (Fine-tuned)
Figure 13: Paraphrasing Verification Example 5. The two input sentences are: (I) "If the MTA appeal to a higher court is successful, the $ 2 bus and subway hase fare won 'I be rolled back "" and (2) "If the MTA appeal is successful, the $ 2 bus and subway base fare won 't change." "
A.2 Vision Transformer (ViT) Heatmaps
To further evaluate the utility of FlowPrune in visual interpretability tasks, we provide additional examples using the Vision Transformer (ViT) model. Specifically, we apply FlowPrune to the DeiT-Small[4]] model on the ILSVRC2012 image classification dataset. The aim to generate attention heatmaps that reflect the relevance between each input image token and the /CLSS token, which is responsible for final classification decisions. In this setting, we compare three interpretability methods: FlowPrune, the original Attention Flow, and Attention Rollout.
23

---

Figures 14 throughh 18 illustrate five representative examples of attention heatmaps produced by these three methods. As shown in the visualizations, FlowPrune yields results that closely match those of the full Attention Flow, successfully highlighting semantically relevant regions in the image. In contrast, Attention Rollout generates more diftuse and less informative maps, often failing localize key visual areas.
These results demonstrate that FlowPrune not only retains the interpretability of the original attention mechanism but also achieves substantial computational savings, making it practical and efficient alternative for ViT visualization.
(a) Original Image
(b) Attention Flow
(c) FlowPrunc
(d) Attention Rollout
Figure 14: VIT Heapmap Example 1.
(a) Original Image
(b) Attention Flow
(c) FlowPrune
(d) Attention Rollout
Figure 15: VIT Heapmap Example 2.
S
(a) Original Image
(b) Attention Flow
(c) FlowPrune
(d) Attention Rollout
Figure 16: ViT Heapmap Example 3.
(a) Original Image
(b) Attention Flow
(c) FlowPrune
(d) Attention Rollout
Figure 17: VIT Heapmap Example 4.
24

---

(a) Original Image
(h) Attention Flow
(c) FlowPrune
(d) Attention Rollout
Figure 18: VIT Heapmap Example 5.
A.3 Question Answer Verification
The Question Answer Verification (QAV) experiment aims to investigate how fine-tuning affects BERT's ability to validate whether a given answer correctly addresses question. Using the QNLI [4) dataset from GLUE, the study fine-tunes a BERT-base model and employs Attention Flows to analyze changes in the model's attention mechanisms. The results show that fine-tuning enables the model to better focus on task-relevant details, such as key connecting words and phrases between the question and answer, thereby improving its accuracy in determining the validity of the answer.
In this experiment, we continue to utilize a 12-layer BERT model. We aim to compare the outcomes of FlowPrune and Attention Rollout with those of the original Attention Flow, both in pre-tre pre-trained and fine-tuned models. The primary focus of this task is the attention metrics from tokens in the initial layer that point to the /cls] token in the final layer. Tokens with the highest computed attention metrics are deemed the most significant.
Table 3: Max-n Overlap for Question Answer Verification. FP represents Flowprune and AR represents Attention RMean + SE)
Max-n
FP Pretrained
AR Pretrained
FP Fine-tuned
AR Fine-tuned
98.05 + 1.11%
64.68 + 2.91%
96.10 + 1.54%
85.71 1 2.59%
n = 5
97.31 + 1.54%
84.48 + 2.14%
94.16 2.31%
76.31 + 2.21%
10
97.01 + 1.72%
66.50 + 1.88%
93.42 2.60%
71.03 + 1.76%
P15
(a) FlowPrune (Pretrained)
(b) Attention Flow (Pretrained) (c) Attention Rollout (Pretrained)
60% Compunter Amentee Fewieedd
Ateicon Aobral Resc
(d) FlowPrunc (Fine-tuned)
(e) Attention Flow (Finc-tuned) (f) Attention Rollout (Finc-tuned)
Figure 19: Question Answer Verification. Visualization of attention flow results across three FlowoussFFoowrunee, original Attention Flow, and Attention Rollout.
We quantify the Max-n Overlap between the most salient tokens derived via FlowPrune/Attention Rollout and the original Attention Flow. For FlowPrune, retain rate of 67% is employed, thereby
25

---

condensing the original 12-layer attention map into 8 layers. Analogous to the Top-K Token Retention method. the Jaccard overlap (IOU) utilized as the evaluation metric. The results are presented in Additionally, an illustrative example from the MPRC dataset is depicted in Figure T9
Table 3
A.4 Layer Selection via Dynamic Programming for FlowPrune
To identify which attention layers compress in the FlowPrune framework, formulate the selection process as a dynamic programming (DP) problem that aims to preserve critical regions in the attention flow graph while reducing computational cost.
As described in the main text, the attention structure Transformers forms a strictly layered directed acyclic graph (DAG)), in which each token in layer i attends all tokens in layer i 1. This layered property enables ayse-wwiee flow analysis and facilitates compression strategies based on edge removal and reconnection.
To estimate the relative importance each layer, we apply a sampling-basee heuristic inspired by the Max-Flow Min-Cut Theorem. Specifically, we sample multiple source-to-sink paths through the attention graph and treat the lowest-capacity edge on each path as a potential min-cul edge. For each intermediate layer we count the number of such edges it contains. A lower count indicates that the layer contributes less to the overall minimum cut and is thus a better candidate for compression.
aL-2) where the first and last layers are excluded from Given these layer- wise scores (a1, ........ compression for structural consistency, our goal select N layers to retain (i.e., not compress) such that the total number of potential min-cut edges in the retained layers is maximized.
We define the dynamic programming state as:
f(i, 0): the maximum total number of potential min-cut edges when the first layers are reduced to j layers, and the i-th layer is compressed.
fi, 1)): the same, but the i-th layer is retained.
The recurrence relations are:
((i,j,0) max (1-
f(i,j,1) max (fii-1,j-100) f(i- f(i-1,
1)) +
The boundary conditions are:
f(0,0,0) 00, f(0,0,1)
f(i,j.0) = -00 if
f(i,j,1) -00 {
f(1,1,0) =0 for
The optimal value given by:
max (L,,L,,,) N,L,
where L is the total number of intermediate layers (excluding the first and last).
To reconstruct the compression scheme, backtrack from the optimal DP state and trace the transition path:
If the current state is (1, n. 0), we move to the maximum of f(l-1, n, 0) or (1-2, and record that layer / is compressed.
If the current state (1, n, 1), we move to the maximum of 10-1,,-11,0 or fl-11, and record that layer is retained.
This dynamic programming approach ensures that layers essential to maintaining attention flow- those with higher concentrations of potential min-cut edges-are preserved. Meanwhile, layers with relatively minor contributions are removed.
26


Reviews:
[{"review_id": "moRryBGFwr", "reviewer": "Reviewer_b9Ta", "summary": "This paper focuses on the computational inefficiency of attention flow, an interpretability technique, when applied to large-scale transformer models. To handle this, the authors exploit two properties of transformers, i.e., the softmax activation in attention map and the strictly layered-structure. Based on these, they propose Edge Pruning to prune less informative token-to-token connections via a predefined threshold and Layer Compression to remove layers with minimal contribution to attention flow. Together, the two strategies reduce the size of attention graph, thereby mitigating the computational complexity. Experiments on LLaMA and LLaVA validate the efficiency of the proposed method.", "questions": "See weaknesses for details.", "limitations": "None.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1.\tThe problem that FlowPrune investigates is of practical impact, which helps to reduce the running time for interpreting behaviors of large-scale transformer models.\n2.\tMultiple well-designed metrics are used to evaluate the efficiency and interpretation quality.\n3.\tThe paper is well organized and easy to follow.\n\n\nWeaknesses:\n1.\tThe pruning strategy of FlowPrune’s is relatively static, especially the threshold used in the Edge Pruning. Besides, the method primarily relies on the properties of softmax sparsity and layer-wise structure, which may limit its effectiveness when applied to specialized Transformer variants (e.g., sparse or hybrid attention models).\n2.\tLack of discussion of recent related work and the compared methods are not recent, e.g., the Attention Rollout in 2020.\n3.\tIs the proposed method sensitive to the edge pruning threshold? Besides, what is the memory complexity?\n4.\tMinor suggestions: “Max-Flow Min-Cut Theorem” and “Highest-Level Preflow Push algorithm” first appear without references. In addition, there are some syntax errors, e.g., “(Line 42) ..., [14] propose ...” and “(Line 94) ..., Push–Relabel methods [33], which maintain ...”", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer:\n\nThank you for your constructive comments, Your valuable comments have been of considerable assistance to us. Then, to your concerns, we address the comments as follows:\n\n**1. Static Threshold.** We use a static threshold primarily because experimental results show that FlowPrune is not particularly sensitive to its value. Among the tested values, $1 \\times 10^{-6}$ consistently performs well and offers a reasonable balance. Given this robustness, using a static threshold simplifies implementation without compromising performance.The table below illustrates the average absolute error for the LLaMA model under different edge pruning thresholds after retaining half of the layers (16 layers).\n\n|1e-8|1e-7|1e-6|1e-5|\n|---|---|---|---|\n|0.0137|0.0099|0.0097|0.0103|\n\n**2. Effectiveness.** Firstly, the dominant Large Language Models (LLMs) and Large Vision-Language Models (LVLMs) currently in use primarily employ standard Transformer blocks, which do not incorporate sparse or hybrid attention mechanisms. Secondly, and more importantly, most existing sparse or hybrid attention mechanisms,  such as those used in Longformer[A] and Natively Trainable Sparse  Attention (NSA)[B], also exhibit softmax sparsity and layer-wise  characteristics. The layer-wise characteristic is applicable to almost  all Transformer architectures. As for softmax sparsity, as long as softmax is utilized in the  architecture, softmax sparsity is inevitable. For example, Longformer  applies softmax to the sparsified attention matrix, and NSA selectively  applies softmax to different attention paths.\n \n[A]  Longformer: The long-document transformer.\n\t\n[B] Natively Trainable Sparse Attention for Efficient Transformer. \n\n**3. Related Studies.** While we may seem to overlook related studies about attention flow or attention rollout, this is because recent research has largely focused on *using* attention flow or rollout as tools for analyzing model behavior or supporting new algorithmic designs, rather than on *improving* the computation of attention flow itself. Thus, these methods are typically treated as fixed, with limited consideration given to its efficiency or scalability. However, FlowPrune is introduced precisely to fill this overlooked gap by offering a more efficient and principled method for attention flow analysis, facilitating its broader application in large-scale interpretability tasks.\n\n**4. Memory Complexity.**  The memory complexity of FlowPrune is proportional to the size of the attention graph, as its primary memory usage arises during graph construction. The attention graph contains $O(LN)$ nodes and $O(LN^2)$ edges, resulting in an overall memory complexity of $O(LN^2 + LN) = O(LN^2)$.  \n\n**5. Minor suggestions.** We apologize for these grammar mistakes and missing citations and thank you for pointing them out. They will be corrected in revision to make overall writing meet the standards."}, {"review_id": "MlM64gzxTU", "reviewer": "Reviewer_suGC", "summary": "This paper presents FlowPrune, a framework designed to accelerate attention flow computation in Transformers by pruning attention graphs prior to applying max-flow algorithms. The approach integrates two techniques: Edge Pruning, which removes low-capacity attention edges, and Layer Compression, which eliminates layers contributing minimally to the flow. These strategies collectively reduce computational complexity while maintaining interpretability fidelity. The authors evaluate FlowPrune on case studies across both NLP and vision tasks, demonstrating that it yields interpretability results consistent with those of the original Attention Flow method.", "questions": "The authors mentioned that \"Several approximations, including Attention Rollout [14], have been proposed to reduce complexity, though they often sacrifice theoretical properties such as flow conservation.\" I’m wondering whether the proposed approach provides any theoretical guarantees on the approximation error.", "limitations": "Although some limitations are acknowledged, a more comprehensive discussion would strengthen the paper.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength: \n- The target problem of accelerating attention flow computation is well-motivated, given its high computational complexity.\n\n- The paper is clearly presented, with good writing quality and logical organization. Additionally, the public release of the code enhances the reproducibility and practical value of the work for the research community.\n\n- To the best of my knowledge, the use of the Max-Flow Min-Cut approach in this context is among the early explorations, making it a novel direction.\n\n\nWeakness: \n- The choice of the Max-Flow Min-Cut approach requires further justification. It would strengthen the paper if the authors could clarify why this method is preferable over other graph sparsification or general sparsification techniques. Including additional ablations comparing against these alternatives would help improve the paper’s rigor.\n\n- The generalizability of the proposed method warrants further evaluation. Currently, the experiments are limited to LLaMA and LLaVA models. Even within this small-scale evaluation, performance patterns vary significantly (e.g., Fig. 4). It remains unclear whether the proposed method consistently performs well across different models and tasks. Broader evaluation on additional benchmarks would be beneficial.\n\n- The paper would benefit from a more thorough discussion of the error distribution in low-attention regions. As shown in Fig. 4, the relative errors are quite high for low attention values. The authors should discuss potential causes, limitations, and possible mitigation strategies for this issue.", "was_revised": true, "final_justification": "The response addresses most of my concerns; however, the justification for choosing the max-flow/min-cut approach over alternative solutions remains insufficient, leaving their design choice without adequate motivation.  Also, the claim regarding theoretical guarantees appears overstated. Based on these points, I will maintain my original score.", "rebuttal": "Dear reviewer:\n\nThank you for your constructive comments, Your valuable comments have been of considerable assistance to us. Then, to your concerns, we address the comments as follows:\n\n **1.Other Sparsification Techniques.** We build FlowPrune on the Max-Flow Min-Cut Theorem because attention flow is inherently formulated as a maximum flow problem over the attention graph[14]. As a cornerstone, the Max-Flow Min-Cut Theorem provides the primary theoretical foundation for designing and analyzing max-flow algorithms. To our knowledge, few alternative frameworks offer comparable guidance for optimizing max-flow computations. Accordingly, in developing FlowPrune, we focused on this theorem in conjunction with exploiting the structural properties of Transformer attention graphs, rather than adopting other graph sparsification techniques that are less directly aligned with the max-flow formulation.\n\n**2. Generalizability.** To evaluate the generalizability of FlowPrune, we have selected two representative large models: the unimodal LLaMA and the multimodal LLaVA. We also assessed its effectiveness on smaller-scale models—BERT for NLP (Sec 4.2.1) and ViT for vision (Sec 4.2.2). In addition to measuring absolute error against the true attention flow (Fig 4), we considered that attention flow is often used to study relative ranking; thus, we introduced corresponding metrics for this purpose (Fig 5, 6). Together, these diverse models and evaluation criteria support the general applicability of our method. In response to your suggestion, we also included experiments on two widely used large models, Qwen (unimodal) and Qwen-VL (multimodal). The following table presents the absolute and relative errors in the high attention regions, which are of the most interest to researchers, for the Qwen and Qwen-VL models, where we can find that the absolute/relative error is not large.\n   \n| Region\\Model| qwen (absolute error) | qwen (relative error) | qwen-vl (absolute error) | qwen-vl (relative error) |\n| --- | --- | --- | --- | --- |\n[0.9-1.0] | 0.003 | 0.31% | 0.001 | 0.12% |\n[0.8-0.9] | 0.014 | 1.63% | 0.012 | 1.41% |\n[0.7-0.8] | 0.017 | 2.28% | 0.013 | 1.71% |\n\nIn addition to the aforementioned experiments, we also included an application scenario experiment on Question Answer Verification, which is also based on BERT. You can find the detailed introduction of the experiment in [22]. We compared the Top-K Token Retention with Attention Flow of FlowPrune and Attention Rollout, and organized the results in a manner similar to *Table 1* in Section 4.2.1. The results show that, compared to Attention Rollout, the FlowPrune method is much closer to the analysis results of Attention Flow, demonstrating FlowPrune outperforms Attention Rollout. The table below presents the results:\n\n|Max-n|FP Pretrained(Mean ± SE)|AR Pretrained(Mean ± SE)|FP Fine-tuned(Mean ± SE)|AR Fine-tuned(Mean ± SE)|\n|---|---|---|---|---|\n| n = 3 | 0.9805 ± 0.0111 | 0.6468 ± 0.0291 | 0.9610 ± 0.0154 | 0.8571 ± 0.0259 |\n| n = 5 | 0.9731 ± 0.0154 | 0.8448 ± 0.0214 | 0.9416 ± 0.0231 | 0.7631 ± 0.0221 |\n| n = 10 | 0.9701 ± 0.0172 | 0.6650 ± 0.0188 | 0.9342 ± 0.0260 | 0.7103 ± 0.0176 |\n\n**3. Difference between llava and llama.** For the differences in absolute error patterns between LLaMA and LLaVA in Fig. 4, we think this is cased by the distinct attention graphs of the two models. LLaMA is a unimodal language model, while LLaVA is a multimodal vision-language model, resulting in inherently different attention graph structures. We hypothesize that the higher errors observed in LLaVA stem from its large number of input vision tokens, which often contain redundant information and contribute to an abundance of low-attention values. As discussed in **4.Low Attention Regions.**, FlowPrune exhibits higher approximation error in low attention regions, thereby increasing the overall error in LLaVA. A similar phenomenon appears in our experiments with Qwen (unimodal) and Qwen-VL (multimodal), reinforcing the notion that attention graph modality plays a key role in shaping the error pattern.\n\nIn contrast, both LLaMA and LLaVA show similar trends in *relative error*. As highlighted in Section 4.2, relative error is of greater practical importance than absolute error, particularly for interpretability and token-level ranking tasks based on attention flow.\n\n**4. Low Attention Regions.** As discussed in Limitations (lines 336–341), FlowPrune exhibits relatively larger approximation error in low-attention regions. Here, \"low-attention regions\" refer to token pairs with small flow values in the original attention flow, indicating weak attentional connections. In practice, such pairs are typically of limited interest. As noted in lines 247–251, researchers are primarily concerned with token pairs exhibiting large attention flow values, where FlowPrune achieves much lower error. Therefore, the observed approximation inaccuracies in low-attention areas have minimal impact on practical applications.\n\n**5. Theoretical Guarantees.** We apologize for any confusion caused by our use of the term *theoretical guarantees*. To clarify, we do not claim a general guarantee on the approximation error, as this depends heavily on the specific numerical structure of the model’s attention graph and is analytically intractable. Rather, our point is that FlowPrune is grounded in the Max-Flow Min-Cut Theorem and thus offers a principled approximation of attention flow as originally defined in [14]. In contrast, Attention Rollout relies on recursively multiplying attention matrices, which is heuristic in nature and diverges significantly from the original flow formulation. Therefore, the *guarantee* we refer to is that FlowPrune adheres to the theoretical definition of attention flow, rather than serving as a purely heuristic alternative."}, {"review_id": "wlvnM6zr9U", "reviewer": "Reviewer_Ttwq", "summary": "This paper proposes FlowPrune, a method to accelerate the computation of attention flow, an interpretability method in NLP and vision which tries to incorporate the attention information from all layers of the model when computing interpretable heatmaps over the inputs. The way FlowPrune achieves speedups is twofold: (1) it prunes edges in the attention graph for which the attention weights fall below a certain threshold and (2) it removes layers of the Transformer for which the edges contained in the layer are unlikely (computed through some heuristic) to contain any of the min-cut edges. The authors test for speedup and accuracy on the LLama and Llava models, and use BERT and ViT for the application domains. Overall, both quantitatively and qualitatively, they find FlowPrune to match the full attention flow results pretty well.", "questions": "**Clarifications**\n1. Line 193: \"... we provided 1,000 inputs ...\" -> \"inputs\" means questions from the dataset here?\n2. Line 212: does \"samples\" mean questions from the dataset here?\n3. Figure 5: are these plots averages over the 1000 samples drawn from the dataset? Because if not, how can the top3 be in the 90s? It should only be able to be 0, 1/3, 2/3, or 1.\n\n**Grammar**\n1. Line 29: change \"propagateed\" to \"propagated\"\n2. Line 136: change \"In Transformer, softmax operation\" to \"In Transformers, the softmax operation\"\n3. Caption of Figure 1: change \"pathes\" to \"paths\" and change \"brwon\" to \"brown\"\n4. Line 162 - 163: \"However, in more realistic ... this assumption less precise.\" -> this sentence doesn't read quite right grammatically, maybe rephrase.\n5. Line 227: change \"may retains\" to \"may retain\"?\n6. Line 306: change \"expriments\" to \"experiments\"", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n### Quality\n1. The quantitative results in Table 1 and 2 are very convincing: they clearly demonstrate FlowPrune is a solid improvement over Attention Rollout compared with the original attention flow. Qualitatively, the heatmaps in Figure 7 and 8 also look quite convincing.\n2. The metrics presented in Section 4.1.1 are well-designed to measure the relevant factors of their method (speedup + error analysis).\n3. Line 247 - 251: I like that the authors connect the low approximation errors for high attention flow values to what actually matters in practice, which are the top source-destination pairs with the highest attention flow values.\n\n### Clarity\n1. Generally, I find the paper to have clear graphs and tables that are easy to read and clearly convey the results.\n\n### Originality\n1. I generally like the authors' idea of using insights from graph theory to improve computational efficiency, and the authors' use of max-flow min-cut seems pretty original to me!\n\n### Significance\n1. Given the strong results, I imagine this could become the current default method when computing attention flow which seems impactful.\n\n## Weaknesses\n### Quality\n1. Line 163 - 165: \"Nevertheless, attention graphs ... to all tokens in layer i + 1\". Two things: (1) typical Transformer architectures (e.g. Llama) are NOT fully connected between layers since they use causal attention masks, and (2) shouldn't this be from i + 1 to i? (1) pointed out above also draws into the question the line that follows: \"This near-uniform connectivity implies that sampled paths are likely to overlap on any edge.\" How are GPT-style models, which don't have near-uniform connectivity, affected here?\n2. Line 166 - 168: \"As a result, the edge ... more likely to lie on the global min-cut, ...\" -> Is this true? Do the authors have evidence or proof for this?\n3. Line 219: \"The empirical results in Figures 2 and 3 generally align with this analysis.\" -> Is this really the case? It doesn't seem to me the curves in Figure 2 are really following the predicted $1/R^{2.5}$ speedup prediction? Maybe the authors could include the theoretical prediction in Figure 2 so this is easier to compare? The authors also state that \"LLaVA's performance closely matches the theoretical expectation\" but I still don't agree based on the $1/R^{2.5}$ factor. \n4. Line 243 - 246: \"(2) In Layer Compression, newly added edge ... not decreases\" -> Why do we set them all to one? This seems to break the constraint on attention weights that they should sum to one? And also the last part says that \"this ensures that the maximum flow from the compressed graph increases, not decreases\", is this desirable? Or why is this not a bad thing? Ideally I thought we want the max flow to stay roughly the same, not increase.\n5. Is it possible to add standard errors to Table 1 and 2?\n6. The authors discuss two techniques for reducing the attention graph: (1) edge pruning and (2) layer reduction. However, it seems there is no experiments specifically investigating (1)?\n\n### Clarity\n1. Line 57 - 58: \"due to the softmax operation, many attention weights are close to zero ...\" -> I think this is stated a few times in the paper, but it's never really made clear *why* this is due to the softmax operation. Could the authors provide a short explanation for this? Is this due to the exponential in the softmax?\n2. Line 127 - 130: \"This theory provides ... without minimum cut.\" -> Why is this true? Or is this just an intuition? Can't removing stuff affect what constitutes the minimum cut?\n3. Figure 1(b): why are the orange lines here the min-cut edges? And which is the source and sink in this diagram?\n4. Line 235: \"... overhead becomes negligible after computing more than 6 casees for Llama and 2 cases for LLaVA.\" How did the authors get these numbers? Can you help me understand how to get these from the graphs?\n5. Line 266 - 268: \"This result is fully sufficient for practical applications with lower precision requirements, whilte the results at a compression rate 50% are adequate for tasks with higher precision demands.\" -> What does \"precision requirements\" mean here? And how do we know its fully sufficient?\n\n### Significance\n1. Line 120-121: \"this high cost makes the method impractical for large-scale Transformer models\" -> Could the authors shed some light on how computationally expensive this would actually be for some large models? For example, in wall clock time, how long woudl it take to compute the full heatmap for a single example say? I listed this under significance since the answer to this question can have some impact to how significant the speedups really are, and hence what impact the method will have in the field.", "was_revised": true, "final_justification": "The weaknesses brought up in my initial review have been sufficiently addressed, and hence I have increased my score to 5. \n\nThe main key issue that was addressed was that I was initially not convinced the paper was really using any property from the Max-Flow Min-Cut Theorem, but the authors explained that lines 127 - 130 are really a corollary, not just an intuition. I would encourage the authors to include this explanation/proof into the next revision of their paper.", "rebuttal": "Dear reviewer:\n\nThank you for your constructive comments, Your valuable comments have been of considerable assistance to us. Then, to your concerns, we address the comments as follows:\n\n**1. Attention Flow is calculated after token-by-token generation instead of during generation.** In fact, when using attention flow for analysis, the entire sequence is input into the Transformer **in a single parallel forward pass** to obtain the corresponding attention graph. Consequently, the causal attention mask used in autoregressive generation is not applied here. Furthermore, since attention flow tracks information propagation from the first to the last layer, the attention graph is formed by edges connecting layer $i$ to layer $i+1$, reflecting the forward flow of information within the model.\n\n**2. Question about Line 166 - 168.** As noted in lines 159–168, this is a heuristic conjecture for which we currently lack a formal proof. Nevertheless, both our fundamental evaluations and experiments support its validity. Specifically, the experiments on LLaMA and LLaVA in Section 4.1 show that FlowPrune have low approximation errors, suggesting that the proposed approach is empirically reliable. \n\n**3. Bound in Lines 216-223.** The bound presented in lines 216–223 serves as an estimate, based on the assumption that when the proportions of pruned edges and nodes in FlowPrune are approximately equal, the bound reasonably approximates the actual speedup. Additionally, the specific structure of the attention graph can influence the extent of acceleration achieved. We will incorporate the above explaination in revision.\n\nIn lines 224–228, we briefly analyze possible reasons why the LLaMA curve deviates from the theoretical prediction. We attribute this discrepancy to differences in input token types across models. Unlike LLaMA, LLaVA is a multimodal vision-language model that processes a large number of visual tokens, many of which are redundant. This redundancy results in a higher proportion of near-zero attention weights, leading to more edges being pruned, making this case more similar to the assumption. As a result, the compressed attention graph in LLaVA aligns more closely with our theoretical assumptions.\n\n**4. Setting edge capacity to 1.** According to the Max-Flow Min-Cut Theorem, increasing the capacity of edges **not in the minimum cut** does not affect the maximum flow value, as the flow is bounded by the capacity of the minimum cut. In FlowPrune, we ensure that the removed parts contain as few minimum cut edges as possible.  To maintain the maximum flow value after compression, newly added shortcut edges are assigned sufficiently large capacities to ensure that they do not become part of the new minimum cut. In our case, a capacity of **1** is considered sufficiently large.\n\nNote that removal involves deleting certain vertices along with their connected edges. For any pair of remaining vertices that were previously connected via a path through the removed nodes, we introduce a new edge with large capacity to restore connectivity. Here, a capacity of 1 is adequate because we do not compress the first attention layer. As a result, the total flow emitted from the source remains bounded by 1, meaning the total maximum flow will also not exceed 1. Consequently, added edges with capacity 1 will not appear in the minimum cut. While this breaks the constraint that attention weights should sum to one, preserving the first layer ensures the attention flow remains within [0, 1], maintaining its interpretability.\n\n**5. Increasing Max-Flow is Better than Decreasing.** Since the maximum flow is bounded above by 1 in both the original and compressed attention graphs, a relatively large original flow value restricts the possible outcomes of FlowPrune to a narrower range, thereby reducing potential error. As discussed in lines 247–251, such cases are considered important, and thus, smaller errors are particularly desirable.\n\n**6. Add Standard Errors.** Thank you for your suggestion. We will add the following results about standard errors in revision.\n\n|Max-n|FP Pretrained(Mean ± SE)|AR Pretrained(Mean ± SE)|FP Fine-tuned(Mean ± SE)|AR Fine-tuned(Mean ± SE)|\n|---|---|---|---|---|\n|n = 3|0.9850 ± 0.0086|0.6340 ± 0.0258|0.9500 ± 0.0150|0.8450 ± 0.0232|\n|n = 5|0.9793 ± 0.0119|0.8471 ± 0.0183|0.9250 ± 0.0226|0.7662 ± 0.0193|\n|n = 10|0.9769 ± 0.0133|0.6638 ± 0.0158|0.9145 ± 0.0258|0.7085 ± 0.0156|\n\n|Max-n|FP(Mean ± SE)|AR(Mean ± SE)|\n|---|---|---|\n|n = 10|1.0000 ± 0.0000|0.2818 ± 0.0306|\n|n = 20|0.9953 ± 0.0012|0.5084 ± 0.0374|\n|n = 30|0.9959 ± 0.0011|0.7852 ± 0.0343|\n\n**7. Edge Pruning Experiment.**  Please see ``**1.Static Threshold.**'' in response to Reviewer b9Ta for more details.\n\n**8. Small Attention Weights.** Your observation is accurate. Due to the exponential nature of the softmax function, only a few    attention values are significantly greater than zero. For example, in our experiments across the four models used in this study (LLaMA, LLaVA, BERT, and ViT), more than 95% of the attention values are smaller than 1 \\times 10^{-6}.\n\n**9. Intuition in Lines 127 - 130.** We acknowledge that the term we used was imprecise; what we intended to convey is more accurately described as a **corollary**, rather than an intuition. Specifically, we refer to the removal of certain vertices from the graph, along with their incident edges. For the remaining vertices, if any pair was previously connected via a path composed of the removed edges, we introduce a new edge with sufficiently large capacity to preserve connectivity. Under this construction, if none of the removed edges belong to the original minimum cut, the minimum cut remains unchanged. This holds because, after the removal, any S–T cut must either  1) includes the newly introduced edges and, due to their sufficiently large capacity, cannot be the minimum cut, or 2) is composed entirely of original edges and thus cannot be smaller than the minimum cut of the original graph. \n\nSince the minimum cut of the original graph indeed exists within the graph, it remains the minimum cut of the graph after the removal operation, thereby preserving the invariance of the maximum flow.\n\n**10. Question About Figure 1(b).** We apologize for the confusion. The figure illustrates our sampling-based heuristic algorithm. The orange lines indicate *potential* min-cut edges, not actual min-cut edges. This algorithm is introduced in lines 153–158, with its motivation discussed in lines 159–168. Note that the figure does not include predefined source or sink nodes; these are determined only when computing the attention flow between a specific pair of tokens.\n\n**11. Overhead of Flowprune.** For a single attention graph, pruning is performed only once, after which the sparsified graph is reused to compute attention flow between all token pairs. Although pruning introduces initial computational overhead, this cost is amortized as the number of evaluated token pairs increases. For example, in the case of LLaMA, computing attention flow for 6 token pairs on the original attention graph takes approximately the same time as applying FlowPrune once and then performing max-flow computations on the sparsified graph for those pairs. As more pairs are calculated, the average time per pair for FlowPrune continues to decrease. This trend is shown in Figure 3, where we plot average per-pair computation time under various compression rates. The intersection with the baseline curve (original attention flow time) indicates the point beyond which FlowPrune offers a computational advantage.\n\n**12. Precision Requirements.** As discussed in lines 68–70, attention flow is often used to assess the relative importance of tokens based on the ranking of flow values. Our notion of precision pertains specifically to this context. At a 50% compression rate, the first instance of incorrect relative ordering appears only after the top 40% of token pairs. In practice, applications that rely on attention flow typically focus on the top 5–10 token pairs. Therefore, an FDR of 0.4 is sufficient for these use cases. \n\nThe comma in lines 266-268 you cited, is a typo. The lines you cited should be two separate sentences. The corrected version should be \"This result is fully sufficient for practical applications with lower precision requirements. While the results at a compression rate 50% are adequate for tasks with higher precision demands.\".\n\nThe first sentence discusses the scenario of retaining only 3 layers, while the second sentence addresses the situation with a compression rate of 50%. We apologize for any confusion this may have caused.\n\n**13. Actual Run Time.** The maximum flow algorithm is a classical graph-theoretic method that typically lacks efficient GPU-accelerated implementations. Our experiments were conducted on a server equipped with two AMD EPYC 7453 28-core processors. Using LLaVA as an example, we measured the runtime for computing attention flow across 100 token pairs in a sample containing 400–500 tokens within a 32-layer LLaVA model. With the original attention flow algorithm, the runtime reaches 112 minutes. In contrast, compressing the attention graph to 12 layers reduces the runtime to approximately 10 minutes.\n\n**14. Input and Samples.** We construct input sequences by concatenating questions with their corresponding answers from the dataset. Each input generates a corresponding attention graph, which we refer to as a sample.\n\n **15.Average Results.** Indeed, that is precisely the case.\n\n **16.Grammar.** Thanks for pointing out these mistakes and we will follow your suggestion to revise them."}, {"review_id": "dxrH55vD8t", "reviewer": "Reviewer_qaDf", "summary": "The paper proposes FlowPrune, an efficient variant of the original Attention Flow interpretability technique. Attention Flow is a method for analyzing how information propagates through all layers of a Transformer model by modeling the entire network as a directed graph. Attention Flow has $O(L^{2.5}N^3)$ complexity for $L$ layers and $N$ tokens, making it impractical for large models. To address this, the authors propose FlowPrune, with two modifications over the original method: Edge Pruning and Layer Compression, both designed to reduce the attention graph while preserving its critical flow structure, giving a complexity of $O(L^{2.5}N^3/R^{2.5})$.\n\nThe authors evaluate FlowPrune on LLaMA and LLaVA across multiple compression rates, measuring time speedup, approximation accuracy, and practical utility in real-world scenarios including paraphrasing verification and Vision Transformer heatmap generation.", "questions": "Does the upper-bound method, Attention Flow, perform comparatively to other established input attribution techniques?\n\nIt is counterintuitive that higher compression rates, close to 1, reflect more layers retained. Typically, I would associate higher compression rates with keeping less number of layers.\n\nl218: \"L/R layers reduces\" -> \"R/L layers reduces\"\n\nTypos:\n\nl29: propagateed -> \"propagated\"\n\nInconsistent use of \"LLama\" and \"Llama\"", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The two proposed techniques, edge pruning and layer compression, are theoretically sound and well-motivated by the structural properties of Transformer architectures.\n\n- The evaluation uses well-designed metrics for comparing against ground-truth attention flow, demonstrating that FlowPrune maintains high fidelity to the original method even at aggressive compression rates (few number of layers retained in the attention graph).\n\nWeaknesses\n\n- The paper does not justify why attention flow is preferable to established attribution methods like gradients or integrated gradients. The original approximation method, \"Attention Rollout\" ([Abnar et al., 2020](https://aclanthology.org/2020.acl-main.385/)) has been shown to underperform compared to alternative methods ([Chefer et al., 2021](https://openaccess.thecvf.com/content/CVPR2021/papers/Chefer_Transformer_Interpretability_Beyond_Attention_Visualization_CVPR_2021_paper.pdf), [Mohebbi et al., 2023](https://aclanthology.org/2023.eacl-main.245/)). A key limitation of attention flow, and by extension FlowPrune, is its inability to provide class-specific explanations, which restricts its applicability. Without a compelling rationale for using attention flow, accelerating this particular method offers limited practical value.\n\n- The layer compression component removes intermediate layers, undermining the method's ability to analyze actual information flow through the original network architecture.\n\n- Results are showed depending on layer compression rate, but not with respect to edge pruning threshold. You say the edge pruning threshold is set to 1e−6 throughout all experiments. How did you choose this threshold? What is the impact of this threshold in the metrics?\n\n- The evaluation is limited to two models (LLaMA, LLaVA), broader evaluation across architectures would strengthen claims.", "was_revised": true, "final_justification": "The authors have provided useful clarifications and additional evaluations that strengthen the work. However, the main concern remains, the paper does not sufficiently justify the choice why attention flow is preferable to established attribution methods. Thus, a rating of 4 is appropriate.", "rebuttal": "Dear reviewer:\n\nThank you for your constructive comments, Your valuable comments have been of considerable assistance to us.\n\nFirst, we would like to emphasize the novelty of our FlowPurne method. This method helps to reduce the running time for interpreting behaviors of large-scale transformer models. Multiple well-designed metrics are used to evaluate the efficiency and interpretation quality.\n\nThen, to your concerns, we address the comments as follows:\n\n**1. Attention flow as interpretability tool.** In [14], Attention Flow is defined by modeling the multi-layer self-attention mechanism of the Transformer as a capacitated flow network, where attention weights and residual connections jointly determine edge capacities. The maximum flow between tokens of interest serves as the interpretability metric, quantifying how much information is transmitted between tokens. This formulation has been adopted in works such as [22] to analyze and compare attention mechanisms. However, computing attention flow is computationally expensive, prompting many researchers to use Attention Rollout as a faster alternative.\n\nHowever, as you have pointed out, attention rollout often performs poorly. This is because it is a heuristic approximation that relies on direct multiplication of attention matrices, which causes the resulting values to deviate from true attention flow. Actually, this is why we propose *FlowPrune*, which is a method grounded in the Max-Flow Min-Cut theorem that can significantly accelerates attention flow computation while providing more accurate approximations of the true flow.\n\n**2. Class-specific Explanations.** For class-specific explanations, our experiments on BERT (lines 276–301, Table 1) were specifically designed to evaluate the effectiveness of FlowPrune in such contexts. These experiments aimed to identify which tokens are most influential in classification decisions. FlowPrune achieved strong performance on this task, significantly outperforming Attention Rollout. Additionally, as demonstrated in Figure 8 and Figures 6–10 of the Appendix, FlowPrune can effectively highlight the visual regions contributing to the final predicted labels in image classification as the original attention flow, while attention rollout fails . These results support the conclusion that FlowPrune provides meaningful class-specific explanations.\n\n**3. Layer Compression.** As discussed in lines 124–129 and 150–158, by leveraging the Max-Flow Min-Cut Theorem alongside the Transformer’s structural properties, we remove layers that have minimal impact on information flow. This selective pruning aligns with our goal of eliminating layers that contribute little to the results.\n\n**4. Edge Pruning Threshold.** In fact, we have conducted experiments on the selection of the edge pruning threshold, but due to page limitations, these results were not included in the paper. The primary reason for omitting the results on the edge pruning threshold is that the experimental outcomes demonstrate that FlowPrune is not particularly sensitive to its value. Among the tested values, $1 \\times 10^{-6}$ consistently performs well and offers a reasonable balance. The table below illustrates the average absolute error for the LLaMA model under different edge pruning thresholds after retaining half of the layers (16 layers).\n\n|1e-8|1e-7|1e-6|1e-5|\n|---|---|---|---|\n|0.0137|0.0099|0.0097|0.0103|\n\n**5. Generalizability.** To evaluate the generalizability of FlowPrune, we have selected two representative large models: the unimodal LLaMA and the multimodal LLaVA. We also assessed its effectiveness on smaller-scale models—BERT for NLP (Sec 4.2.1) and ViT for vision (Sec 4.2.2). In addition to measuring absolute error against the true attention flow (Fig 4), we considered that attention flow is often used to study relative ranking; thus, we introduced corresponding metrics for this purpose (Fig 5, 6). Together, these diverse models and evaluation criteria support the general applicability of our method. In response to your suggestion, we also included experiments on two widely used large models, Qwen (unimodal) and Qwen-VL (multimodal). The following table presents the absolute and relative errors in the high-flow regions, which are of the most interest to researchers, for the Qwen and Qwen-VL models, where we can find that the absolute/relative error is not large.\n   \n  | Region\\Model| qwen (absolute error) | qwen (relative error) | qwen-vl (absolute error) | qwen-vl (relative error) |\n| --- | --- | --- | --- | --- |\n[0.9-1.0] | 0.003 | 0.31% | 0.001 | 0.12% |\n[0.8-0.9] | 0.014 | 1.63% | 0.012 | 1.41% |\n[0.7-0.8] | 0.017 | 2.28% | 0.013 | 1.71% |\n\nIn addition to the aforementioned experiments, we also included an application scenario experiment on Question Answer Verification, which is also based on BERT. You can find the detailed introduction of the experiment in [22]. We compared the Top-K Token Retention with Attention Flow of FlowPrune and Attention Rollout, and organized the results in a manner similar to *Table 1* in Section 4.2.1. The results show that, compared to Attention Rollout, the FlowPrune method is much closer to the analysis results of Attention Flow, demonstrating FlowPrune outperforms Attention Rollout. The table below presents the results:\n\n|Max-n|FP Pretrained(Mean ± SE)|AR Pretrained(Mean ± SE)|FP Fine-tuned(Mean ± SE)|AR Fine-tuned(Mean ± SE)|\n|---|---|---|---|---|\n| n = 3 | 0.9805 ± 0.0111 | 0.6468 ± 0.0291 | 0.9610 ± 0.0154 | 0.8571 ± 0.0259 |\n| n = 5 | 0.9731 ± 0.0154 | 0.8448 ± 0.0214 | 0.9416 ± 0.0231 | 0.7631 ± 0.0221 |\n| n = 10 | 0.9701 ± 0.0172 | 0.6650 ± 0.0188 | 0.9342 ± 0.0260 | 0.7103 ± 0.0176 |\n\n**6. Input Attribution Techniques.** Your question is indeed very insightful. In fact, there has been some debate over the merits of attention-based versus gradient-based attribution methods about 5-6 years ago [A, B]. Wiegreffe and Pinter [B] argued that attention can be \"an explanation, not the explanation.\" This implies that for the same prediction, there may be multiple equally valid and reasonable explanations. We concur with this conclusion and believe that both types of methods have their own strengths.\n\nResearchers generally view these two types of methods as representing different categories of explanations. Attention-based methods focus on Intrinsic Interpretability, utilizing the inherent architectural characteristics of the model to provide explanations. They offer a built-in window into the model's internal \"focus\" [C]. On the other hand, gradient-based attribution methods focus on Post-Hoc Explanations, generating explanations by probing the model's behavior. This category is primarily dominated by saliency and attribution techniques.\n\nAlso, attention-based methods have the advantage of requiring only forward propagation, making them simpler to use and more intuitive [D]. In contrast, gradient-based methods necessitate backward propagation, which demands additional GPU computation and is not supported in some scenarios. For instance, many models undergo quantization after training, and quantized models do not support backward propagation. In such cases, attention-based methods are more convenient. Moreover, for attention flow[14], it has about 1000 citations in Google Scholar, suggesting it has considerable influence and impact in the community.\n\n[A] Attention is not Explanation\n\n[B] Attention is not not Explanation\n\n[C] The Role of Attention Mechanisms in Enhancing Transparency and Interpretability of Neural Network Models in Explainable AI\n\n[D] A Comprehensive Guide to Explainable AI (XAI) for Generative Models\n\n**7.  Compression Rate.** Thank you for pointing out this issue. In our paper, the compression rate is defined as the ratio of the number of **retained layers** to the total number of layers. Therefore, when the compression rate approaches 1, it implies that the number of retained layers is close to the total number of layers, meaning that more layers are preserved. This name is indeed counterintuitive. Following your suggestion, we will rename the compression rate to **retain rate** to align the terminology with its definition. We will correct this name in the revision.\n\n**8. L/R Layers.** In fact, in the context of our paper, it should be $L/R$ layers in line 218. The context means that retaining only $L/R$ layers will reduces the computational time, where $L$ denotes the total number of layers in the model, and $1/R$ represents the compression rate, or retain rate in the revision."}]

Author rebuttal:

Meta-review: This paper uses max-flow algorithms to analyse and improve the attention mechanism in transformers. The paper is found to be standing on solid grounds by the reviewers, and even the skeptic reviewer acknowledged that his concerns were responded well by the authors. I am wondering about the complexity of the algorithms and whether they are tangible for large networks. The authors should be more careful with their writing.
Decision: Accept (poster)