Persistent Topological Features in Large Language Models

We thank the reviewer for their useful feedback. They recognize the novel approach of using zigzag persistence in the context of interpretability of NNs and the clarity of presentation. Additionally, they raise a few points which we address element-wise below:

Claim of “Identification of Phases of Prompt processing” is too broad - prompts can’t be truly represented with only one token (as it is done in the paper).

While we agree with the referee that an analysis of whole sentences would allow for finer-grained information about phases, we chose to simplify the analysis computationally by adopting the common practice in studies of the geometry of representations of looking at the last token only.

Why was used sliding window of 5, why no other sizes were reported?

Figure 15 shows the analysis for different sliding windows, which we use to zoom in on a phase in which accuracy drops significantly. In other ranges, the size of the window would not change results significantly, if kept small enough.

Is there any difference (in terms of identified 4 phases) for encoder-only models (in the paper only decoder models are covered)?

We believe that a difference might indeed be seen for encoder-only models. For those models, considering the last token as a proxy for the whole sentence would not be correct. Along these lines, previous studies of protein language models (see e.g. Valeriani et al. 2023) have considered averaging over tokens. We deserve these studies to future work.

Only one task of layer pruning was explored. Some other tasks, especially those focusing on topology of individual texts, would be much appreciated.

We thank the referee for suggesting exploring different downstream tasks. Indeed previous work has considered using geometric quantities (e.g. intrinsic dimension) to identify human text from AI-generated text (Tulchinski et al 2024). An analysis of individual sentences would imply using our algorithm on the tokens of a prompt, i.e. without reducing the whole prompt to a single prompt. Such analysis would be interesting, though quite different from the present work since the manifolds in the two cases would represent rather different dynamics.

Have you performed your main experiments only with 0-, and 1-dimensional topological features?

We have performed our calculations for up to 3-dimensional holes. However, as mentioned in section 4.2, 0-,2- and 3-dimensional holes have relatively low number counts, making it hard to draw solid conclusions from them. A brief discussion on which homology dimension contains relevant information is included in Appendix C.

From what dimensionality topology of kNN-complexes becomes trivial?

We find non-zero counts up to 3-dimensional holes. This is specific to kNN complexes which are typically more connected as compared to other complex such as e.g. Vietoris-Rips complexes.

Could you please provide for Figures 4 plot of average length of intervals in (classic) persistence diagrams for each of the layers for comparison?

We apologize with the referee, but we did not understand this question. Are they asking to perform persistent homology at each layer for comparison?

We thank the reviewer for their insightful feedback. They recognize the novelty of our approach to analyzing LLMs and its importance in enhancing understanding of these models. They also acknowledge a proper review of existing literature. As for feedback on points to improve, we reply element-wise below:

the shape of the plots in Figure 3 changes significantly with alpha

The variability in alpha is a feature of the method since it allows to explore different regimes (i.e. short-lived vs long-lived features). This is explained in the paragraphs where phases are described. As a side note, we show that this behavior is consistent across models in App. D.

the paper only reports Birth's relative frequency plots for variations of Llama (do not consider other LLMs). [...] Also, in Appendix D1, I expected to see the same plots for different datasets and models.

We have computed these plots, but did not add them to the manuscript. They can be accessed here: https://anonymous.4open.science/r/conferenceProject-019A/src/plots/Rebuttal-plots.md

[...] it is unclear if this is the case --- e.g., the number of tuples of persistence = 5 seems similar for l=5 and l=25 (or persistence = 10 for l=5 and l=20).

We agree that this statement is hard to confirm by eye in Fig. 2. We are open to rephrase it as: “Figure 2 shows that features born after the first half of the model’s depth have a higher tendency to be long-lived with respect to features born earlier on.”

[..] This is less than found in other papers on layer pruning like "ShortGPT: Layers in Large Language Models are More Redundant Than You Expect".

We have run additional benchmarks, results included here: https://anonymous.4open.science/r/conferenceProject-019A/src/plots/Rebuttal-plots.md.

how does it decide which simplex dies when two simplices get merged in the subsequent layer? Does it matter?

The topological feature that dies is uniquely determined by the decomposition of the zigzag module provided in Equations 11-12, removing any ambiguity in its identification.

[..] the specific form of Eq. (4) is not properly motivated. In addition, why is "smoothness as a function of layers" important for the application at hand?

