More Like this

1. Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis

   Yang, Hongru; Kailkhura, Bhavya; Wang, Zhangyang; Liang, Yingbin (December 2024, Conference on Neural Information Processing Systems (NeurIPS 2024))

   Understanding the training dynamics of transformers is important to explain the impressive capabilities behind large language models. In this work, we study the dynamics of training a shallow transformer on a task of recognizing co-occurrence of two designated words. In the literature of studying training dynamics of transformers, several simplifications are commonly adopted such as weight reparameterization, attention linearization, special initialization, and lazy regime. In contrast, we analyze the gradient flow dynamics of simultaneously training three attention matrices and a linear MLP layer from random initialization, and provide a framework of analyzing such dynamics via a coupled dynamical system. We establish near minimum loss and characterize the attention model after training. We discover that gradient flow serves as an inherent mechanism that naturally divide the training process into two phases. In Phase 1, the linear MLP quickly aligns with the two target signals for correct classification, whereas the softmax attention remains almost unchanged. In Phase 2, the attention matrices and the MLP evolve jointly to enlarge the classification margin and reduce the loss to a near minimum value. Technically, we prove a novel property of the gradient flow, termed \textit{automatic balancing of gradients}, which enables the loss values of different samples to decrease almost at the same rate and further facilitates the proof of near minimum training loss. We also conduct experiments to verify our theoretical results. 

   more » « less
2. Compositionally restricted attention-based network for materials property predictions

   https://doi.org/10.1038/s41524-021-00545-1  

   Wang, Anthony Yu-Tung; Kauwe, Steven K.; Murdock, Ryan J.; Sparks, Taylor D. (December 2021, npj Computational Materials)

   null (Ed.)

   Abstract In this paper, we demonstrate an application of the Transformer self-attention mechanism in the context of materials science. Our network, the Compositionally Restricted Attention-Based network (), explores the area of structure-agnostic materials property predictions when only a chemical formula is provided. Our results show that ’s performance matches or exceeds current best-practice methods on nearly all of 28 total benchmark datasets. We also demonstrate how ’s architecture lends itself towards model interpretability by showing different visualization approaches that are made possible by its design. We feel confident that and its attention-based framework will be of keen interest to future materials informatics researchers. 

   more » « less
3. On the connection between mpnn and graph transformer.

   Cai, C; Hy, TS; Yu, R; Wang, Y (July 2023, MLResearch Press)

   Graph Transformer (GT) recently has emerged as a new paradigm of graph learning algorithms, outperforming the previously popular Message Passing Neural Network (MPNN) on multiple benchmarks. Previous work shows that with proper position embedding, GT can approximate MPNN arbitrarily well, implying that GT is at least as powerful as MPNN. In this paper, we study the inverse connection and show that MPNN with virtual node (VN), a commonly used heuristic with little theoretical understanding, is powerful enough to arbitrarily approximate the self-attention layer of GT. In particular, we first show that if we consider one type of linear transformer, the so-called Performer/Linear Transformer, then MPNN+ VN with only depth and width can approximate a self-attention layer in Performer/Linear Transformer. Next, via a connection between MPNN+ VN and DeepSets, we prove the MPNN+ VN with width and depth can approximate the self-attention layer arbitrarily well, where is the input feature dimension. Lastly, under some assumptions, we provide an explicit construction of MPNN+ VN with width and depth approximating the self-attention layer in GT arbitrarily well. On the empirical side, we demonstrate that 1) MPNN+ VN is a surprisingly strong baseline, outperforming GT on the recently proposed Long Range Graph Benchmark (LRGB) dataset, 2) our MPNN+ VN improves over early implementation on a wide range of OGB datasets and 3) MPNN+ VN outperforms Linear Transformer and MPNN on the climate modeling task. 

   more » « less
4. A Global Geometric Analysis of Maximal Coding Rate Reduction

   Wang, Peng; Liu, Huikang; Pai, Druv; Yu, Yaodong; Zhu, Zhihui; Qu, Qing; Ma, Yi (June 2024, International Conference on Machine Learning)

   The maximal coding rate reduction (MCR2) objective for learning structured and compact deep representations is drawing increasing attention, especially after its recent usage in the derivation of fully explainable and highly effective deep network architectures. However, it lacks a complete theoretical justification: only the properties of its global optima are known, and its global landscape has not been studied. In this work, we give a complete characterization of the properties of all its local and global optima, as well as other types of critical points. Specifically, we show that each (local or global) maximizer of the MCR2 problem corresponds to a low-dimensional, discriminative, and diverse representation, and furthermore, each critical point of the objective is either a local maximizer or a strict saddle point. Such a favorable landscape makes MCR2 a natural choice of objective for learning diverse and discriminative representations via first-order optimization methods. To validate our theoretical findings, we conduct extensive experiments on both synthetic and real data sets. 

   more » « less
5. Global Optimality Beyond Two Layers: Training Deep ReLU Networks via Convex Programs.

   Ergen, T.; Pilanci, M. (January 2021, International Conference on Machine Learning (ICML) 2021)

   Understanding the fundamental mechanism behind the success of deep neural networks is one of the key challenges in the modern machine learning literature. Despite numerous attempts, a solid theoretical analysis is yet to be developed. In this paper, we develop a novel unified framework to reveal a hidden regularization mechanism through the lens of convex optimization. We first show that the training of multiple threelayer ReLU sub-networks with weight decay regularization can be equivalently cast as a convex optimization problem in a higher dimensional space, where sparsity is enforced via a group `1- norm regularization. Consequently, ReLU networks can be interpreted as high dimensional feature selection methods. More importantly, we then prove that the equivalent convex problem can be globally optimized by a standard convex optimization solver with a polynomial-time complexity with respect to the number of samples and data dimension when the width of the network is fixed. Finally, we numerically validate our theoretical results via experiments involving both synthetic and real datasets. 

   more » « less