From spikes to heavy tails: unveiling the spectral evolution of neural networks

Vignesh Kothapalli, Tianyu Pang, Shenyang Deng, Zongmin Liu, Yaoqing Yang

Transactions on Machine Learning Research 2025

Summary: This paper uncovers the mechanism behind the emergence of heavy-tailed empirical spectral densities (ESDs). We show that heavy-tailed ESDs arise from the interaction between an ESD spike, caused by feature learning, and the bulk, resulting from iid random weight initialization. This theory reveals several surprising facts: the emergence of a heavy-tailed spectrum (1) does not require SGD noise during training, (2) does not need the model to be overparameterized or sufficiently interpolated, and (3) does not require more than one spike in the ESD.

Full paper  |  Code

Eigenspectrum analysis of neural networks without aspect ratio bias

Yuanzhe Hu, Kinshuk Goel, Vlad Killiakov, Yaoqing Yang

ICML 2025

Summary: Our group has published several papers on using spectral analysis of weight matrices to identify critical "under-trained" layers in neural networks. This insight led to the development of methods like Temperature Balancing and AlphaPruning. However, we recently uncovered a key oversight in our own work: when two weight matrices have different aspect ratios—that is, different m/n ratios for an m-by-n matrix—it is inaccurate to compare their eigenspectra directly. The empirical spectral density is biased by this shape difference, a factor we had overlooked. Correcting this bias leads to consistent improvements across methods like Temperature Balancing and AlphaPruning.

Full paper  |  Code

LIFT the veil for the truth: principal weights emerge after rank reduction for reasoning-focused supervised fine-tuning

Zihang Liu, Tianyu Pang, Oleg Balabanov, Chaoqun Yang, Tianjin Huang, Lu Yin, Yaoqing Yang, Shiwei Liu

ICML 2025

Summary: This paper investigates parameter-efficient fine-tuning (PEFT) of LLMs for reasoning tasks. While PEFT is more efficient than full fine-tuning, it often sacrifices reasoning accuracy. We show how to bridge this gap using a sparse fine-tuning method that delivers both efficiency and accuracy. The key insight is striking: the most important weights for fine-tuning are those with large magnitudes after applying a low-rank approximation to the weight matrices. We call this method “LIFT” because it reveals the useful weights by removing the noisy low-rank components.

Full paper  |  Code

Why LLM safety guardrails collapse after fine-tuning: a similarity analysis between alignment and fine-tuning datasets

Lei Hsiung, Tianyu Pang, Yung-Chen Tang, Linyue Song, Tsung-Yi Ho, Pin-Yu Chen, Yaoqing Yang

ICML 2025 Workshop on Data in Generative Models - The Bad, the Ugly, and the Greats

Summary: LLMs are often vulnerable to jailbreak attacks, especially after downstream fine-tuning. Existing mitigation strategies overlook a key factor: the role of the original safety-alignment data. This paper investigates how safety guardrails degrade by examining the representation similarity between upstream alignment datasets and downstream fine-tuning tasks. Our experiments show that high similarity between the two significantly weakens guardrails and increases jailbreak risk, while low similarity leads to more robust models.

Full paper  |  Blog

Mitigating memorization in language models

Mansi Sakarvadia, Aswathy Ajith, Arham Khan, Nathaniel Hudson, Caleb Geniesse, Kyle Chard, Yaoqing Yang, Ian Foster, Michael W. Mahoney

ICLR 2025

Summary: Language models can memorize training data, encoding it in their weights so that inference-time queries trigger verbatim regurgitation. This raises concerns when the data are private or sensitive. In this work, we compare seventeen memorization mitigation methods: three regularizer-based, three fine-tuning-based, and eleven machine unlearning methods, five of which are novel contributions. We also introduce TinyMem, a suite of small, efficient LMs for fast development and evaluation of mitigation strategies. Our experiments reveal that regularizer-based methods are slow and ineffective, fine-tuning methods reduce memorization but are prohibitively expensive when accuracy matters, and unlearning methods are both faster and more effective, enabling precise removal of memorized content before inference. Notably, our proposed unlearning method, BalancedSubnet, outperforms all others in eliminating memorized information while preserving task performance.

Full paper  |  Code  |  Blog

LossLens: diagnostics for machine learning through loss landscape visual analytics