PIs are usually linked to a choice of kernel. In our case, the PD is already amenable to a density grid, the value of density being feature counts. Intersection layers are only necessary for defining the filtration (e.g. one could use the union equivalently). This is desirable since we are interested in studying statistically the LLM layers. The smoothness is related to smoothing formigram (Prop. C.2) in https://arxiv.org/abs/1712.04064 where they smooth barcodes defined on X (all the layers including intersection our case) to barcodes on S(X) (just the LLM layers)

Overall, the results regarding layer pruning are not impressive. Indeed, since the competitive methods seem simpler (angular distance and bi-score), I am not sure if one would prefer this approach.

See a reply to a similar feedback in response to reviewer zXjm.

Replies to minor issues:

the choice for (Gromov et al, 2024) and (Men at al, 2024) is not motivated.

These methods were chosen as it is easy to compare with ours, since the algorithms require as an input the number of blocks to be removed, which is what our criterion based on inter-layer persistence outputs.

Based on Fig. 1, the range of (b,d) should be {0, \dots, 2(N - 1)}.

Correct, this is a typo caused by Python notation. In fact, there are 2N-1 snapshots.

Is $l_i =i$ ? in all cases? If so, why use $l_i$?

The notation $\ell_i$ was used to be suggestive of the fact that the indices refer to layers.

Replies to “Questions for Authors”:

1. As noted at the beginning of Sec. 4.2, experiments are performed for 0,1,2,3-dim holes, however only 1-dim holes have large number counts so that a statistical study is stable.

