Loss in the Crowd: Hidden Breakthroughs in Language Model Training

• There are various hyperparameters introduced here (number of Hessian basis vectors to obtain at each iteration, number of training steps between iterations, choice of clustering method and whatever hyperparameters are associated with that clustering method, eg number of clusters) I would appreciate more discussion of how the hyperparameters settings were chosen for the experiments in the paper and robustness of the results to these settings.
• The process of associating an interpretation to data points corresponding to each cluster, per basis vector, per output token, seems to be labor-intensive and possibly error-prone.
• For natural language, POLCA is not compared with any baselines. (For the addition task, there is a comparison with clustering solely on per-token loss, which is good.) Another natural baseline to compare against is with randomly sampled basis vectors, instead of eigenvectors of the loss Hessian.
• Although it is interesting that the POLCA trajectory clusters correspond to human-interpretable skills (at least for the examples given), it is not clear how this interpretation is useful in practice.
  • The paper says, "POLCA analysis could also be combined with model editing methods to better remove abilities learned by the model." I would be interested to see more elaboration on how this could be done + some preliminary results in this direction. For example: does ablating a POLCA basis vector destroy model performance on specifically the skills associated to corresponding clusters?
• Causal experiments like the example in the above point ^ would also help to increase confidence in the current results, which are all observational.

1. While POLCA is shown to perform well, the paper lacks a robust comparison with existing methods for interpretability or skill recovery in LLM training, such as traditional reinforcement learning or trajectory clustering methods. A comparative study could have strengthened the claims about POLCA's uniqueness and efficiency.
2. While the method is demonstrated on a smaller transformer and Wikipedia dataset, the paper lacks discussion on how POLCA would scale to larger LLMs. Given the computational overhead of Hessian-based calculations, discussing the feasibility of POLCA in large-scale training scenarios would be valuable.
3. Although POLCA is a promising approach, the paper could benefit from a more comprehensive discussion of its limitations. For example, the reliance on Hessian eigenvectors might introduce noise or dependencies in high-dimensional parameter spaces. Addressing these potential issues would strengthen the claims.

Despite the work is quite solid in its presented setups, I am concerned about the scalability of the approach. This matters because the practical utility of training dynamics research are commonly found in large-scale models, like Xia et al. 2023 the authors cited.

• The experiments are performed on quite small transformer language models (a 2-layer transformer model).

• If I understand correctly, finding the basis (in Sec. 3.1.1) requires computing the Hessian matrix of size $d \times d$, which is infeasible even for a relatively "small" LLM with 1B parameter, because the Hessian has the same cost as storing 1 billion 1B models. Is it the main reason why the authors only performed experiments on small transformers?

• If so, I suggest the authors to discuss the scalability issue in the Limitations section, and consider viable workarounds like approximations.

• I think the computational efficiency of the entire method is bottlenecked by the computation cost of Hessian matrix in Sec. 3.1.1, as loss decomposition along the basis (Sec. 3.1.2) & clustering are quite computationally efficient. Could you provide more justification about the choice of using the Hessian in Sec. 3.1.1?

For me, the scalability does not quite block the paper from acceptance, but this is certainly a weakness that should be discussed.

• Overall, I think the presentation of the paper can be improved:

  • Section 4 is pretty opaque - it took me several tries to understand the experimental setting described by the authors, mostly due to not immediately understanding what the POLCA vectors being clustered represent and how many skills the model was being examined for. I believe having a separate algorithm table describing the entire process, as well as explaining more succinctly what the precise skills that are being evaluated here are would help.

  • I'm having trouble understanding the point that the Figures in Section 4 aim to convey. These figures contain the main experimental result of this Section, and according to the authors show that clustering using the POLCA vectors demonstrates the skills being learned by the model (Figure 2 in particular). However, this is not immediately clear to me, mainly because a) the authors label the bars in the barplot with some of the skills of the model, but don't actually explain how this label is derived, and b) in Figures 2c and 2i a significant portion of the data points are labeled as outliers by the clustering algorithm. I would be grateful if the authors could elaborate on these points.