Tiankai Xie, Jiaqing Chen, Yaoqing Yang, Caleb Geniesse, Ge Shi, Ajinkya Chaudhari, John Kevin Cava, Michael W. Mahoney, Talita Perciano, Gunther H. Weber, Ross Maciejewski

IEEE Computer Graphics & Applications 2024

Summary: Modern machine learning relies on optimizing a neural network’s parameters via a loss function to learn complex features. Examining this loss function with respect to the network’s parameters—the loss landscape—can reveal key insights into the architecture and learning process. While local landscape structures near individual solutions are well-studied, the global structure, with its many local minima, remains challenging to understand and visualize. To address this, we introduce LossLens, a visual analytics framework for exploring loss landscapes across multiple scales. LossLens integrates local and global metrics into a unified visual representation, making it easy for practitioners to conduct model diagnostics. We showcase its utility through two case studies: visualizing the role of residual connections in ResNet-20, and analyzing how physical parameters affect a physics-informed neural network solving a convection PDE problem.

Full paper  |  Code  |  Video

Model balancing helps low-data training and fine-tuning

Zihang Liu*, Yuanzhe Hu*, Tianyu Pang, Yefan Zhou, Pu Ren, Yaoqing Yang

EMNLP 2024

Summary: As shown in our previous work Temperature Balancing, different layers of an LLM can be highly imbalanced—some layers are much better trained than others. Here, we reveal that this imbalance is even more pronounced at the start of training, especially during fine-tuning with limited downstream data. This paper is also the first in which we apply layer-balancing techniques to scientific machine learning tasks.

Full paper  |  Code

AlphaPruning: using heavy-tailed self regularization theory for improved layer-wise pruning of large language models

Haiquan Lu*, Yefan Zhou*, Shiwei Liu, Zhangyang Wang, Michael W. Mahoney, Yaoqing Yang

NeurIPS 2024

Summary: Most deep neural networks have complex multilayer structures, often viewed as a barrier to transparency. Our research uncovers a key insight: these layers are not equally well-trained. This holds true even for LLaMA-scale LLMs. Building on this, we show that existing LLM pruning methods often fall short. Since layers differ in how well they are trained, applying the same pruning ratio across them is not ideal. This paper echoes “Temperature Balancing,” another work from our group, but here we scale HT-SR-based training methods to LLaMA-scale models for the first time. That means 70 billion parameters! Check out the paper—the experimental results are impressive.

Full paper  |  Code

Sharpness-diversity tradeoff: improving flat ensembles with SharpBalance

Haiquan Lu*, Xiaotian Liu*, Yefan Zhou*, Qunli Li*, Kurt Keutzer, Michael W. Mahoney, Yujun Yan, Huanrui Yang, Yaoqing Yang

NeurIPS 2024

Summary: This paper studies ensemble learning and reveals an intriguing tradeoff between two key factors that affect ensemble performance: the sharpness of the local minima where ensemble members converge and the prediction diversity among those members. Reducing sharpness in the loss landscape often comes at the cost of diversity, which can limit the ensemble's overall gain over its individual members. To address this, we introduce a new ensemble method called “SharpBalance” that balances sharpness and diversity. We also provide theoretical insights by analyzing moments of random matrices to support our empirical observations.

Full paper  |  Code

MD tree: a model-diagnostic tree grown on loss landscape

Yefan Zhou*, Jianlong Chen*, Qinxue Cao, Konstantin Schürholt, Yaoqing Yang

ICML 2024

Summary: Is a test accuracy of 69.5% good or bad? How can you decide the best way to improve your model? This paper tackles these questions by formally identifying the most “critical” failure source of an underperforming model. We use loss landscape metrics to assess the model and map them to one of three failure sources: bad data, bad model, or bad optimization. Our main technical contribution is a simple tree-based model that enables this classification and significantly outperforms traditional methods that rely solely on validation error.

Full paper  |  Code

Temperature balancing, layer-wise weight analysis, and neural network training

Yefan Zhou*, Tianyu Pang*, Keqin Liu, Charles H. Martin, Michael W. Mahoney, Yaoqing Yang

NeurIPS 2023

Summary: Most deep neural networks have complex multilayer structures, often seen as a barrier to transparency. In our research, we reveal a significant insight: these layers are not uniformly well-trained. By identifying and addressing underperforming layers, we enhance the overall network quality. Our approach introduces a "model diagnostic" tool for improving training. We demonstrate its effectiveness across various benchmarks, datasets, and network architectures, outperforming more than five existing methods, all rooted in our ability to dissect and diagnose network imbalances.