2. Yes, the method can be used to analyze models with variable embedding dimension (e.g. CNNs), this is mostly because our method is distance-based and we can expect neighborhoods to shift smoothly in CNNs (see eg https://arxiv.org/pdf/2007.03506)

3. Even though our method allows us to cut non-consecutive layers, our results are in line with the compared methods cutting consecutive layers. If tested on more specialized prompts, our algorithm might cut non-consecutive layers for a low enough threshold (see e.g. Java code in Fig. 9).

4. The key point of our work is to build an interpretable framework that allows studying the internal reps of NNs as a whole system, rather than collecting summaries of snapshots and then combining them a posteriori. This is important in for interpreting different phases of prompt processing, as argued in our work.

5. The computational cost of our method is higher than other methods, as detailed in the reply to 4Kx7. While it is an important worry for scaling up our analysis, for this specific task, it is not an issue given pruning is performed only once.

Reply to zXjm

We thank the reviewer for their useful feedback. They highlight the novelty of using zigzag persistence from topological data analysis to understand feature evolution in LLMs, appreciating the accessible presentation. They raise a few concerns, which we address below:

My biggest question is why the idea of Zigzag is important/interesting to the ML community. I'm not an expert in TDA, but this looks like a niche concept to me

Within the TDA community, there exist several works addressing time-varying point clouds. The reason for the relatively slower progress in this field has been mainly computational: in a similar way as in multi-persistence, the complexity grows rapidly when adding filtration parameters, and real-world applications are harder to implement. Nevertheless, recent advances in fast algorithms have allowed these methods to be applied more widely. In the specific case of Zigzag persistence, the fast zigzag algorithm of [1] allowed us to run the pipeline on large and high-dimensional datasets (see reply to 4Kx7 for computational complexity). Consequently, these algorithms have become feasible to apply in the framework of interpreting dynamical changes in internal representations of neural networks.

In short, we would argue that recent computational advances in analysing time-varying point clouds with TDA approaches have greatly increased relevance for the broader ML community.

The improvement in the accuracy of the authors' method compared to the pruning method looks marginal

As the review remarks earlier on, as the work’s main objective is to connect a relatively underexploited mathematical algorithm (zigzag) to LLMs interpretability, the scope of the layer pruning analysis is to show that our method can be used on a downstream task with performances on par with state of the art methods.

Is it possible to report SE so the interpretation of the results can be more robust?

There is relatively little variability in the benchmarks used for layer pruning, as it comes from boostrapping. Consequently, our computation reveals a standard deviation of ~1e-4. We note that whenever the accuracy score is the same for our work and the cited references, the results are exactly equal, i.e. the methods suggest pruning the same layers.

Also the computational cost of the authors' method seems expensive.

We refer to our reply to 4Kx7 for computational costs of the algorithm. As for the specific comparison to the layer pruning methods cited in our work, our algorithm is relatively slower, but it contains strictly more information (beyond layer pruning), as it analyzes all layers at the same time, rather than considering pairwise comparisons among layers.

[1] Dey, T. K. and Hou, T. Fast Computation of Zigzag Persistence. In Chechik, S., Navarro, G., Rotenberg, E., and Herman, G. (eds.), 30th Annual European Symposium on Algorithms (ESA 2022), volume 244 of Leibniz International Proceedings in Informatics

Reply to 4Kx7

We thank reviewer 4Kx7 for their useful feedback. While recognizing the novelty of our approach in combining zigzag persistence for LLM interpretability, the reviewer correctly points out that our work should have included a broader discussion of similar TDA approaches. We agree with this comment and recognize it would improve the effectiveness of our work. Here we reply to direct points raised by the referee in this regard, and we would be open to including part of this discussion in an edited version of the manuscript

Dynamic Persistence [1], related to zigzag persistence, should be added in references and compared with.

The reference is already cited in the main text as “Kim and Memoli, 2021.”

The integration of the proposed method into filtration, and the tracking of the evolution of internal representations both are desired, which deserve to be appended for further validating the proposed zigzag persistence based method.

Given the large number of points considered in this analysis, we adopt a statistical approach to study the effect of layer transformations on internal representations, where the identity of individual points is not essential. The idea of tracking the zigzag representatives using [3] could certainly be investigated. Testing this pipeline on an example in a more controlled setting would be useful in validating the proposed zigzag framework and it would allow in principle to track more closely the evolution of features across layers. We deserve these studies for future work.

The comparison to related TDA approach is insufficient, How does the proposed method compare to dynamic persistence?

Here are a few points of comparison:

1. We choose kNN filtration (whereas it is Rips in [1]) since it is more suitable to high dimensional data, especially for LLM reps. The stability of kNN filtrations is discussed in [4] in the context of persistent homology, we recognize that we missed this reference in the submitted manuscript, we would be open to add it in an edited version.
2. [1] varies both a time and a scale parameter whereas we fix the scale. Their main summary statistic (the rank invariant) involves calculating a 6-dimensional summary statistic (4 across layers and 2 across scale) which can make analysis hard for both computational reasons as well as finding the informative summary statistic. Nevertheless, varying also the scale parameter is definitely worth investigating in this context, and the techniques in [1] would be a starting point for implementing it.
3. The work [5] from the same authors is more related to our work since the maximal group diagram and the persistence clustergram described in the above paper are "annotated" (with the representative topological features) barcodes. In this work, they fix a scale similar to our case.

The runtime efficiency evaluation for the proposed method is missing, which is quite desirable to know the efficiency performance of the proposed method.

A broad assessment of computational efficiency is present in the current version (Appendix B). We agree this is an important point to address, thus here are more details on it: the theoretical complexity of FastZigZag is O($m^{\omega}$) with $\omega < 2.37286$ for a sequence of $m$ deletion and additions that is an advancement from previous algorithms with a cubic complexity, which allowed us to do our experiments with point clouds of 10k points over ~30 layers. More details on the algorithm are described in the paper by Dey and Hou [2]. As for the knn graph, its creation with a greedy algorithm is O($n^2$) with n being the data points. In total, we have a O($n^2 * N_{layers}$) + O($m^{\omega}$). We would be open to adding this discussion in an edited version.

[1] Kim, Woojin, and Facundo Mémoli. "Spatiotemporal persistent homology for dynamic metric spaces." Discrete & Computational Geometry 66 (2021): 831-875.

[2] Dey, T. K. and Hou, T. Fast Computation of Zigzag Persistence. In Chechik, S., Navarro, G., Rotenberg, E., and Herman, G. (eds.), 30th Annual European Symposium on Algorithms (ESA 2022), volume 244 of Leibniz International Proceedings in Informatics

[3] Dey, Tamal K., Tao Hou, and Dmitriy Morozov. "A Fast Algorithm for Computing Zigzag Representatives." Proceedings of the 2025 Annual ACM-SIAM Symposium on Discrete Algorithms (SODA). Society for Industrial and Applied Mathematics, 2025.

[4] Le, Minh Quang, and Dane Taylor. "Persistent homology with k-nearest-neighbor filtrations reveals topological convergence of PageRank." Foundations of Data Science 7.2 (2025): 536-567.

[5] Kim, Woojin, and Facundo Mémoli. "Extracting persistent clusters in dynamic data via möbius inversion." Discrete & Computational Geometry 71.4 (2024): 1276-1342.