• I also find that the experimental portion of the paper is not convincing enough:

  • In Section 4 (experiment on synthetic data), as mentioned above only clusters for select POLCA trajectory vectors are presented, with a significant portion being labeled as outliers. The additional settings presented in the appendix for the same experiment show that in some cases, more vectors will end up being labeled as outliers than being part of actual clusters. This undermines the usefulness of the method in my opinion. I believe that some statistics on how often clustering is "succesful" (in the sense that the algorithm returns useful clusters) would elucidate this point (such as what is the average percentage of points being marked as outliers when clustering).

  • In Section 5, the experiment concluded is mostly qualitative, and focuses only on two tokens. While I understand that analysing every single output token is not practical, the analysis performed here is quite basic, and not particularly convincing. More specifically, it is not clear how distinct these skills are (e.g. some contexts of the token "and" have overlap between the "List" and "Token after comma" skills). It is not explicitly stated how many elements are being used to decide which skill corresponds to each cluster (the authors quantify it as a proportion of the elements in the cluster, but the fact that the clustering can have a large amount of outliers as seen in the previous sections makes the absolute amount unclear). I would be grateful if the authors could clarify this.

Overall, while I believe the method is interesting and shows promise, the results presented in this paper don't appear as convincing.

1. Insufficient Discussion of Related Work and Ambiguous Terminology

   The paper does not sufficiently position POLCA within the broader context of related work, missing key opportunities to discuss previous methods in phase transition analysis, loss dynamics, and interpretability—even though interpretability is a central claimed benefit. This gap makes it difficult to assess the novelty and significance of POLCA. Although the background section lists some relevant points, these feel more like brief justifications than a thoughtful review of prior literature. Consequently, the section reads more like a positioning statement than a critical examination of foundational work.

   Additionally, the paper introduces numerous terms without adequate definitions, which creates challenges in readability and comprehension. Terms such as “breakthroughs,” “shifts,” “structures,” “breakthrough elisions,” “conceptually meaningful,” “loss Hessian,” “base,” “basis,” and “skill” are used without prior explanation, making it hard for readers to fully grasp POLCA’s approach and its connections to existing research. Providing clear definitions for these terms would greatly enhance the paper’s clarity and accessibility.

2. Unnecessary Complexity in Methodology

   POLCA’s reliance on second-order Hessian approximations appears overly complex, particularly when simpler methods could likely achieve similar insights. For example, if the main goal is to cluster training trajectories for individual samples, established methods like saliency maps, activation pattern analysis, or influence functions might provide a more straightforward and interpretable alternative. While influence functions also involve Hessian computations, they are a widely recognized technique and may offer a more practical and accessible alternative to POLCA’s decomposition approach.

3. Inconsistent and Confusing Notation in Algorithm Formulation

   The mathematical notation in POLCA’s formulation is inconsistent and unclear in several places, which hampers the paper’s readability. For instance, the symbol $\phi$ in line 127 is ambiguous—if it represents an empty set, it is unclear how an inverse of an empty set is computed in line 128. Notational inconsistencies, such as $\mathcal{L}$ in line 130 versus $L$ in line 148, are also left unexplained, making the flow of the loss formula challenging to follow. Additionally, symbols like $\perp$ (presumably indicating orthogonality) are not defined. If $\perp$ is meant to denote orthogonality of $\Pi$ to the matrix $\mathcal{H}$, this should be clarified, along with citing the method (e.g., Lanczos method) used to find eigenvectors, as referenced in line 131.

4. Insufficient Experimental Validation

   The experimental setup appears limited and lacks rigor, especially in the natural language setting, which raises questions about the practical utility of POLCA. For instance, the choice of a two-layer transformer model with an embedding size of 512 is not justified, nor is there any variation in model configurations. Why were these specific parameters chosen, and would POLCA yield similar results with different configurations, such as deeper layers or alternate embedding sizes? The limited variation in hyperparameters weakens confidence in the generalizability of the findings. Additionally, while POLCA seems versatile, it is tested on only language modeling task; applying it to more diverse tasks would have better demonstrated its applicability. Finally, the absence of baseline （see 2）comparisons further undermines confidence in the purported advantages of POLCA.

5. Unclear Practical Relevance

   The practical relevance of detecting these “hidden breakthroughs” is not well-argued, and without concrete examples, it remains unclear how this information would enhance training efficiency, model interpretability, or overall performance. Strengthening the argument for the real-world utility of POLCA would make its contributions to applied machine learning more compelling.