Full paper  |  Code

When are ensembles really effective?

Ryan Theisen, Hyunsuk Kim, Yaoqing Yang, Liam Hodgkinson, Michael W. Mahoney

NeurIPS 2023

Summary: This study examines when ensembles are "really" effective in improving the test accuracy of learning models. Our theoretical analysis establishes that ensembling improves test accuracy when the "disagreement" is high compared to the average error rate of individual learners. We establish this conclusion based on a condition known as "competence," which helps eliminate abnormal cases that often restrict conventional analysis on ensembling. Empirical findings validate the theory and highlight the more significant benefit of ensembling in non-interpolating models, such as tree-based methods, compared to interpolating models.

Full paper 

Evaluating natural language processing models with generalization metrics that do not need access to any training or testing data

Yaoqing Yang, Ryan Theisen, Liam Hodgkinson, Joseph E. Gonzalez, Kannan Ramchandran, Charles H. Martin, Michael W. Mahoney

KDD 2023

Summary: We provide the first large-scale correlational studies on the generalization measures for natural language processing models. This paper focuses on the measures derived from the heavy-tail self regularization (HT-SR) theory, which does not need access to training or testing data to calculate. Also, we show that these measures can perform uniformly better than existing norm-based measures if we aim to predict test-time performance instead of the "generalization gap", which is the difference between training and test accuracies. We use the WeightWatcher toolbox to analyze the HT-SR measures.

Full paper  |  Code  |  Video

A three-regime model of network pruning

Yefan Zhou, Yaoqing Yang, Arin Chang, Michael W. Mahoney

ICML 2023

Summary: Recent research has emphasized the intricate relationship between training hyperparameters and the ability to prune machine learning models. However, accurately predicting how adjusting a specific hyperparameter impacts pruning remains challenging. To address this gap, a phenomenological model based on the statistical mechanics of learning is introduced, using "temperature-like" and "load-like" parameters to represent the influence of hyperparameters on pruning performance. The study identifies a transition phenomenon, where the effect of increasing the temperature-like parameter depends on the value of the load-like parameter, leading to different pruning outcomes. The findings are then applied to three practical scenarios, including optimizing hyperparameters for improved pruning and selecting the most suitable model for pruning.

Full paper  |  Code

Two sides of the same coin: Heterophily and oversmoothing in graph convolutional neural networks

Yujun Yan, Milad Hashemi, Kevin Swersky, Yaoqing Yang, Danai Koutra

ICDM 2022

Summary: Graph convolutional neural networks may perform worse when we increase the number of layers (oversmoothing problem) and when we feed in heterophilous graphs (heterophily problem). In this work, we show it theoretically and empirically that these two seemingly unrelated problems are closely related.

Full paper  |  Code

Neurotoxin: durable backdoors in federated learning

Zhengming Zhang*, Ashwinee Panda*, Linyue Song, Yaoqing Yang, Michael W. Mahoney, Prateek Mittal, Kannan Ramchandran, Joseph E. Gonzalez

ICML 2022

Summary: We propose Neurotoxin, a simple one-line modification to existing backdoor attacks in federated learning. Our attack can double the durability of state of the art backdoors.

Full paper  |  Code

Taxonomizing local versus global structure in neural network loss landscapes

Yaoqing Yang, Liam Hodgkinson, Ryan Theisen, Joe Zou, Joseph E. Gonzalez, Kannan Ramchandran, Michael W. Mahoney

NeurIPS 2021

Summary: This paper experimentally demonstrates the long-standing conjecture that "local properties" of a loss landscape cannot dictate generalization. The study taxonomizes learning problems into "phases" by analyzing various generalization metrics obtained from the loss landscapes of neural networks, and it provides a formal way to divide and conquer typical failure modes of learning in the different phases.

Full paper  |  Code  |  Video

Improving semi-supervised federated learning by reducing the gradient diversity of models

Zhengming Zhang*, Yaoqing Yang*, Zhewei Yao*, Yujun Yan, Joseph E. Gonzalez, Kannan Ramchandran, Michael W. Mahoney

IEEE BigData 2021

Summary: Cell phone users who participate in federated learning often do not have the time to provide labels to their private data, making semi-supervised learning a practical alternative. This paper shows that the large dissimilarity between model gradients from different users could arise from the semi-labeled data and become an obstacle to semi-supervised federated learning.

Full paper  |  Code

Boundary thickness and robustness in learning models

Yaoqing Yang, Rajiv Khanna, Yaodong Yu, Amir Gholami, Kurt Keutzer, Joseph E. Gonzalez, Kannan Ramchandran, Michael W. Mahoney

NeurIPS 2020

Summary: This paper introduces the notion of "boundary thickness" and shows that thin decision boundaries lead to overfitting (e.g., measured by the robust generalization gap between training and testing) and lower robustness. Also, welcome to check Dominic's thesis and see how we use boundary thickness to reveal "backdoors" hidden in a neural network.

Full paper  |  Code

Foldingnet: Point cloud auto-encoder via deep grid deformation

Yaoqing Yang, Chen Feng, Yiru Shen, Dong Tian

CVPR 2018

Summary: In this work, a novel auto-encoder is proposed to address the challenge of unsupervised learning on point clouds. A novel folding-based decoder is used to deform a canonical 2D grid onto a point cloud's underlying 3D object surface. The proposed decoder structure is proved, in theory, to be a generic architecture that can reconstruct an arbitrary point cloud from a 2D grid.

Paper  |  Code  |  Video

Mining point cloud local structures by kernel correlation and graph pooling

Yiru Shen*, Chen Feng*, Yaoqing Yang, Dong Tian

CVPR 2018

Summary: Existing ML models on point clouds do not take full advantage of a point’s local neighborhood that contains fine-grained structural information. In this paper, we present novel operations to exploit local structures in a point cloud.

Paper  |  Code

Serverless straggler mitigation using local error-correcting codes

Vipul Gupta*, Dominic Carrano*, Yaoqing Yang, Vaishaal Shankar, Thomas Courtade, Kannan Ramchandran

ICDCS 2020

Best Paper Finalists

Summary: Inexpensive cloud services, such as serverless computing, are often vulnerable to straggling nodes that increase end-to-end latency. We propose and implement simple yet principled coding approaches for straggler mitigation.

Full paper  |  Code

Coded elastic computing

Yaoqing Yang, Matteo Interlandi, Pulkit Grover, Soummya Kar, Saeed Amizadeh, Markus Weimer

ISIT 2019

Summary: Cloud providers have recently introduced new offerings whereby spare computing resources are accessible at discounts compared to on-demand computing. Exploiting such an opportunity is challenging since such resources are accessed with low priority and can elastically leave (through preemption) and join the computation at any time. This paper designs a new technique called coded elastic computing, enabling distributed computations over these elastic resources.

Full Paper  |  Code

Coded iterative computing using substitute decoding

Yaoqing Yang, Malhar Chaudhari, Pulkit Grover, Soummya Kar

ISIT 2018

Summary: Applying conventional linear codes to large-scale matrix operations can make sparse matrices dense, and codes with low-density generator matrices (LDGM) are often preferred. In this paper, we show a novel way of using LDGM codes called "substitute decoding". Applications of this new coding scheme include power iterations, truncated singular value decompositions, and gradient descent in the distributed setting.

Conference Paper  |  Full Paper

Coded distributed computing for inverse problems

Yaoqing Yang, Pulkit Grover, Soummya Kar

NeurIPS 2017

Summary: In this paper, we utilize the emerging idea of "coded computation" to design a novel technique for solving linear inverse problems under specific iterative methods in a parallelized implementation affected by stragglers. The applications studied in this paper include personalized PageRank and sampling on graphs.

Paper  |  Arxiv Version

Computing linear transformations with unreliable components

Yaoqing Yang, Pulkit Grover, Soummya Kar

Transactions on Information Theory 2017

Summary: The work provides the first coding strategies that provably require fewer gates in scaling sense than replication for computing finite-field linear transforms with all computational nodes being error-prone. The main insight is that allowing all nodes to be error-prone necessitates repeated error suppression through the embedding of decoders inside the computation, resulting in a "coded computation" setup.

Full paper  |  Code

Rate distortion for lossy in-network linear function computation and consensus: Distortion accumulation and sequential reverse water-filling

Yaoqing Yang, Pulkit Grover, Soummya Kar

Transactions on Information Theory 2017

Summary: The work provides fundamental limits as well as achievable strategies on "distortion accumulation" in distributed linear computing problems.  By successfully characterizing the overall distortion-rate function with accumulated distortion in a high-rate regime, we tighten earlier cut-set bounds by a factor that can be arbitrarily large even in simple line networks.

Full paper