The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training
Extracted source text · arXiv 2502.10927
\documentclass[11pt,twocolumn]{article}
\setlength{\columnsep}{20pt}
\usepackage[a4paper,bindingoffset=0.in,
left=.8in,right=.8in,top=1in,bottom=1in,
footskip=.25in]{geometry}
\usepackage{tikz-cd}
\usepackage{pifont}
\usepackage{graphicx}
\usepackage{float}
\usepackage{xcolor}
\definecolor{darkred}{rgb}{0.5,0.2,0.2}
\definecolor{darkblue}{rgb}{0.2,0.2,0.6}
\usepackage[hidelinks,colorlinks=true,linkcolor=darkred,citecolor=darkred,urlcolor=darkblue]{hyperref}
\usepackage{titlesec}
\titleformat{\section}[hang]{\bfseries\Large}{\thesection.}{0.3em}{}
\titleformat{\subsection}[hang]{\bfseries\large}{\thesubsection.}{0.3em}{}
\usepackage{gensymb}
\usepackage{textcomp}
% \usepackage{caption}
% \DeclareCaptionFormat{myformat}{#1#2#3\hrulefill}
% \captionsetup[figure]{format=myformat}
% \usepackage[labelfont=bf]{caption}
% For theorems and such
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
\usepackage{dsfont}
\usepackage{txfonts}
\usepackage[euler]{textgreek}
\theoremstyle{plain}
\newtheorem{theorem}{Theorem}[section]
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
\theoremstyle{definition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
\theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
\let\bm\mathbf
\usepackage{enumitem}
\usepackage[capitalize,noabbrev]{cleveref}
\usepackage{booktabs}
\usepackage{natbib}
\bibliographystyle{plainnat}
\usepackage{authblk}
\newcommand\blfootnote[1]{%
\begingroup
\renewcommand\thefootnote{}\footnote{#1}%
\addtocounter{footnote}{-1}%
\endgroup
}
%%%%%%%%%%%%%%%%%%%%
\title{The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training}
\author[1,*]{Matteo Saponati}
\author[1,2,*]{Pascal Sager}
\author[1]{Pau Vilimelis Aceituno}
\author[1,2,3]{\\Thilo Stadelmann}
\author[1,4]{Benjamin Grewe}
\affil[1]{\small Institute of Neuroinformatics, ETH Zürich and University of Zürich, Zürich, Switzerland}
\affil[2]{\small Centre for Artificial Intelligence, Zürich University of Applied Sciences, Winterthur, Switzerland}
\affil[3]{\small ECLT European Centre for Living Technology, Venice, Italy}
\affil[4]{\small ETH AI Center, Zürich, Switzerland}
\affil[*]{\small equal contribution}
\date{}
%%%%%%%%%%%%%%%%%%%%
\begin{document}
\maketitle
\begin{abstract}
\normalsize \input{ICML2025-rebuttal/sections/abstract}
\end{abstract}
\section{Introduction}
%
\blfootnote{Correspondence: \href{mailto:masapo@ini.ethz.ch}{masapo@ini.ethz.ch}, \href{mailto:bgrewe@ethz.ch}{bgrewe@ethz.ch}}
%
\label{sec-introduction}
\input{ICML2025-rebuttal/sections/introduction}
\section{Autoregressive and bidirectional training leads to directional and symmetric weight updates}
\label{sec-results-math}
\input{ICML2025-rebuttal/sections/results-math}
%
\section{Symmetric and directional structures are predominant in Transformer models}
\label{sec-results-pretrained-models}
\input{ICML2025-rebuttal/figures/figure-pretrained-models-language}
\input{ICML2025-rebuttal/sections/results-pretrained-models}
%
\section{Experiments}
\label{sec–results–experiments}
\input{ICML2025-rebuttal/sections/results-experiments}
\section{Related work}
\label{sec-related-work}
\input{ICML2025-rebuttal/sections/realted-work}
\section{Discussion}
\label{sec–discussion}
\input{ICML2025-rebuttal/sections/discussion}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection*{Acknowledgments}
%
MS was supported by an ETH Postdoctoral fellowship (reference nr. 23-2 FEL-042).
%
PVA was supported by a UZH Postdoctoral fellowship (reference nr. K-76207-01-01).
%
TS and PS were supported by DIZH Fellowship ``Stability of self-organizing net fragments as inductive bias for next-generation deep learning'' from the ZHAW.
%
We thank Yassine Taoudi-Benchekroun, the Grewe lab, Alexander Nedergaard, Thomas Hofmann, and the Data Analytics lab for their helpful comments and insightful discussions.\\
%
\subsection*{Authorship contributions}
%
Conceptualization: MS, TS, BG.
%
Mathematical analysis: MS, PVA.
%
Analysis of pre-trained models: MS, PS.
%
Experiments: PS.
%
Writing: MS, PS, BG.
%
Supervision: TS, BG.
%
\subsection*{Code availability}
%
The code to reproduce our results and all the Figures in the main text and Supplementary Materials is freely available at \href{https://github.com/matteosaponati/attention-geometry}{github.com/matteosaponati/attention-geometry}.
%
%
\subsection*{Declaration of Interests}
%
The authors declare no competing interests.
\bibliography{ICML2025-rebuttal/bib}
\newpage
\setcounter{section}{0}
\setcounter{subsection}{0}
\setcounter{equation}{0}
\setcounter{figure}{0}
\setcounter{table}{0}
\setcounter{footnote}{0}
\makeatletter
\renewcommand{\thesection}{S\arabic{section}}
\renewcommand{\thesubsection}{S\arabic{section}.\arabic{subsection}}
\renewcommand{\theequation}{S\arabic{equation}}
\renewcommand{\thefigure}{S\arabic{figure}}
\renewcommand{\thetable}{S\arabic{table}}
% Fix hyperlinks so that section "S2" links to the correct place
\renewcommand{\theHsection}{S\arabic{section}}
\renewcommand{\theHsubsection}{S\arabic{section}.\arabic{subsection}}
\makeatother
\onecolumn
\section{Mathematical proofs}
\label{sec-proof-statements}
\input{ICML2025-rebuttal/supplementary/math-proofs}
\section{Experimental Details}
\label{sec-experimental-details}
\input{ICML2025-rebuttal/supplementary/experimental-details}
\section{Enforcing symmetry at initialization for vision transformers}
\label{sec-experimental-results-vit}
\input{ICML2025-rebuttal/supplementary/vision-transformer-experiments}
\section{Supplementary Figures}
\label{sec-supp-figures}
\input{ICML2025-rebuttal/supplementary/supplementary-figures}
\section{Supplementary Tables}
\label{sec-supp-tables}
\input{ICML2025-rebuttal/supplementary/supplementary-tables}
\end{document}
%
% Transformer models learn powerful representations by integrating contextual information through self-attention.
%
% Yet, the distinct impacts of autoregressive (decoder-only) and bidirectional (encoder-only) training on self-attention matrices are not well understood.
%
%
% In this work, we propose a novel approach: instead of examining individual attention matrices, we formalize the product of query and key (the composite matrix) as a bilinear form and study its learning dynamics.
%
% Additionally, we show that these structural properties emerge dynamically during training, with distinct patterns across different layers.
%
% Interpreting self-attention matrices as bilinear forms reveals a striking difference between encoder-only and decoder-only models and the structures learned during training.
%
Self-attention is essential to Transformer architectures, yet how information is embedded in the self-attention matrices and how different objective functions impact this process remains unclear.
%
We present a mathematical framework to analyze self-attention matrices by deriving the structures governing their weight updates.
%
Using this framework, we demonstrate that bidirectional training induces symmetry in the weight matrices, while autoregressive training results in directionality and column dominance.
%
Our theoretical findings are validated across multiple Transformer models — including ModernBERT, GPT, LLaMA3, and Mistral — and input modalities like text, vision, and audio.
%
Finally, we apply these insights by showing that symmetric initialization improves the performance of encoder-only models on language tasks.
%
This mathematical analysis offers a novel theoretical perspective on how information is embedded through self-attention, thereby improving the interpretability of Transformer models.
%
% This mathematical analysis of self-attention strengthens our understanding of Transformer models and provides a new theoretical perspective on self-attention mechanisms.
%Transformer models are usually trained with autoregressive or bidirectional objectives, yet the distinct impacts of these two training schemes on self-attention matrices are poorly understood.
%
%This work presents a new approach to studying self-attention matrices, introducing a novel perspective on these matrices and deriving the mathematical structures that govern their weight updates.
%
%Through this lens, we prove the connections between self-attention matrices and bilinear forms, showing that bidirectional training induces symmetry while autoregressive training enforces directionality and column dominance.
%
%We validate these theoretical findings across multiple Transformer models — including ModernBERT, GPT, LLaMA3, and Mistral — and input modalities such as text, vision, and audio.
%
%Additionally, we show that symmetry and directionality emerge dynamically during training, with distinct patterns across different layers.
%
%Finally, we demonstrate a practical application of these insights by showing that symmetric initialization enhances the performance of encoder-only models on language tasks.
%
%% New introduction
Transformer models now achieve state-of-the-art performance across a wide range of tasks and domains \citep{radfordLanguageModelsAre2019, dosovitskiyImageWorth16x162021, radford_robust_2023}.
%
Despite their success, the internal mechanisms governing their decision-making processes remain poorly understood, raising concerns regarding model alignment, reliability, and safety \citep{wang_aligning_2023, yao_survey_2024}.
%
A key challenge in understanding these models is unraveling the structures of self-attention, which is essential to Transformer architectures.
%
Current literature largely overlooks the nature of the self-attention weight matrices during autoregressive training, where the model predicts the next token in a sequence given previous ones \citep{radfordLanguageModelsAre2019,blackGPTNeoLargeScale2021,touvronLLaMAOpenEfficient2023} and bidirectional training, where the model predicts a missing token given the full sequence \citep{devlinBERTPretrainingDeep2019, baoBEiTBERTPreTraining2022,warnerSmarterBetterFaster2024}.
%
Understanding self-attention requires answering two fundamental questions:
%
%
%% general introduction to self-attention
%
%Transformer networks now deliver state-of-the-art performance across various domains \citep{}.
%
%Pretraining with self-supervised objectives is crucial for their success, typically optimizing models to either (a) predict the next token in a sequence given the previous ones (autoregressive training \citep{radfordLanguageModelsAre2019,blackGPTNeoLargeScale2021,touvronLLaMAOpenEfficient2023}), or (b) predict a missing token given the full sequence (bidirectional training \citep{devlinBERTPretrainingDeep2019, baoBEiTBERTPreTraining2022,warnerSmarterBetterFaster2024}).
%
%The intrinsic differences in these training schemes suggest that Transformers might learn different self-attention matrices in the two cases.
%
%% statement of the problem/knowledge gap
%However, the nature of the self-attention weight matrices that Transformers learn remains largely unexplored.
%
How can we interpret the structures learned in the self-attention matrices?
%
What is the impact of different objective functions on these matrices?
%
%Identifying such structures is fundamental not only for improving the performance of Transformer models, but also for their safety, alignment, and interpretability \citep{olahMechanisticInterpretabilityVariables2022}.
%
%% limitations of other mechanistic interpretability work
%
Previous work used sparse auto-encoders to identify interpretable features \citep{Huben_Cunningham_Smith_Ewart_Sharkey_2024, bricken2023monosemanticity}, circuit analysis to interpret Transformer components \citep{olahZoomIntroductionCircuits2020,elhageMathematicalFrameworkTransformer2021,olahMechanisticInterpretabilityVariables2022}, and techniques like the logit lens to analyze self-attention mechanisms \citep{Geva_Schuster_2021, darAnalyzingTransformersEmbedding2023} (for a detailed discussion, see Section \ref{sec-related-work}).
%
%
%
%Previous work demonstrates that sparse auto-encoders can identify units that consistently activate for specific input features, suggesting that Transformers can be described as compositions of interpretable feature detectors \citep{Huben_Cunningham_Smith_Ewart_Sharkey_2024, bricken2023monosemanticity}.
%
%While effective, this approach is data-dependent and does not reveal the structure of weight matrices.
%
%Other approaches focus on circuit analysis, mapping Transformer components onto interpretable computational graphs to study their functions \citep{olahZoomIntroductionCircuits2020,elhageMathematicalFrameworkTransformer2021,olahMechanisticInterpretabilityVariables2022}.
%
%Additionally, methods like the \emph{logit lens} project the query $\bm{W}_q$ and key $\bm{W}_k$ matrices into the vocabulary space to study how self-attention mixes and transfers information \citep{Geva_Schuster_2021, darAnalyzingTransformersEmbedding2023}.
%
However, these methods do not reveal the structural patterns in self-attention matrices or the transformations they encode.
%
Crucially, how autoregressive and bidirectional training shape specific weight structures remains unclear.
%% results/contributions
%
To address this gap, we introduce a novel framework for analyzing self-attention matrices and understanding how different objective functions define their weight updates.
%
We then use this framework to derive understandable mathematical structures that should emerge from such updates.
%
Finally, we verify these interpretable structures numerically on many pre-trained and custom models across different modalities, supporting the universality of our results.
%
Identifying these universal structures is fundamental not only for improving the performance of Transformer models, but also for their safety, alignment, and interpretability \citep{olahMechanisticInterpretabilityVariables2022}.
%We then use this framework to connect the matrix $\bm{W}_{qk} = \bm{W}_q\bm{W}_k^\top$ of self-attention with bilinear forms, offering novel insights compared to studying query and key matrices alone.
%
Specifically, we connect the matrix $\bm{W}_{qk} = \bm{W}_q\bm{W}_k^\top$ of self-attention with bilinear forms, offering novel insights compared to studying query and key matrices alone.
%
We reveal structured patterns in the implicit weight updates of $\bm{W}_{qk}$, uncovering key differences between encoder-only and decoder-only models:
%
\begin{enumerate}[noitemsep,nolistsep]
%
\item \emph{Decoder-only models}: Training with an autoregressive objective produces a few columns with disproportionately high norms, introducing directionality in $\bm{W}_{qk}$.
\item \emph{Encoder-only models}: Bidirectional optimization induces symmetric structures in $\bm{W}_{qk}$, reflecting the balanced nature of the training objective.
\item We validate these theoretical findings across diverse Transformer architectures and input modalities, showing that they generalize across models and tasks.
\item Empirically, we find that symmetric structures in $\bm{W}_{qk}$ enhance training efficiency for encoder-only models, leading to higher accuracy and faster convergence in language tasks.
%
\end{enumerate}
%
%
In this section, we introduce a novel framework that links the self-attention matrices $\bm{W}_{qk}$ to bilinear forms, enabling us to analyze how an objective function influences their structure.
%
This approach reveals fundamental patterns in $\bm{W}_{qk}$ that are not apparent when examining $\bm{W}_q$ and $\bm{W}_k$ separately.
%
For example, we prove that autoregressive and bidirectional training induce directional and symmetric structures in the $\bm{W}_{qk}$ matrices.
%
In the following sections, we define and formalize these concepts.
%
\subsection{Interpreting self-attention with bilinear forms}
\label{sec:self-attention-bilinear-form}
%
Self-attention \citep{vaswaniAttentionAllYou2017,radfordLanguageModelsAre2019} is
a type of score function $A: \mathbf{R}^{N,d} \times \mathbf{R}^{N,d} \rightarrow \mathbf{R}^{N,N}$ that maps a sequence of $N$ token embeddings with dimension $d$ into a matrix of attention scores.
%
Except for the row-wise softmax function $\sigma(\cdot)$, self-attention is a linear transformation of the embedded tokens.
%
In particular,
%
\begin{equation}
\label{eq:results:self-attention-bilinear-form}
\begin{split}
\bm{A}(\bm{X})
& = \sigma\left(\frac{1}{\sqrt{d}}\,\hat{\bm{A}}(\bm{X})\right) \\
& = \sigma\left(\frac{1}{\sqrt{d}}\,\bm{Q} \bm{K}^T\right) \\
& = \sigma\left(\frac{1}{\sqrt{d}}\,\bm{X} \bm{W}_{qk}\bm{X}^T\right) \,,
\end{split}
\end{equation}
%s
where $\hat{\bm{A}}(\bm{X})$ is the linear part of self-attention (raw unscaled attention scores), $\bm{X} = [\bm{x}^\top_1, \dots, \bm{x}^\top_N] \in \mathbb{R}^{N,d} $ is the sequence of $N$ token embeddings $\bm{x}_i \in \mathbb{R}^d$, and $\bm{W}_{qk} = \bm{W}_q\bm{W}_k^\top \in \mathbb{R}^{d,d}$.
%
This equation shows that the linear transformation $\bm{W}_q$ and $\bm{W}_k$ are always combined to compute attention scores with one single matrix $\bm{W}_{qk}$.
%
While the matrices $\bm{W}_q$ and $\bm{W}_k$ are defined separately for computational efficiency, this formulation remains mathematically equivalent \citep[see also][]{elhageMathematicalFrameworkTransformer2021, olssonIncontextLearningInduction2022,darAnalyzingTransformersEmbedding2023}.
%
We observe that $\bm{X} \bm{W}_{qk}\bm{X}^\top$ corresponds to a bilinear form (see Definition \ref{def-bilinear-form}).
%
Specifically, the entry $\hat{\alpha}_{ij} = [\hat{\bm{A}}]_{ij}$ can be formulated in two equivalent ways: (1) as the canonical dot product between a query $\bm{q}_i$ and a key $\bm{k}_j$ (like in standard Transformer models), or (2) as the dot product between tokens $\bm{x}_i$ and $\bm{x}_j$ under the bilinear form $\bm{W}_{qk}$,
%
\begin{equation}
\hat{\alpha}_{ij} = \langle \bm{q}_i, \,\bm{k}_j \rangle = \langle \bm{x}_i, \bm{W}_{qk} \bm{x}_j \rangle = \langle \bm{x}_i, \bm{x}_j \rangle_{\bm{W}_{qk}} \,.
\end{equation}
%
Intuitively, this equivalence shows that the matrices $\bm{W}_q$ and $\bm{W}_k$ together define an alternative metric in the embedding space, which quantifies the score of $\bm{x}_i$ and $\bm{x}_j$ without explicitly constructing the query and key vectors.
%
Let us consider a specific input token $\bm{x}_i \in \mathbb{R}^d$.
%
The self-attention layer maps it to an updated token $\hat{\bm{x}}_i \in \mathbb{R}^d$ as follows,
%
\begin{equation}
\hat{\bm{x}}^\top_i = \bm{x}^\top_i + \sum_{j= 1}^N \hat{\alpha}_{ij}\,\bm{x}_j\bm{W}_v \,,
\label{eq-prop-self-attention-bilinear-form}
\end{equation}
%
where the coefficients $\{\hat{\alpha}_{ij}\}$ in the second term represent the projection of $\bm{x}_i$ onto $\text{span}\{\bm{X}\}$ (the subspace spanned by $\bm{X} = \{\bm{x}_0^\top, \bm{x}_1^\top, \dots, \bm{x}_N^\top\})$ in the transformed embedding space defined by $\bm{W}_{qk}$,
%
\begin{equation}
\sum_j\hat{\alpha}_{ij}\, \bm{x}_j = \sum_j \langle \bm{x}_i, \bm{x}_j \rangle_{\bm{W}_{qk}} \,\bm{x}_j \,.
\label{eq-proof-self-attention-bilinear-form}
\end{equation}
%
Because the softmax function preserves the order of its input values, the ranking of the raw attention scores $\hat{\alpha}_{ij}$ is maintained in the final normalized scores $\alpha_{ij}$,
%
\begin{equation}
\alpha_{ij} < \alpha_{ij'} \, \Leftrightarrow \,\hat{\alpha}_{ij} < \hat{\alpha}_{ij'}\quad \forall i,j, j' \,.
\end{equation}
%
As a result, the attention weights ${\alpha_{ij}}$ define a convex combination of the token vectors ${\bm{x}_j}$, meaning the output lies within the convex hull of the input sequence $\text{Conv}(\bm{X}) \subset \text{span}\{\bm{X}\}$ (see also Appendix \ref{supp-math-self-attention-bilinear-form}).
%
We note that $\bm{W}_v$ is a linear transformation that is applied independently to each projection in the sum and thus does not influence our derivation.
% It follows that $\bm{W}_{qk}$ defines the geometric structure that encodes the relevance of the token embedding $\bm{x}_j$ in predicting $\bm{x}_i$ through projections within the subspace (see Equation \eqref{eq-prop-self-attention-bilinear-form}).
% %
% Therefore, token prediction depends on learning an optimal bilinear form $\bm{W}_{qk}$ to project tokens in the embedding space effectively.
In the following sections, we demonstrate how this equivalent formulation of self-attention provides a useful framework for analyzing the training of Transformer models.
%
Specifically, we show that the choice of the objective function such as autoregressive prediction \citep{radfordLanguageModelsAre2019} or bidirectional training \citep{devlinBERTPretrainingDeep2019} produces distinct structural patterns in $\bm{W}_{qk}$.
%
\subsection{Deriving the gradients of self-attention with bilinear forms}
\label{sec:self-attention-derivation-gradient}
%
\input{ICML2025/figures/figure-illustration-one-column}
%%
%%
%
To show the connection between the objective function and the structures of self-attention matrices, we derive a convenient formulation for the weight update of $\bm{W}_{qk}$.
%
We first formulate a sequence modeling problem with self-supervised training as follows.
%
Let \( U = \{t_1, \dots, t_N\} \) be a sequence of \( N \) tokens.
%
For each position \( i \), let \( \mathcal{D}_i \subseteq \{1, 2, \dots, N\} \) denote the \emph{conditioning set}, that is, the set of indices corresponding to the tokens used to predict the token \( t_i \).
%
This formulation makes it possible to isolate and analyze the contribution of each token \( t_j \in \mathcal{D}_i \) to the prediction of \( t_i \).
%
Let $\mathcal{L}(U)$ be the negative log-likelihood of each token $t_i$, expressed as,
%
\begin{equation}
\mathcal{L}(U;\mathcal{W}) = \sum_{i=1}^N \mathcal{L}(t_i)= - \sum_{i=1}^N \log p_{\mathcal{W}}(t_i \,|\, t_{\mathcal{D}_i}) \,,
\end{equation}
%
where $\mathcal{W}$ is the set of trainable parameters, and $p_{\mathcal{W}}(t_i \,|\, t_{\mathcal{D}_i}) $ are the conditional densities parameterized by the model (see also Appendix \ref{supp-math-self-attention-derivation-gradients}).
%
Here, we demonstrate that the updates to the matrix $\bm{W}_{qk}$ follow a structured pattern: the contribution of the token $t_j $ in the prediction of the embedding $t_i$ results in adding a rank-1 matrix $\bm{K}_{ij}$ to the matrix $\bm{W}_{qk}$.
%
As a result, the total weight update to $\bm{W}_{qk}$ is expressed as a linear combination of these rank-1 matrices.
%
We formalize this observation in the following proposition.
%
\\
\input{ICML2025-rebuttal/math/statements/prop-gradient-self-attention}
%
We provide proof for this proposition with related remarks in Appendix \ref{supp-math-gradients-self-attention}, and an illustrative description of the forward and backward pass in Figure \ref{fig-illustration}.
%
\subsection{How context and prediction impact the gradient differently}
%
Next, we show how the formulation of $\Delta \bm{W}_{qk}$ enables the analysis of the contribution of any given token $t^*$ to the weight updates and how this affects the properties of $\bm{W}_{qk}$.
%
Indeed, Proposition~\ref{prop-gradients-self-attention} indicates that a token $t^*$ impacts the updates of $\bm{W}_{qk}$ differently when serving as context for predicting other tokens or being itself predicted.
%
When a token $t^*$ serves as context ($t_j = t^*$), the embeddings of all predicted tokens contribute to the column space of $\bm{W}_{qk}$, where the update of every $k$-th column $\bm{w}_{\cdot, k}$ is given by
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} \Big|_{t_j = t^*}= [\bm{x}_{t^*}]_k \left(\sum_{i \in \mathcal{P}_{t}}\beta_{i*}\bm{x}_i\right)\,.
\end{equation}
%
Only the embedding of $t^*$ is instead added to the row space, where the update of every $m$-th row $\Delta \bm{w}_{m, \cdot}$ is given by
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} \Big|_{t_i = t^*} = \left(\sum_{i \in \mathcal{P}_{t}}\beta_{i*}[\bm{x}_{i}]_k\right)\bm{x}_{t^*}\,.
\end{equation}
%
Intuitively, using $t^*$ as context increases the dimensionality of the column space proportionally to the embeddings of the predicted tokens, while reducing the row space along the direction of the embedding of $t^*$.
%
Conversely, when $t^*$ is being predicted ($t_i = t^*$), all token embeddings from the context are added to the row space of $\bm{W}_{qk}$, while only the embedding of $t^*$ is added to the column space.
%
Consequently, the role a token plays during training affects its contribution to the weight update differently.
%
We formalize this in the following proposition.
%
\input{ICML2025-rebuttal/math/statements/prop-gradient-columns-rows}
%
We provide proof for this proposition with related remarks in Appendix \ref{supp-math-prop-gradient-column-rows}.
%
Note that all mathematical results derived so far rely solely on the structure of self-attention and make no assumptions about the input data.
%
In the next section, we build on these results to link structural patterns in \( \bm{W}_{qk} \) with the specific form of the objective function.
\subsection{The relation between objective functions and structures in self-attention matrices}
%
In this final section, we show how to relate these properties to the specific objective function.
%
Crucially, the number of times a token \( t^* \) appears as context or as a prediction depends on the training objective.
%
Autoregressive training implicitly introduces directionality by predicting each token solely on its preceding tokens.
%
In contrast, bidirectional training uses tokens as context and as predictions symmetrically.
%
This fundamental difference affects the weight updates of \( \bm{W}_{qk} \), and consequently, the structures encoded in its rows and columns.
%
We formalize this relationship in the next two theorems, under the following assumptions: (a) tokens exhibit statistical correlations, leading to partial alignment in their embeddings; (b) the entries of \( \bm{W}_{qk} \) are i.i.d. at initialization with finite mean and variance; (c) some tokens tend to occur earlier in the sequence and are more predictive of future content; (d) bidirectional training induces approximately symmetric error signals.
%
In this setting, we prove how the objective function influences the internal structure encoded by self-attention.
%
\input{ICML2025-rebuttal/math/statements/theo-gradient-directionality}
%
\input{ICML2025-rebuttal/math/statements/theo-gradient-symmetry}
%
We provide a proof of these two Theorems, and the related Propositions and Lemmas in Appendix \ref{supp-math-directionality-symmetry}.
%
%
\begin{figure}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/figure-illustration-one-column.pdf}
\end{center}
%
\caption{
%
%
Illustration of the computation of the self-attention score between token $\bm{x}_i$ and token $\bm{x}_j$ (forward pass, see Equation \eqref{eq-prop-self-attention-bilinear-form}), and its corresponding contribution to the weight update of $\bm{W}_{qk}$ (backward pass, see Equation \eqref{eq:weight-update-bilinear-form} and Equation \eqref{eq:weight-update-bilinear-form-prediction}).
%
The symbols "$\oplus$", "$\otimes$", and "$\odot$" refer to the addition, multiplication, and dot product operations, respectively.
%
}
%
\label{fig-illustration}
%
\end{figure}
\begin{proposition}
\label{prop-gradients-self-attention}
%
(\textbf{The implicit weight update as sum of rank-1 matrices}).
%
Let \( U = \{t_1, \dots, t_N\} \) be a sequence of \( N \) tokens, and let \( \mathcal{L}(U; \mathcal{W}) \) denote the negative log-likelihood of the sequence under a Transformer model with learnable parameters \( \mathcal{W} \),
%
\begin{equation}
\mathcal{L}(U; \mathcal{W}) = \sum_i \mathcal{L}(t_i) = - \sum_i \log p_{\mathcal{W}}(t_i \mid t_{\mathcal{D}_i})\,,
\end{equation}
%
where \( \mathcal{D}_i \subseteq \{1, \dots, N\} \) is the set of indices used to condition the prediction of token \( t_i \).
%
Let the self-attention function be defined as in Equation~\eqref{eq:results:self-attention-bilinear-form}.
%
Following the gradient of \( \mathcal{L}(U; \mathcal{W}) \), the implicit gradient-based update to the bilinear weight matrix \( \bm{W}_{qk}^l \) at the \( l \)-th self-attention layer is proportionally equivalent to:
\begin{enumerate}[noitemsep]
\item
The sum of contributions from each token \( t_j: j \in \mathcal{C}_i \) used to predict \( t_i \),
\begin{equation}
\label{eq:weight-update-bilinear-form}
\Delta \bm{W}_{qk}^l \propto \sum_i \sum_{j \in \mathcal{C}_i} \Delta \bm{W}_{qk}^l\Big|_{t_i \leftarrow t_j} = \sum_i \sum_{j \in \mathcal{C}_i} \beta^l_{ij} \bm{K}^{l-1}_{ij}\,,
\end{equation}
where \( \mathcal{C}_i \) denotes the set of context indices for predicting \( t_i \).
%
\item
The sum of contributions from each prediction target \( t_i: i \in \mathcal{P}_j \) when predicted with \( t_j \),
\begin{equation}
\label{eq:weight-update-bilinear-form-prediction}
\Delta \bm{W}_{qk}^l \propto \sum_{i \in \mathcal{P}_j} \sum_j \Delta \bm{W}_{qk}^l\Big|_{t_i \leftarrow t_j} = \sum_{i \in \mathcal{P}_j} \sum_j \beta^l_{ij} \bm{K}^{l-1}_{ij}\,,
\end{equation}
where \( \mathcal{P}_j \) denotes the set of tokens for which \( t_j \) is included in the context.
\end{enumerate}
%
Here, \( \beta^l_{ij} \) is a scalar that quantifies the contribution of the token embedding \( \bm{x}_j \) to the prediction error for \( \bm{x}_i \) at the \( l \)-th layer, and the matrix \( \bm{K}^{l-1}_{ij} \in \mathbb{M}_d \) is a rank-1 matrix defined by the outer product of the token embeddings at the previous layer,
%
\begin{equation}
\bm{K}^{l-1}_{ij} = \bm{x}^{l-1}_i {\bm{x}^{l-1}_j}^\top\,.
\end{equation}
%
\end{proposition}
\begin{proposition}
\label{prop-gradient-columns-rows}
%
(\textbf{Different implicit updates for context and prediction}).
%
%
Let $U = \{t_1, \dots, t_N\}$ be a sequence of tokens and let $\bm{x}_i \in \mathbb{R}^d$ be the embedding of the $i$-th token.
%
Let $\Delta \bm{W}_{qk}$ be the weight update from Proposition \ref{prop-gradients-self-attention}.
%
Let $t^*$ be a given token in $U$.
%
When using $t^*$ as context, the $k$-th column of $\Delta \bm{W}_{qk}$ is given by
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} \bigg|_{t_j = t^*}= [\bm{x}_{t^*}]_k \left(\sum_{i \in \mathcal{P}_{t^*}}\beta_{i*}\bm{x}_i\right)\,,
\end{equation}
%
while the $m$-th row of $\Delta \bm{W}_{qk}$ is given by
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} \bigg|_{t_j = t^*}= \left(\sum_{i \in \mathcal{P}_{t^*}}\beta_{i*}[\bm{x}_i]_m\right)\bm{x}_{t^*} \,.
\end{equation}
%
%where $\P_{t^*}$ is the set of tokens predicted by $t^*$.
%
When predicting $t^*$, the $k$-th column of $\Delta \bm{W}_{qk}$ is given by
%
\begin{equation}
\Delta \bm{w}_{k, \cdot} \bigg|_{t_i = t^*} = \left(\sum_{j \in \mathcal{C}_{t^*}}\beta_{*j}[\bm{x}_j]_k\right)\bm{x}_{t^*} \,,
\end{equation}
%
while the $m$-th row of $\Delta \bm{W}_{qk}$ is given by
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} \bigg|_{t_i = t^*} = [\bm{x}_{t^*}]_m \left(\sum_{j \in \mathcal{C}_{t^*}}\beta_{*j}\bm{x}_j^\top\right)\,.
\end{equation}
%
%where $C_{t^*}$ is the set of context tokens for $t^*$.
%
\end{proposition}
%
% First we consider the gradient of a column for a given predictor-predicted pair $ij$, and derive the gradient associated to it
% \begin{equation}
% \Delta \bm{w}_{\cdot, k}(i,j) = \Delta \bm{w}_{\cdot, k}\Big|_{\bm{t}_i \leftarrow \bm{t}_j}
% =
% (\beta_{ij}[\bm{x}_j]_k)\,\bm{x}_i \,,
% \end{equation}
% where clearly the direction is given by the $\bm{x}_i$
% \begin{equation}
% \Delta \bm{w}_{\cdot, k}(i,j) \mathbin{\|} \bm{x}_i.
% \end{equation}
% Second, we note that the norm is proportional to the error ($\bm{t}_i - \sigma(\bm{z}_i)$ and is also proportional to the occurrence of $j$ before $i$. Assuming that most words have low predictive power, we would have that for each $i$ there are few $j$ with high gradients.
% Now we look at the norm of the gradient for one column for the rest (for one of the high-gradient $j$)
% \begin{equation}
% \dfrac{\|\Delta \bm{w}_{\cdot, k}(i,j)\|}{\sum_k \|\Delta \bm{w}_{\cdot, c}(i,j)\|}
% \end{equation}
% and since $\beta_{ij}$, is the same scalar for all columns and the $\bm{x}_i$ is the same vector for all columns,
% \begin{equation}
% \dfrac{\|\Delta \bm{w}_{\cdot, k}(i,j)\|}{\sum_c \|\Delta \bm{w}_{\cdot, k}(i,j)\|}
% = \dfrac{|[\bm{x}_j]_k|}{\sum_c |[\bm{x}_j]_k|}.
% \end{equation}
% Thus, for each pair $i,j$ the gradient would grow proportionally to $\bm{x}_j$ in the direction of $\bm{x}_i$.
% %Stop point 1: We assume that there is a low-rank structure in the data: few j that are predicting i. Then
% Now we study how does the gradient change the attention, which is quantified
\begin{theorem}
\label{theo-gradient-directionality}
%
(\textbf{Autoregressive training induces directionality})
Let \( \mathcal{V} = \{t_0, \dots, t_V\} \) be a vocabulary of tokens, and let \( \mathcal{U} \subset \mathcal{V}^N \) denote the sample space of all sequences of length \( N \). Let \( U = \{t_1, \dots, t_N\} \in \mathcal{U} \) be a random variable with \( U \sim P(U) \). Define \( \Pr[t_j = t^*] \) as the marginal probability that the token at position \( j \) equals \( t^* \in \mathcal{V} \).
%
Let \( \{\bm{x}_1, \dots, \bm{x}_N\} \) be the token embeddings corresponding to the elements of \( U \), where each embedding \( \bm{x}_i \sim \mathcal{D} \) is drawn i.i.d. from a distribution \( \mathcal{D} \) with zero mean and covariance matrix $\text{Cov}(\bm{x}_i) = \Sigma$..
%
Let \( \bm{W}_{qk} \) be query-key matrix of a self-attention mechanism, and let \( \Delta \bm{W}_{qk} \) denote its gradient update as defined in Proposition~\ref{prop-gradients-self-attention}, computed under an autoregressive objective as in Definition~\ref{def-objective-functions}.
%
It follows that the contribution of the token \( t^* \) after training satisfies,
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ \|\Delta \bm{w}_{\cdot, k}\|^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ \|\Delta \bm{w}_{m, \cdot}\|^2 \right]} > 1
\quad \forall k, \,\,\forall m\,\, \text{s.t.} \,\, \Sigma_{m,m} < \frac{\mathrm{Tr}(\Sigma)}{d} \,.
\end{equation}
%
Moreover, there exists a scalar \( \gamma \in \mathbb{R}_{>0} \), proportional to the product of the row and column norm variances, $
\gamma \propto \mathrm{Var}\left(\|\bm{w}_{m,\cdot}\|\right) \cdot \mathrm{Var}\left(\|\bm{w}_{\cdot,k}\|\right)$,
such that for all \( w > \gamma \),
\begin{equation}
\Pr\left[\|\bm{w}_{\cdot, k}\|^2 > w\right] > \Pr\left[\|\bm{w}_{m, \cdot}\|^2 > w\right] \,,
\end{equation}
that is, columns of \( \bm{W}_{qk} \) are more likely than rows to have high norm under autoregressive training, indicating a directional bias in gradient updates.
%
\end{theorem}
\begin{theorem}
\label{theo-gradients-symmetry}
%
(\textbf{Bidirectional training induces symmetry}).
%
Let $U = \{t_1, \dots, t_N\}$ be a sequence of tokens.
%
Let \( \bm{W}_{qk} \) be the query-key matrix of a self-attention mechanism, and let \( \Delta \bm{W}_{qk} \) denote its gradient update as defined in Proposition~\ref{prop-gradients-self-attention}, computed under a bidirectional objective as in Definition~\ref{def-objective-functions}.
%
It follows that,
%
\begin{equation}
\Delta \bm{W}_{qk} = \sum_{i = 1}^N \sum_{j = 1}^N \beta_{ij} \bm{K}_{ij} \,,
\end{equation}
%
and that every pair $(i,j)$ with $i \neq j$ contributes to the weight update with a term
%
\begin{equation}
\Delta \bm{W}_{qk}\big|_{\bm{t}_i \leftrightarrow \bm{t}_j} = \beta_{ij}\bm{K}_{ij} + \beta_{ji}{\bm{K}_{ij}}^T \,,
\end{equation}
%
that is approximately symmetric,
%
\begin{equation}
\Delta \bm{W}_{qk}\big|_{\bm{t}_i \leftrightarrow \bm{t}_j} \approx \Delta \bm{W}_{qk}\big|^\top_{\bm{t}_i \leftrightarrow \bm{t}_j} \,.
\end{equation}
%
\end{theorem}
%
\begin{figure*}[!ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/figure-pretrained-models-language.pdf}
\end{center}
%
\caption{
%
%
\textbf{a})
%%
Left) Median symmetry score of the matrix $\bm{W}_{qk}$ as a function of the total number of parameters.
%
Each dot corresponds to the median and the interquartile range across layers of a given pre-trained model (see Tables in Appendix \ref{table:directionality-score-models}).
%
Right) Example of structures in the $\bm{W}_{qk}$ matrix of an encoder-only model (BERT Tiny, layer 1 \citep{turcWellReadStudentsLearn2019})
%
\textbf{b})
%%
Left) Same as in \textbf{a} for the median directionality score of the matrix $\bm{W}_{qk}$.
%
Right) Example of structures in the $\bm{W}_{qk}$ matrix of a decoder-only model (TinyStories GPT, layer 1 \citep{eldanTinyStoriesHowSmall2023})
}
%
\label{fig-pretrained-models}
%
\end{figure*}
In this section, we validate our theoretical findings by quantifying empirically the degree of symmetry and directionality in different families of open-source Transformer models.
%
To do so, we define two scores for symmetry and directionality in square matrices.
%
First, we define the symmetry score $s \in \mathbb{R}$ as follows,
%
\input{ICML2025/math/statements/def-symmetry-score}
%
Here, positive and negative symmetry scores indicate the presence of symmetric and skew-symmetric structures, respectively (see Appendix \ref{supp-math-symmetry-score}).
%
Second, we define the directionality score $d \in \mathbb{R}$ as follows,
%
\input{ICML2025/math/statements/def-directionality-score}
%
Here, positive and negative directionality scores indicate the dominance of high norm rows or columns, respectively (see Appendix \ref{supp-math-directionality-score}).
%
Finally, we compute the matrix $\bm{W}_{qk}$ for every layer, calculate the median symmetry and directionality score across layers, and analyze the differences between encoder- and decoder-only variants.
We find that encoder-only models remarkably show a higher degree of symmetry than decoder-only (Figure \ref{fig-pretrained-models}a).
%
This difference is consistent across multiple families of models and input modalities, such as BERT \citep{devlinBERTPretrainingDeep2019}, GPT \citep{radfordImprovingLanguageUnderstanding2018, radfordLanguageModelsAre2019}, LLAMA3 \citep{touvronLLaMAOpenEfficient2023}, Phi \citep{hughesPhi2SurprisingPower2023,abdinPhi3TechnicalReport2024}, MISTRAL \citep{jiangMistral7B2023}, ModernBERT \citep{warnerSmarterBetterFaster2024}, and many others (see Figure \ref{suppfig-pretrained-models} for vision and audio models).
%
Strikingly, we observe that decoder-only models have higher degrees of directionality than encoder-only models (Figure \ref{fig-pretrained-models}b).
%
Again, this difference is consistent across all the models and input modalities we consider.
%
We show in Figure \ref{suppfig-encoder-decoder-models} that a similar pattern is observed when including full encoder-decoder Transformers (e.g. the language T5 models \citep{xueMT5MassivelyMultilingual2021}), despite these models having an overall lower degree of directionality.
\begin{definition}
\label{def-symmetry-score}
%
(\textbf{\emph{Symmetry score}}).
%
Given a square matrix $\bm{M} \in \mathbb{M}_n$ we define the symmetry score,
%
\begin{equation}
\label{eq:definition-symmetry-score}
s = 2 \,\dfrac{||\bm{M}_s||_F^2 - ||\bm{M}_n||_F^2}{||\bm{M}||_F^2} \,,
\end{equation}
%
where $|| \cdot ||_F$ is the Frobenious norm, and $\bm{M}_s$ and $\bm{M}_n$ are the symmetric and skew-symmetric parts of the Toeplitz decomposition of $\bm{M}$, respectively,
%
\begin{equation}
\begin{split}
\bm{M}_s = \frac12\big(\bm{M} + \bm{M}^\top\big) \quad ; \quad
\bm{M}_n = \frac12\big(\bm{M} -\bm{M}^\top\big) \,.
\end{split}
\end{equation}
%
\end{definition}
\begin{definition}
\label{def-directionality-score}
%
(\textbf{\emph{Directionality score}}).
%
Given a square matrix $\bm{M} \in \mathbb{M}_n$ we define the directionality score,
%
\begin{equation}
\label{eq:definition-directionality-score}
d = \dfrac{\bar{r}_{\bm{M}} - \bar{c}_{\bm{M}}}{\bar{r}_{\bm{M}} + \bar{c}_{\bm{M}}} \,,
\end{equation}
%
where $\bar{c}_{\bm{M}}$ is the sum of the norm of the columns that are higher than a given threshold, as follows,
%
\begin{equation}
\bar{c}_{\bm{M}} = \sum_{k \in \overline{\mathcal{C}}} ||\bm{m}_{\cdot, k}||_2
\,\,\, \text{with} \,\,\ \overline{\mathcal{C}} = \{k \mid \|\mathbf{m}_{\cdot,k}\|_2 > \mu_c + \gamma \sigma_c\}\,,
\end{equation}
%
where $|| \cdot ||_2$ is the L2 norm, $||\bm{m}_{\cdot, k}||_2$ is the norm of the $k$-th column, $\mu_c = \mathbb{E}[||\bm{m}_{\cdot, k}||_2]$
and $\sigma_c = \sqrt{\text{Var}||\bm{m}_{\cdot, k}||_2}$ are the mean and standard deviation of the column norms, $\gamma$ is a scaling factor, and similarly $\bar{r}_{\bm{M}}$ is the sum of the norm of the rows that are higher than a given threshold,
%
%
\begin{equation}
\bar{r}_{\bm{M}} = \sum_{k \in \overline{\mathcal{R}}} ||\bm{m}_{k, \cdot}||_2
\,\,\, \text{with} \,\,\ \overline{\mathcal{R}} = \{k \mid \|\mathbf{m}_{k, \cdot}\|_2 > \mu_r + \gamma \sigma_r\}\,,
\end{equation}
%
with $||\bm{m}_{k, \cdot}||_2$ as the norm of the $k$-th row, $\mu_r = \mathbb{E}[||\bm{m}_{k,\cdot}||_2]$
and $\sigma_r = \sqrt{\text{Var}||\bm{m}_{k, \cdot}||_2}$ as the mean and standard deviation of the row norms.
%
\end{definition}
%
In this final section, we test if using structural priors based on our previous results can improve the pretraining of Transformer models.
%
To do so, we train Transformer models from scratch and perform a series of experiments to analyze how symmetric and directional structures develop during training across layers.
%
\subsection{Evolution of symmetric and directional structures during learning}
\label{sec-evolution-structures}
%
\input{ICML2025-rebuttal/figures/figure-training-layers}
%
To test the applicability of our result, we first train 12-layer transformer models in both encoder and decoder modes and quantify the median symmetry and directionality scores across epochs.
%
At initialization, the symmetry and directionality score of the matrix $\bm{W}_{qk}$ at any layer is zero (see Definition \ref{def-symmetry-score} and Definition \ref{def-directionality-score} and related Appendix \ref{supp-math-symmetry-score} and \ref{supp-math-directionality-score}).
%
The incremental update of $\bm{W}_{qk}$ we described in the previous sections predicts that decoder-only models develop high-norm columns incrementally during training (see Theorem \ref{theo-gradient-directionality}).
%
Likewise, as symmetric weight updates are added to $\bm{W}_{qk}$ in encoder-only models, Theorem \ref{theo-gradients-symmetry} predicts that symmetric structures emerge incrementally during training.
%
% It follows from Definition \ref{def-symmetry-score} and Definition \ref{def-directionality-score} that the bilinear form at any $l$-th layer $\bm{W}^l_{qk}$ has both a symmetry score and a directionality score of zero at initialization (see Appendix \ref{supp-math-symmetry-score} and \ref{supp-math-directionality-score}).
%
% Theorem \ref{theo-unidirectional-structures}
% predicts that decoder-only models develop high-norm columns incrementally during training.
% %
% Likewise, as symmetric weight updates are added to $\bm{W}_{qk}$ in encoder-only models, Theorem \ref{theo-symmetric-structures} predicts that symmetric structures emerge incrementally during training.
% %
% To test these predictions, we train 12-layer transformer models in both encoder and decoder modes and quantify the median symmetry and directionality scores across epochs (Figure \ref{fig-training-layers}a-b).
% %
Consistent with our results on pre-trained models, encoder-only models show a higher degree of symmetry than decoder-only models (Figure \ref{fig-training-layers}a).
%
In contrast, decoder-only models have a higher directionality score (Figure \ref{fig-training-layers}b).
%
We observe this difference on all datasets we tested (Jigsaw \citep{jigsaw_challenge}, Wikipedia \citep{wikidump}, Red Pajama \citep{together2023redpajama}, see Figure \ref{suppfig-training-custom-models}).
%
Furthermore, late layers of encoder-only models are more symmetric and converge faster than early layers when training bidirectionally. At the same time, decoder-only models learn almost non-symmetric matrices with strong skew-symmetric matrices in the middle layers (Figure \ref{fig-training-layers}c).
%
When training unidirectionally, both encoder and decoder models show a higher degree of directionality for late layers, which is remarkably stronger for decoder-only models (Figure \ref{fig-training-layers}d).
%
We observe similar differences across layers with all the datasets we tested (Figure \ref{suppfig-training-custom-models-layers}), despite these models having
less significant differences in directionality scores.
%
See Appendix~\ref{sec-experimental-details} for a detailed description of the experiments.
%
\subsection{Enforcing symmetry at initialization improves the training of encoder-only models}
\label{sec-symmetric-initialization}
%
\input{ICML2025-rebuttal/tables/table-results-symmetric-initialization}
%
The previous section showed that symmetric structures incrementally emerge during training in the $\bm{W}_{qk}$ matrices of encoder-only models.
%
Here, we first provide evidence that these findings can be exploited to speed up training using symmetry as an inductive bias.
%
Specifically, we explore how symmetric initialization influences the training dynamics of the model and whether it enhances learning efficiency and overall performance.
%
We train 4-layer, 12-layer, and 24-layer encoder-only models, comparing two initialization strategies: initialize the self-attention matrices independently versus initializing the $\bm{W}_q$ and $\bm{W}_k$ in each self-attention layer to ensure that $\bm{W}_{qk}$ is symmetric (see Appendix \ref{sec-experimental-details}).
%
We report the results of our experiments in Table \ref{table:symmetric-initialization}.
%
We observe that enforcing symmetry at initialization leads to lower loss values at the end of training for most of the models.
%
Importantly, symmetric initialization significantly accelerates convergence, reaching the final loss value faster than those with random initialization (up to $~75\%$ faster for 4-layer models, $~35\%$ faster for 12-layer models, and $~44\%$ faster for 24-layer models; see also Figure \ref{suppfig-symmetric-initialization}a).
%
Moreover, we observe that self-attention matrices initialized symmetrically lose symmetry during training but converge to higher symmetry levels than random initialization (Figure \ref{suppfig-symmetric-initialization}b)
%
This symmetric initialization decreases the gap in symmetric scores between layers compared to random initialization (Figure \ref{suppfig-symmetric-initialization}c).
%
When we tested symmetric initialization on vision Transformers, the improvements were not significant compared to the gains observed in language models (see Appendix~\ref{sec-experimental-results-vit} for a detailed analysis).
%
% Interestingly, we also observe that enforcing symmetry at initialization improves accuracy on small, linear self-attention models (Figure \ref{suppfig:linear-attention-configs-symmetric-data}).
%
These results highlight that embedding symmetry as an inductive bias across all Transformer layers can enhance training efficiency and model performance.
%
% While symmetric initialization significantly improved training dynamics for text models, we did not observe comparable gains in the vision domain; see Appendix~\ref{sec-experimental-results-vit} for a detailed analysis.
% %
%
\begin{figure*}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/figure-training-layers.pdf}
\end{center}
%
\caption{
%
\textbf{a})
%
Evolution of symmetry score during training.
%
We train a Bert-base-uncased model with 12 layers on the Wikipedia dataset \citep{wikidump} in encoder-only (blue) and decoder-only (purple) mode, see legend.
%
Shown are the median and the interquartile range.
%
Shown are the median and the interquartile range.
%
\textbf{b})
%
Same as in panel \textbf{a}
for the median directionality score.
%
\textbf{c})
%
Evolution of the median symmetry score across layers of the encoder-only (left) and decoder-only (right) models.
%
Each layer is color-coded as shown on the legend.
%
\textbf{d})
%
Same as panel $\textbf{c}$ for the median directionality score.
%
}
%
\label{fig-training-layers}
%
\end{figure*}
\begin{table}[H]
\caption{The final loss at the end of training and the speed-up for the 4, 12, and 24-layer models trained on the Jigsaw dataset \citep{jigsaw_challenge}, Wikipedia \citep{wikidump}, and Red Pajama \citep{together2023redpajama}, with and without symmetry initialization (see Appendix~\ref{sec:exp_bertmodels_models}).
%
Speed-up (\%) is calculated by subtracting the epoch at which the symmetrically initialized model reaches the non-symmetric model’s final loss from the total number of epochs, and then dividing by the total number of epochs.
%
For example, a 50\% speed-up means that the model with symmetric initialization achieves the final loss of the non-symmetric model in half the number of training epochs.
%
}
\label{table:symmetric-initialization}
\vskip 0.15in
\begin{center}
\begin{small}
\begin{sc}
\begin{tabular}{lcccr}
\toprule
Model & Loss & Speed-up \\
\midrule\midrule
4-layer model & & \\
\midrule\midrule
Jigsaw & 2.782 & \\
Jigsaw (+ symm) & \textbf{2.758} & 26 \% \\
\midrule
Wikipedia & 0.984& \\
Wikipedia (+ symm) & \textbf{0.812} & 73 \%\\
\midrule
Red Pajama & 1.106 &\\
Red Pajama (+ symm) & \textbf{0.907} & 69 \%\\
\midrule\midrule
12-layer model & & \\
\midrule\midrule
Jigsaw & 1.419 & \\
Jigsaw (+ symm) & 1.430 & 0 \%\\
\midrule
Wikipedia & 0.256 & \\
Wikipedia (+ symm) & \textbf{0.247} & 20 \%\\
\midrule
Red Pajama & 0.297&\\
Red Pajama (+ symm) & \textbf{0.274} & 35 \%\\
\midrule\midrule
24-layer model & & \\
\midrule\midrule
Jigsaw & 0.786 & \\
Jigsaw (+ symm) & \textbf{0.739} & 14\%\\
\midrule
Wikipedia & 0.192 & \\
Wikipedia (+ symm) & \textbf{0.166} & 44\%\\
\midrule
Red Pajama & 0.209 &\\
Red Pajama (+ symm) & \textbf{0.189} & 34\%\\
\bottomrule
\end{tabular}
\end{sc}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
% Related work -> new version for main text instead of appendix
%
\paragraph{Mechanistic Interpretability (MI)}
%
In contrast to interpretability approaches that focus on explaining specific data instances by analyzing features \citep{Wu_Chen_2020, Lundstrom_Huang_2022}, attention scores \citep{Hoover_Strobelt_2020, Barkan_2021, Yeh_Chen_2023}, or output variations \citep{Jin_Jin_2020, Wang_Xu_2022}, mechanistic interpretability (MI) seeks to provide a more general understanding of Transformer models.
%
MI is based on the study of ``circuits,'' analyzing the between activations across different components of a Transformer \citep{olahZoomIntroductionCircuits2020}.
%
Following the categorization by \citep{Rai_Zhou_Feng_Saparov_Yao_2024}, MI techniques include:
(i)
The \emph{logit lens} \citep{nostalgebraist_2020, Geva_Schuster_2021} projects layer activations or weights into the vocabulary space $\mathcal{V}$, allowing the derivation of logits and revealing the influence of individual components on the prediction.
This technique can also be applied to query-key matrices $\bm{W}_{qk}$ to study how attention heads transform source into target tokens \citep{darAnalyzingTransformersEmbedding2023};
%
(ii)
\emph{Probing techniques} allow to identify correlations between layer activations and features by training a linear classifier or shallow neural network (the probe) to predict the presence of a feature in layer activations \citep{Dalvi_Durrani_2019, gurnee2023finding};
%
(iii)
\emph{Sparse autoencoders} map activations into a higher-dimensional yet sparse representation, facilitating the identification of independent (monosemantic) features \citep{Huben_Cunningham_Smith_Ewart_Sharkey_2024, bricken2023monosemanticity};
%
(iv)
\emph{Visualization techniques} facilitate the analysis of attention scores \citep{olssonIncontextLearningInduction2022, Lieberum_Rahtz_Kramár_Nanda_Irving_Shah_Mikulik_2023} and neuronal activity patterns \citep{Elhage_Softmax_2023} by rendering them in a graphical format. However, their utility often depends on human comprehension of the resulting visualizations;
%
(v)
\emph{Ablation studies} assess the importance of model components by systematically removing or modifying them and observing the resulting behavioral changes \citep{olssonIncontextLearningInduction2022, wangInterpretabilityWildCircuit2022};
%
(vi)
\emph{Causal mediation analysis (CMA)} analyzes the importance of components \citep{Vig_Gehrmann_Belinkov_Qian_Nevo_Singer_Shieber_2020, Meng_Bau_Andonian_Belinkov_2022} or connections \citep{wangInterpretabilityWildCircuit2022, Goldowsky-Dill_MacLeod_Sato_Arora_2023} between them by introducing perturbations (e.g., noise) and selectively patching them to measure the recovery of lost capabilities.
%
% What we do differently
Our work adds a new perspective on MI by providing a scalable and generalizable approach to the mechanistic understanding of self-attention.
In contrast to existing work, it is not limited to analyzing fully trained models but investigates the influence of learning and can analyze models of different sizes across all modalities.
\paragraph{MI for model enhancement}
Insights from MI have been instrumental in various applications, including improving model safety \citep{Belrose_Furman_Smith_Halawi_Ostrovsky_McKinney_Biderman_Steinhardt_2023}, updating the model's learned knowledge \citep{Meng_Bau_Andonian_Belinkov_2022}, and guiding the generation process of generative models \citep{Geva_Caciularu_2022}.
%
One of the most exciting applications is leveraging MI techniques to improve model performance.
For instance, \citep{Marks_Rager_Michaud_Belinkov_Bau_Mueller_2024} identify and eliminate spurious features, leading to improved generalization in classification tasks.
%
Similarly, \citep{trockman_mimetic_2023} observe that query-key matrices ($\bm{W}_{qk}$) frequently exhibit a pronounced negative diagonal, prompting them to initialize it with approximately the identity matrix, leading to enhanced accuracy in image classification tasks.
%
% What we do
Similar to \citep{trockman_mimetic_2023}, we demonstrate that findings about the structure of $\bm{W}_{qk}$ can inform initialization strategies that improve Transformer performance. However, since the identity matrix is one instance of a symmetric matrix, we consider the work by \citep{trockman_mimetic_2023} as a special instance of our broader approach, confirming our findings in the image domain.
%
In this work, we demonstrate how bidirectional and autoregressive objective functions influence the structure of the query-key matrix $\bm{W}_{qk}$ in self-attention, enhancing our understanding of Transformer models.
%
Our mathematical framework shows that bidirectional training induces symmetric structures in $\bm{W}_{qk}$, whereas autoregressive training results in matrices characterized by directionality and column dominance.
%
To empirically validate our analysis, we develop and apply symmetry and directionality scores to various Transformer encoder and decoder models across multiple modalities, including text, audio, and images.
%
Our results reveal that bidirectionally trained encoder models exhibit high symmetry, while autoregressively trained decoder models demonstrate strong directionality, thereby supporting the predictions of our mathematical framework.
%
%To validate our analysis empirically we define a symmetry and directionality score, showing consistent results across different Transformer architectures, input modalities, and layers within a model.
%
This suggests that self-attention inherently reflects these structural properties, contributing to the mechanistic interpretability of Transformer models.
%
%Our results thus indicate fundamental characteristics of the self-attention mechanism, contributing to a deeper mechanistic interpretability of Transformer models.
%
Finally, we leverage our findings to improve convergence speed during bidirectional training by initializing $\bm{W}_q$ and $\bm{W}_k$ matrices such that $\bm{W}_{qk}$ is symmetric.
%
%This result shows a practical consequence of our findings beyond their implications for interpretability.
%
While our findings mark an initial step toward leveraging symmetry for more efficient Transformer training, further research is required to assess the scalability of symmetric initialization in large-scale models and across diverse domains.
%
% Future research should validate our initialization strategy across diverse models and test its scalability in larger architectures.
%
Furthermore, it is important to explore strategies for leveraging the directionality structures of decoder-only models.
%
For instance, incorporating structural constraints into the objective function or weight regularization could enhance training efficiency and stability in autoregressive settings.
%
By bridging theoretical insights with practical improvements, our work not only advances the interpretability of self-attention but also provides a foundation for optimizing Transformer architectures.
%
Ultimately, these findings contribute to a deeper understanding of the mechanisms governing self-attention, paving the way for more reliable and efficient Transformer-based models.
%% PRELIMINARIES
%% (2) background and preliminaries
\subsection{Preliminaries}
%
Following the notation in \citep{vaswaniAttentionAllYou2017, radfordImprovingLanguageUnderstanding2018}, we define a Transformer architecture as,
%
\input{ICML2025/math/statements/def-transformer-model}
%
Furthermore, we use the following definition of a bilinear form,
\input{ICML2025/math/statements/def-bilinear-form}
%
Finally, we provide the following definition of autoregressive and bidirectional training objectives,
%
\input{ICML2025-rebuttal/math/statements/def-objective-functions}
%% RELATED REMARKS TO SELF-ATTENTION AND BILINEAR FORM
\subsection{Related remarks to Section \ref{sec:self-attention-bilinear-form}}
\label{supp-math-self-attention-bilinear-form}
%
\input{ICML2025-rebuttal/math/statements/remark-self-attention-bilinear-form-projection}
\input{ICML2025-rebuttal/math/statements/remark-self-attention-bilinear-form-multihead}
%% REMARKS ON OBJECTIVE FUNCTION, AUTOREGRESSIVE and BIDIRECTIONAL TRAINING, and BERT MLM
\subsection{Related remarks to Section \ref{sec:self-attention-derivation-gradient}}
\label{supp-math-self-attention-derivation-gradients}
%
\input{ICML2025-rebuttal/math/statements/remark-likelihood-pseudo-likelihood}
%% PROOF PROPOSITION GRADIENTS AND RANK-1 MATRIX
\subsection{Proof of Proposition \ref{prop-gradients-self-attention} and related remarks}
\label{supp-math-gradients-self-attention}
%
\input{ICML2025-rebuttal/math/proofs/proof-prop-gradient-self-attention}
%
In standard Transformer models, the bilinear form $\bm{W}_{qk}$ is not directly computed, and as such, it is not explicitly updated through gradient descent.
%
Nonetheless, $\bm{W}_{qk}$ is implicitly updated with a combination of the weight updates of $\bm{W}_q$ and $\bm{W}_k$ having the same form as in Proposition \ref{prop-gradients-self-attention}, see the following.
%
\input{ICML2025/math/statements/remark-self-attention-gradient-WqWk}
\subsection{Formal proofs of Theorem \ref{theo-gradient-directionality} and \ref{theo-gradients-symmetry}: the connection between objective functions, directionality, and symmetry}
\label{supp-math-directionality-symmetry}
%
In this section, we provide formal proofs of Theorems~\ref{theo-gradient-directionality} and~\ref{theo-gradients-symmetry}, which characterize the directional and symmetric structures that emerge during autoregressive and bidirectional training, respectively.
%
To this end, we introduce a series of intermediate Propositions and Lemmas, following these steps:
\begin{enumerate}[noitemsep,nolistsep]
\item We begin by showing that when \( t^* \) is used as a context token to predict other tokens, all predicted tokens contribute to the \emph{column space} of \( \bm{W}_{qk} \), while only \( t^* \) contributes to the \emph{row space}.
Conversely, when \( t^* \) is the token being predicted, only \( t^* \) contributes to the column space, whereas all context tokens contribute to the row space (Section \ref{supp-math-prop-gradient-column-rows}).
%
\item Next, we show that under reasonable assumptions about the statistical distribution of token embeddings, these structural properties imply that the expected norm of column updates exceeds that of row updates when \( t^* \) is used as context. In contrast, row updates dominate when \( t^* \) is the predicted token (Section \ref{supp-math-prop-gradient-asymmetric-growth-rows-columns}).
%
\item We then show that in autoregressive training, the expected number of times a token \( t^* \) appears as context can differ from the expected number of times it is predicted, since these two quantities are affected differently by statistical correlations between tokens.
%
In contrast, under bidirectional training, the expected counts are always equal, regardless of token correlations (Section~\ref{supp-math-prop-counting-prediction-context}).
%
\item Finally, we combine the above results to prove our main theorems:
(1) the weight updates to \( \bm{W}_{qk} \) induce \emph{column dominance} under autoregressive training (Section \ref{supp-math-lemma-tail-probabilities}-\ref{supp-math-theo-gradient-directionality}), and
(2) the weight updates to \( \bm{W}_{qk} \) induce \emph{symmetry} under bidirectional training (Section \ref{supp-math-theo-symmetry-gradients}).
\end{enumerate}
%
%
%% (1) PROPOSITION COLUMN AND ROWS
%
\subsubsection{Proof of Proposition \ref{prop-gradient-columns-rows}}
\label{supp-math-prop-gradient-column-rows}
%
First, we show that a token $t^*$ contributes differently to the updates of $\bm{W}_{qk}$ depending on whether it serves as context for predicting other tokens or is itself predicted.
%
\input{ICML2025-rebuttal/math/proofs/proof-prop-gradient-columns-rows}
%
%
%%
%% (2) ASYMMETRIC GROWTH OF COLUMNS AND ROWS
\subsubsection{Asymmetric growth of rows and columns during weight update}
\label{supp-math-prop-gradient-asymmetric-growth-rows-columns}
%
Next, we demonstrate that under reasonable assumptions about the statistical distribution of token embeddings, the expected norm of column updates exceeds that of row updates when the token \( t^* \) is used as context to predict other tokens.
%
Conversely, when $t^*$ is being predicted by other tokens, the row updates become dominant.
%
To do so, we make the following assumptions:
%
\begin{itemize}[noitemsep,nolistsep]
%
\item The token embeddings $\bm{x}_i$ are independent and identically distributed (i.i.d.) random vectors drawn from a probability distribution $\mathcal{D}$ with zero mean, $\mathbb{E}[\bm{x}_i] = \mathbf{0}$, and covariance matrix $\text{Cov}(\bm{x}_i) = \Sigma$.
%
This assumption holds at initialization for any Transformer model with learnable embeddings.
%
\item The covariance matrix is not isotropic, i.e., $\Sigma \neq \sigma^2 \mathbb{I}$,.
%
More specifically, we posit that there is partial alignment between the embeddings $\bm{x}_i$ due to the semantic and predictive relationships between tokens, which typically emerge during training.
%
\end{itemize}
%
Similar to Proposition \ref{prop-gradient-columns-rows}, the scenarios where $t^*$ is used to predict other tokens and where $t^*$ is being predicted by other tokens are complementary.
%
In the following Proposition, we focus solely on the case where $t^*$ serves as context to predict a set of tokens.
%
A formal derivation for the opposite case where $t^*$ is predicted by other tokens is provided in a subsequent Corollary.
\input{ICML2025-rebuttal/math/statements/prop-gradient-asymmetric-growth-rows-columns}
\input{ICML2025-rebuttal/math/proofs/proof-prop-gradient-asymmetric-growth-rows-columns}
%
%
Again, a complementary argument can be made about the asymmetric weight update when a set of tokens $U = \{t_1, \dots, t_N\}$ is used to predict a given token $t^*$. We formalize this in the following Corollary.
%
\input{ICML2025-rebuttal/math/statements/corollary-gradient-asymmetric-growth-rows-columns}
\input{ICML2025-rebuttal/math/proofs/proof-corollary-gradient-asymmetry-growth-rows-columns}
%
%% (3) PROPOSITION COUNTING + PROOF
\subsubsection{Expected contribution for context and prediction and related remarks}
\label{supp-math-prop-counting-prediction-context}
%%
%%
%% PROVIDE ASSUMPTIONS ON MASKING AND BIDIRECTIONAL TRAINING
%
Additionally, we show that in autoregressive training, the expected number of tokens predicted by a given token \( t^* \) can differ from the number of tokens that predict \( t^* \), with this ratio influenced by token correlations, while in bidirectional training, the two quantities are always equal regardless of such correlations, yielding a ratio of 1.
%
We formalize this in the following proposition and illustrate, in the subsequent remark, an example from autoregressive training where the expected ratio exceeds 1.
%
\input{ICML2025-rebuttal/math/statements/prop-counting-prediction-context}
\input{ICML2025-rebuttal/math/proofs/proof-prop-counting-prediction-context}
%
The following remark highlights how, during autoregressive training, the ratio of expected token counts can vary depending on the correlations and conditional dependencies among the tokens.
%
\input{ICML2025-rebuttal/math/statements/remark-ratio-autoregressive-training}
% %
% We provide some examples where tokens that can tend to appear earlier in the sequence.
% %
% As a first example, consider a position-biased mixture distribution \( P \), where tokens are independent across positions, but the probability of observing a specific token \( t^* \) depends on the position \( k \). Specifically, \( \Pr[t_k = t^*] = \alpha q_k \), where \( \alpha \in (0,1] \) and \( \{q_k\} \) is a non-increasing sequence in \( k \).
% %
% For instance, \( q_k \) may follow a linear decay \( q_k \propto (N + 1 - k) \), or an exponential decay \( q_k \propto e^{-\lambda(k - 1)} \), such that \( q_1 \ge q_2 \ge \cdots \ge q_N \).
% %
% In such cases, any token whose probability decreases with position \( k \) will exhibit an expected count ratio greater than 1.
% %
% Next, consider a distribution \( P \) defined by a Zipf-tail model governed by a two-state hidden Markov chain. Assume that in state 0, the token \( t^* \) is emitted with a small probability \( \theta \), while in state 1, it is emitted with a higher probability \( \Theta > \theta \), and the chain remains in state 1 with probability \( \varphi \).
% %
% Additionally, suppose that bursts are more likely to begin near the start of the sequence. For example, the initial state is set to 1 with probability \( p_{\mathrm{init}} \propto k^{-\beta} \), decaying with \( k \) according to a Zipf law.
%% (5.0) LONG-TAIL DISTRIBUTION AND HIGHER NUMBER OF OUTLIERS
\subsubsection{Tail probabilities of distributions with equal mean and different variances}
\label{supp-math-lemma-tail-probabilities}
%
Before proving column dominance in the next Theorem, we first establish that a probability distribution with higher variance is more likely to produce large values.
%
Specifically, when two distributions share the same mean but differ in variance, the one with higher variance has a greater probability of generating samples that exceed a given threshold.
%
We formalize this in the following lemma.
%
\input{ICML2025-rebuttal/math/statements/lemma-probability_same_mean_different_variance}
\input{ICML2025-rebuttal/math/proofs/proof_lemma-probability_same_mean_different_variance}
%
%% (5) THEOREM DIRECTIONALITY
\subsubsection{Proof of Theorem \ref{theo-gradient-directionality}}
\label{supp-math-theo-gradient-directionality}
%
Here, we show that the weight updates of $\bm{W}_{qk}$ induce column dominance during autoregressive training, leading to the emergence of directionality.
%
To do so, we make the following assumptions:
%
\begin{itemize}[noitemsep,nolistsep]
%
\item We adopt the same assumptions stated in Propositions~\ref{prop-gradient-asymmetric-growth-rows-columns} and~\ref{prop-counting-prediction-context}.
%
\item At initialization, the entries of \( \bm{W}_{qk} \) are drawn from a probability distribution \( \mathcal{P} \) with finite mean \( \mu \) and variance \( \sigma^2 \).
This is satisfied by any standard machine learning initialization scheme. This second assumption implies that each $k$-th column $\bm{w}_{\cdot,k}$ and $m$-th row $\bm{w}_{m, \cdot}$ have the same mean $\mu$ and variance $\sigma^2/\sqrt{n}$.
%
Furthermore, the squared norm of both rows and columns has equal mean $n(\sigma^2 + \mu^2)$ and variance $n(2\sigma^4 + 4\mu^2)$.
%
\end{itemize}
%
\input{ICML2025-rebuttal/math/proofs/proof-theo-gradient-directionality}
%
%% THEOREM SYMMETRY
\subsubsection{Proof of Theorem \ref{theo-gradients-symmetry} and related remarks}
\label{supp-math-theo-symmetry-gradients}
%
Finally, we prove that the weight updates of $\bm{W}_{qk}$ induce symmetry during bidirectional training.
%
Importantly, the column dominance is present only during autoregressive training.
%
Indeed, it follows from Proposition \ref{prop-counting-prediction-context} that, during bidirectional training, the net increase of the norm of the columns is equal to the net increase of the norm of the rows.
%
We formalize this in the following Corollary.
%
\input{ICML2025-rebuttal/math/statements/corollary-gradient-directionality}
\input{ICML2025-rebuttal/math/proofs/proof-corollary-gradient-directionality}
%
%
We now prove Theorem~\ref{theo-gradients-symmetry}, showing that the bidirectional training objective induces approximate symmetry in the gradient contributions of each token pair \( (i, j) \).
%
Specifically, we assume that predicting token \( t_i \) from \( t_j \) is correlated with predicting \( t_j \) from \( t_i \), that is, the terms \( \beta_{ij} \) and \( \beta_{ji} \) from Proposition~\ref{prop-gradients-self-attention} are correlated, leading to similar contributions in both directions.
%
\input{ICML2025-rebuttal/math/proofs/proof-theo-gradient-symmetry}
%
Encoder-only models are typically not trained to predict every token in a sequence, but rather a random subset of tokens, and the model can attend to tokens bidirectionally.
%
This is usually called Masked Language Modeling (MLM) \citep{devlinBERTPretrainingDeep2019, liuRoBERTaRobustlyOptimized2019, lanALBERTLiteBERT2020, warnerSmarterBetterFaster2024}.
%
Therefore, only a subset of terms in the double summation of Equation \eqref{eq:gradients-symmetric-terms} has the symmetric properties described above.
%
We generalize the proof to this case in the following.
%
\input{ICML2025/math/statements/remark-gradient-symmetry-MLM}
%
Let $|M| = pN$ with $0<p<1$ being the percentage of tokens to be predicted during bidirectional training.
%
The total number of pairs in the second term of Equation \eqref{eq:factorization-summation} is given by a binomial coefficient, thus the total number of elements in the summation is $pN(pN - 1)$.
%
The total number of elements in the third term is instead the product $pN(N - pN)$.
%
Therefore, the percentage of symmetric weight updates from the second term over the total number of updates in the third term is given by
%
\begin{equation}
\frac{pN(pN - 1)}{pN(N - pN)} \approx \, \frac{pN}{(N - pN)}\,,
\end{equation}
%
in the limit of large $N$.
%
In practice, $p$ is set to be around 15\%-30\% \citep{devlinBERTPretrainingDeep2019, liuRoBERTaRobustlyOptimized2019,lanALBERTLiteBERT2020,warnerSmarterBetterFaster2024}, leading to $ \approx 25\%$ of symmetric weight updates on average.
%
%% PROOF PROPERTIES OF SYMMETRY SCORE
\subsection{Properties of the symmetry score in Definition \ref{def-symmetry-score} and related proofs}
\label{supp-math-symmetry-score}
%
The score $s$ we introduce in Section~\ref{sec-results-pretrained-models} indicates the \emph{degree} of symmetry of a matrix $\bm{M}$ by quantifying the contribution to the Frobenious norm of its symmetric and skew-symmetric parts.
%
In particular, $s$ equals 1 and -1 for a fully symmetric and skew-symmetric matrix.
%
Accordingly, positive (negative) values of $s$ indicate the presence of symmetric (skew-symmetric) structures.
%
Here, we provide a proof for these properties.
%
First, we show that the Frobenious norm of any square matrix $\bm{M}$ can be decomposed in the sum of the Frobenious norm of its symmetric and skew-symmetric components, as in the following Lemma,
%
\input{ICML2025/math/statements/lemma-decomposition-norm-symmetric-skew-symmetric}
\input{ICML2025/math/proofs/proof-lemma-decomposition-norm-symmetric-skew-symmetric}
%
Next, we formulate the properties of the symmetry score as follows,
%
\input{ICML2025/math/statements/prop-score-symmetry-properties}
\input{ICML2025/math/proofs/proof-prop-symmetry-score-properties}
%
%% PROOF PROPERTIES OF DIRECTIONALITY SCORE
\subsection{Properties of the directionality score in Definition \ref{def-directionality-score} and related proofs}
\label{supp-math-directionality-score}
%
The score $d$ we introduce in Section~\ref{sec-results-pretrained-models} quantifies the directional bias of a square matrix $\bm{M}$ by comparing the total norm of the "outliers" rows and columns, that is, that are higher than $\gamma$ times the standard deviations of the norms.
%
A directionality score $d$ of 1 indicates the presence of rows with high ``outlier'' norms and the absence of outliers in the distribution of the column norms.
%
The opposite is true for a directionality score $d$ of -1.
%
Accordingly, positive (negative) values of $d$ indicate the presence of row (column) dominance in the matrix.
%
Here, we provide a proof for these properties.
%
\input{ICML2025/math/statements/prop-score-directionality-properties}
\input{ICML2025/math/proofs/proof-prop-directionality-score-properties}
%
\begin{definition}
\label{def-transformer-model}
%
(\textbf{\emph{Transformer architecture}}).
%
%
Let $\bm{U} \in \mathbb{R}^{N,V}$ be a matrix representing the sequence of $N$ one-hot encoded tokens of dimension $V$.
%
A Transformer architecture consists of $L$ stacked attention blocks, each one composed of a self-attention layer and a feedforward layer, as follows,
%
\begin{equation}
\label{eq:def:transformer-model}
\begin{cases}
%
\bm{X}_0(\bm{U}) = \bm{U} \bm{W_e} + \bm{W}_p \\[5pt]
\vspace{5pt}
\begin{aligned}
\Bigg\{
\begin{aligned}
&\hat{\bm{X}}_l = \bm{X}_{l-1} + a_l(\bm{X}_{l-1}; \bm{W}_q^l, \bm{W}_k^l, \bm{W}_v^l) \\[2pt]
&\bm{X}_l = \hat{\bm{X}}_l + m_l(\hat{\bm{X}}_l ; \bm{W}_1^l, \bm{W}_2^l)
\end{aligned} & \,\,\, \forall l \in [1, L]
\end{aligned} \\
\sigma\big(\bm{Z}\big) = \sigma\big(\bm{X}_L \bm{W_u}\big) \,,
%
\end{cases}
\end{equation}
%
where $\bm{W_e}\in \mathbb{R}^{V,d}$ represents the linear transformation from the vocabulary space to the embedding space of dimension $d$, $\bm{W_p}\in \mathbb{R}^{V,d}$ represents the positional encoding, $\bm{X}_0\in \mathbb{R}^{N,d}$ is the initial embedding of the sequence,
$a_l(\cdot)$ is a self-attention function given by
%
\begin{equation}
a(\bm{X}_{l-1}) = \bm{A}^l(\bm{X}_{l-1})\bm{V}^l(\bm{X}_{l-1})
\end{equation}
%
where the matrix of attention scores $\bm{A}^l(\bm{X}_{l-1})$ is given by
%
\begin{equation}
\label{eq:results:self-attention-transformers}
\begin{cases}
%
\bm{Q}^l(\bm{X}_{l-1}) = \bm{X}_{l-1}\bm{W}^l_q \\
\bm{K}^l(\bm{X}_{l-1}) = \bm{X}_{l-1}\bm{W}^l_k \\
\bm{A}^l(\bm{X}_{l-1}) = \sigma\left(\frac{1}{\sqrt{d}}\,\bm{Q}^l {\bm{K}^l}^T\right) \,,\\
%
\end{cases}
\end{equation}
%
where $1/\sqrt{d}$ is a constant normalization factor, and $\bm{W}_q^l\in \mathbb{R}^{d,d}$ and $\bm{W}^l_k\in \mathbb{R}^{d,d}$ represent linear transformations within the embedding space,
$\bm{V}^l(\bm{X}_{l-1}) = \bm{X}_{l-1}\bm{W}^l_v$ represents a linear transformation within the embedding space, $m_l(\cdot)$ is a position-wise feedforward layer with hidden dimension $d_f$ and learnable matrices $\bm{W}^l_1 \in \mathbf{R}^{d,d_f}$ and $\bm{W}^l_2\in \mathbf{R}^{d_f,d}$, $\bm{W_u}\in \mathbb{R}^{d, V}$ represents the linear transformation from the embedding space back to the vocabulary space, $\sigma(\cdot)$ is the row-wise softmax function, and $\sigma\big(\bm{Z}\big)\in \mathbb{R}^{N,V}$ is the estimated probability distribution over the vocabulary.
%
We omit layer normalization and biases for simplicity (see also \citep{elhageMathematicalFrameworkTransformer2021}).
%
\end{definition}
%
\begin{definition}
\label{def-bilinear-form}
(\textbf{\emph{Bilinear form}}).
%
A bilinear form on a vector space $V$ over a field $F$ is a map $M: V \times V \to F$ that is linear in each argument separately, that is,
%
\begin{equation}
\begin{split}
& M(a\bm{x} + b\bm{y}, \bm{z}) = aM(\bm{x}, \bm{z}) + bM(\bm{y}, \bm{z}) \\
& M( \bm{x}, a\bm{y} + b\bm{z}) = aM(\bm{x}, \bm{y}) + bM(\bm{x}, \bm{z}) \,,\\
\end{split}
\end{equation}
%%
for all \(\bm{x}, \bm{y}, \bm{z} \in V \) and \( a, b \in F \).
%
Let $\{\bm{e}_1,\dots,\bm{e}_d\}$ be a basis for the vector space $V$.
%
The matrix $\bm{M}$ such that $[\bm{M}]_{ij} = M(\bm{e}_i,\bm{e}_j)$ is the matrix of the bilinear form on this basis, and it follows
%
\begin{equation}
M(\bm{x},\bm{y}) = \bm{x}^\top\bm{M}\bm{y}
\end{equation}
%
\end{definition}
\begin{definition}
\label{def-objective-functions}
%
(\textbf{\emph{Autoregressive and bidirectional objectives}}).
%
Let $U = \{t_1, \dots, t_N\}$ a sequence of tokens.
%
The joint probability of $U$ is factorized autoregressively as follows,
%
\begin{equation}
\Pr[U] = \Pr[t_1,\dots,t_N] = \Pi_{i=1}^N \Pr[t_i|t_1,\dots,t_{i-1}]\,.
\end{equation}
%
During autoregressive training, a model with a set of parameters $\mathcal{W}$ is optimized to minimize the following negative log-likelihood
%
\begin{equation}
\mathcal{L}(U; \mathcal{W}) = - \sum_{i=1}^{N} \log p_{\mathcal{W}}(t_i \,|\, \{t_j\} : j < i)\,,
\end{equation}
%
where the conditional probabilities are modeled with learnable parameters $\mathcal{W}$.
%
% The join probability of $U$ can also be factorized bidirectionally as follows,
% %
% \begin{equation}
% \Pr[U] = \Pr[t_1,\dots,t_N] = \Pi_{i=1}^N \Pr[t_i|t_1,\dots,t_{i-1}, t_{i+1}, \dots, t_N]\,.
% \end{equation}
%
During bidirectional training, a model with a set of parameters $\mathcal{W}$ is optimized to minimize the following negative log-likelihood
%
\begin{equation}
\mathcal{L}(U; \mathcal{W}) = - \sum_{i=1}^{N} \log p_{\mathcal{W}}(t_i \,|\, \{t_j\} : j \neq i)\,.
\end{equation}
%
In practice, only a subset $M\in [1,\dots,N]$ of the tokens are predicted as in Masked Language Modelling (MLM) \citep[see ][]{devlinBERTPretrainingDeep2019,warnerSmarterBetterFaster2024}, leading to the following negative log-likelihood
%
\begin{equation}
\mathcal{L}(U; \mathcal{W}) = - \sum_{i\in M} \log p_{\mathcal{W}}(t_i \,|\, \{t_j\} : j \notin M)\,.
\end{equation}
%
\end{definition}
\begin{remark}
%
Each element of the sum in Equation \eqref{eq-proof-self-attention-bilinear-form} implicitly defines an operator in the subspace spanned by $\bm{X} = \{\bm{x}_0^\top, \bm{x}_1^\top, \dots, \bm{x}_N^\top\}$ given the transformed embedding space defined by $\bm{W}_{qk}$,
%
\begin{equation}
\sum_j \alpha_{ij}(\bm{W}_{qk})\, \bm{x}_j = \sum_j \bm{x}_i^\top \bm{W}_{qk} \bm{x}_j\,\bm{x}_j = \sum_j P_{\bm{W}_{qk}}(\bm{x}_i ,\bm{x}_j)\,,
\end{equation}
%
where $P_{\bm{W}_{qk}}(\,\cdot\,, \bm{x}_j)$ are operators over the subset $\bm{X}$.
%
The set $\bm{X}$ is in general linearly dependent, and the number of tokens $N$ in the sequence differs from the embedding space dimension $d$.
%
Furthermore, $\bm{W}_{qk}$ represents a general bilinear form (see Definition \ref{def-bilinear-form}), which may not satisfy all the defining axioms of a formal inner product — namely, linearity, conjugate symmetry, and positive definiteness.
%
Finally, the operators $P_{\bm{W}_{qk}}(\cdot \bm{x}_j)$ are not formal projection operators since they are not nilpotent ($P_{\bm{W}_{qk}}(\,\cdot\, \bm{x}_j) \circ P_{\bm{W}_{qk}}(\,\cdot\, \bm{x}_j) \neq P_{\bm{W}_{qk}}(\,\cdot\, \bm{x}_j)$).
%
Nonetheless, the bilinear map $\bm{W}_{qk}: \mathbb{R}^d \times \mathbb{R}^d \rightarrow \mathbb{R}$ still associates any pair of vectors with a scalar value quantifying their alignment as determined by the geometric relations encoded in $\bm{W}_{qk}$.
%.
% Nevertheless, the sum in Equation \ref{eq:math:linear-decomposition-subspace-tokens} represents a decomposition of $\bm{x}_i$ on the subspace $\text{span}\{\bm{X}\}$ where each coefficient $\alpha_{ij}(\bm{W}_{qk})$ quantifies the aligment of $\bm{x}_i$ with the corresponding element of $\bm{X}$ following the bilinear form $\bm{W}_{qk}$.
% %%
Therefore, self-attention computes a generalized decomposition of $\bm{x}_i$ on $\text{Conv}(\bm{X})$ in the transformed embedding space defined by $\bm{W}_{qk}$.
%
A convex combination ensures that the resulting vector remains within the region enclosed by the basis vectors $\bm{X}$.
%
\end{remark}
\begin{remark}
%
Following Definition \ref{def-transformer-model}, multi-head attention consists of parallelizing the self-attention operation across $H$ different heads with an embedding space $d_h < d$,
%
\begin{equation}
\hat{\bm{X}}(\bm{X}) = \bm{X} + \text{concat}\big(\bm{A}_1\bm{V}_1, \bm{A}_2\bm{V}_2, \dots, \bm{A}_h\bm{V}_h\big)\bm{W}_o
\end{equation}
%
where $\bm{A}_h = \sigma(d^{-1/2}\,\,\bm{X} \,\bm{W}_{q,h}\,\bm{W}^\top_{k,h}\,\bm{X}^T)$ is the self-attention of the $h$-th head, $\bm{W}_{q,h}\in \mathbb{R}^{d,d_h}$, $\bm{W}_{k,h}\in \mathbb{R}^{d,d_h}$ and $\bm{W}_{v,h}\in \mathbb{R}^{d,d_h}$ are the query, key, and value matrices of the $h$-th attention head, respectively, and $\bm{W}_o \in \mathbb{R}^{d,d}$ is a linear transformation \citep{vaswaniAttentionAllYou2017}.
%
Operationally, the self-attention computation is performed in parallel by factorizing the $\bm{W}_q$ and $\bm{W}_k$ matrices into $H$ rectangular blocks, as follows,
%
\begin{equation}
\begin{split}
& \bm{W}_q = \big[\bm{W}_{q,1} \big| \bm{W}_{q,2} \big| \dots \big| \bm{W}_{q,H}\big] \\
& \bm{W}_k = \big[\bm{W}_{k,1} \big| \bm{W}_{k,2} \big| \dots \big| \bm{W}_{k,H}\big] \,,
\end{split}
\end{equation}
%
and performing the matrix multiplication $\bm{W}_{q,h}\bm{W}_{k,h}^\top$ per every $h$-th head independently in one step.
%
It follows that the full $\bm{W}_{qk}$ matrix is given by the sum of the bilinear forms $\bm{W}_{qk,h}$ of every head, as follows,
%
\begin{equation}
\bm{W}_{qk} = \bm{W}_q\bm{W}_k^\top = \sum_h \bm{W}_{q,h}\bm{W}^\top_{k,h} = \sum_h \bm{W}_{qk,h}
\end{equation}
%
where each $\bm{W}_{qk,h} \in \mathbb{R}^{d,d}$ is a square matrix with $\text{rank}(\bm{W}_{qk,h}) \leq d_h$.
%
Therefore, each head perform independent projections onto $\text{Conv}(\bm{X})$ that are then summed together, as follows,
%
\begin{equation}
\sum_j\hat{\alpha}_{ij}\, \bm{x}_j = \sum_j \bm{x}_i^\top \bm{W}_{qk} \bm{x}_j\,\bm{x}_j =
\sum_j \bm{x}_i^\top \Big(\sum_h\bm{W}_{qk,h}\Big) \bm{x}_j\,\bm{x}_j = \sum_j \sum_h \langle \bm{x}_i, \bm{W}_{qk,h}\bm{x}_j \rangle \,\bm{x}_j \,,
\end{equation}
%
thus performing the same operations as in Equation \eqref{eq-proof-self-attention-bilinear-form}.
%
\end{remark}
%
Let \( \mathcal{V} = \{t_1, \dots, t_V\} \) be a vocabulary of tokens, and let \( D \) be a dataset of sequences \( U_u \in \mathcal{V}^N \) of length \( N \), where $D = \left\{ U_u = (x^u_1, \dots, x^u_N) \right\}_{u=1}^D \subset \mathcal{V}^N$.
%
For a single sequence \( U_u \) with joint distribution \( p(U_u) \), the chain rule implies that for any permutation \( \sigma \) of \( \{1, \dots, N\} \), the joint probability can be written as,
\begin{equation}
p(U_u) = \prod_{t=1}^{N} p\big(x^u_{\sigma(t)} \mid x^u_{\sigma(1)}, \dots, x^u_{\sigma(t-1)}\big).
\label{eq:chain}
\end{equation}
%
Given a set of conditional parents \( \mathcal{D}_i \subset \{1, \dots, N\} \setminus \{i\} \), define the **local negative log-likelihood** as:
\begin{equation}
\mathcal{L}(U; \mathcal{W}) = -\sum_{i=1}^{N} \log p_{\mathcal{W}}(x_i \mid x_{\mathcal{D}_i}),
\label{eq:localNLL}
\end{equation}
where \( \mathcal{W} \) denotes the model parameters.
%
Let \( G \) be a directed graph where each edge is \( \mathcal{D}_i \to i \). We distinguish two cases:
%
\begin{enumerate}[noitemsep,nolistsep]
\item
If \( G \) is a directed acyclic graph (DAG), then a topological ordering \( \tau \) exists such that \( \mathcal{D}_{\tau(i)} \subseteq \{\tau(1), \dots, \tau(i-1)\} \). In this case, the joint distribution $p_{\mathcal{W}}(x) = \prod_{i=1}^N p_{\mathcal{W}}(x_{\tau(i)} \mid x_{\mathcal{D}_{\tau(i)}})
$ is properly normalized and the local negative log-likelihood in \eqref{eq:localNLL} is equal to the true negative log-likelihood.
%
For example, during autoregressive training, each token is conditioned only on its predecessors,
\begin{equation}
\mathcal{L}(U; \mathcal{W}) = -\sum_{i=1}^N \log p_{\mathcal{W}}(x_i \mid x_1, \dots, x_{i-1}) = -\sum_{i=1}^N \log p_{\mathcal{W}}(x_i \mid x_{<i}).
\label{eq:ARloss}
\end{equation}
\item If \( G \) contains directed cycles, the product $\prod_{i=1}^N p_{\mathcal{W}}(x_i \mid x_{\mathcal{D}_i})$ is not guaranteed to be normalized, and thus does not define a proper joint likelihood.
%
In this case, one can instead minimize the following pseudo log-likelihood,
\begin{equation}
\hat{\mathcal{L}}(U; \mathcal{W}) = -\sum_{i=1}^{N} \log p_{\mathcal{W}}(x_i \mid x_1, \dots, x_{i-1}, x_{i+1}, \dots, x_N) = -\sum_{i=1}^{N} \log p_{\mathcal{W}}(x_i \mid x_{-i}),
\label{eq:PLL}
\end{equation}
where \( x_{-i} \) denotes all tokens in the sequence except \( x_i \).
%
Following Besag's theorem \cite{besag1975statistical}, the maximum pseudo-likelihood estimator is consistent under standard regularity conditions, and \eqref{eq:PLL} remains compatible with stochastic-gradient
optimization even though it is not the true negative log-likelihood \cite{yang2019xlnet}.
\end{enumerate}
\begin{proof}
%
Let $U = \{t_1, \dots, t_N\}$ be a sequence of $N$ tokens.
%
Let $\mathcal{L}(U ; \mathcal{W})$ be the negative log-likelihood of each token $t_i$, expressed as
%
\begin{equation}
\label{eq:results:self-supervised-pretraining}
\mathcal{L}(U;\mathcal{W}) =\sum_{i=1}^N \mathcal{L}(t_i)= - \sum_{i=1}^N \log p_{\mathcal{W}}(t_i \,|\, t_{\mathcal{D}_i}) \,,
\end{equation}
%
%
where $\mathcal{D}_i \subset [0,1,\dots, N]$ is the conditioning set, that is, the set of indices defining the set of tokens $\{t_j: j \in \mathcal{D}_i\}$ of the conditional probability distribution.
%
Let $\bm{U} = [\bm{t}_0, \bm{t}_1, \dots, \bm{t}_N]$ be the sequence of $N$ one-hot encoded tokens $\bm{t}_i \in \mathbb{R}^V$ associated with $U$, where $V$ is the dimension of the vocabulary.
%
Let $\mathcal{L}(\bm{t}_i)$ be the cross-entropy of the one-hot encoded token $\bm{t}_i$ and the estimated probability distribution $\sigma(\bm{z}_i) \in \mathbb{R}^V$, as follows,
%
\begin{equation}
\mathcal{L}(\bm{U}) = \sum_{i=1}^N \mathcal{L}(\bm{t}_i) = \sum_{i=1}^N \bm{t}_i\log(\sigma(\bm{z}_i))\,,
\end{equation}
%
where we let $\bm{z}_i$ be the prediction of the $i$-th token $\bm{t}_i$ from the representations in the last layer of a Transformer model, following Definition \ref{def-transformer-model},
%
%Let the self-attention function in Definition~\ref{def-transformer-model} be linear, that is,
%
% \begin{equation}
% \bm{A}^l(\bm{X}_{l-1}) = \,\bm{Q}^l {\bm{K}^l}^T
% \end{equation}
%
% where we neglect the scaling factor $\sqrt{d}$ for simplicity.
% %
% The self-attention function thus maps the token embeddings $\{\bm{x}_i^{l-1}\}$ from layer $l-1$ to the embeddings $\{\hat{\bm{x}}^l_i\}$ as follows,
% %
% \begin{equation}
% {\bm{\hat{x}}^{l-1}_i}^\top = {\bm{x}^{l-1}_i}^\top + \sum_{j\in C_i}
% \alpha^l_{ij}\,
% {\bm{x}^{l-1}_j}^\top \bm{W}^l_v \,,
% \end{equation}
% %
% where $C_i \subset \bm{U}$ is the subset of tokens that defines the conditional probability distribution.
%
\begin{equation}
\begin{cases}
%
{\bm{x}^0_i}^\top = \bm{t_i}^\top \bm{W_e} + \bm{W}_p \\[5pt]
\vspace{5pt}
{\bm{x}_i^l}^\top = \mathcal{F}_l(\bm{x}_i^{l-1})\quad \forall l \in [1, L] \\
\sigma\big(\bm{z}_i\big) = \sigma\big({\bm{x}_i^L}^\top \bm{W_u}\big)\,,
%
\end{cases}
\label{eq-proof-gradient-bilinear-form-transformer-model}
\end{equation}
%
where ${\bm{x}^l_i}^\top = \mathcal{F}_l(\bm{x}_i^{l-1})$ is a short notation for the self-attention and multi-layered perception transformation of the $l$-th layer,
%
\begin{equation}
\mathcal{F}_l(\bm{x}_i^{l-1}) =
\left\{
\begin{aligned}
& \hat{\bm{x}}_i^{l^{\top}} = \hat{\bm{x}}_i^{{l-1}^{\top}} + a_l(\bm{x}_i^{l-1}) \\
& \bm{x}_i^{l^{\top}} = \hat{\bm{x}}_i^{l^{\top}} + m_l(\hat{\bm{x}}^l_i)
\end{aligned}
\right.\,,
\end{equation}
%
where the self-attention function is given by
%
\begin{equation}
a_l(\bm{x}_i^{l-1}) = \sum_{j\in \mathcal{D}_i}
\alpha^l_{ij}(w^l)\,
{\bm{x}^{l-1}_j}^\top \bm{W}^l_v \,.
\end{equation}
%
Let attention coefficients $\alpha^l_{ij} \equiv \alpha^l_{ij} (w^l)$ of the $l$-th layer be parameterized with a general parameter $w^l$.
%
% We divide the derivation of the gradient of $\mathcal{L}(\bm{t}_i)$ w.r.t. the self-attention function into two steps.
%
Let the gradient of $\mathcal{L}(\bm{t}_i)$ w.r.t. $w_l$ (the parameterization of the attention scores) be factorized as follows,
%
\begin{equation}
\nabla_{w^l} \mathcal{L}(\bm{t}_i) = \frac{\partial \mathcal{L}(\bm{t}_i)}{\partial \alpha^l_{ij}} \frac{\partial \alpha^l_{ij}}{\partial w^l}\,.
\end{equation}
%
It follows that,
%
\begin{equation}
\begin{split}
\frac{\partial \mathcal{L}(\bm{t}_i)}{\partial \alpha^l_{ij}}
& = \frac{\partial \mathcal{L}(\bm{t}_i)}{\partial \bm{z}_i} \, \frac{\partial \bm{z}_i}{\partial \bm{x}^L_i} \, \frac{\partial \bm{x}^L_i}{\partial \hat{\bm{x}}^l_i}\,\frac{\partial \hat{\bm{x}}^l_i} {\partial \alpha^l_{ij}}
\\
& = (\bm{t}_i - \sigma(\bm{z}_i))^\top \bm{W}_u^\top \, \frac{\partial \bm{x}^L_i}{\partial \hat{\bm{x}}^l_i}\,{\bm{W}^l_v}^\top \sum_{j\in \mathcal{D}_i} \bm{x}^{l-1}_j \,,
\end{split}
\end{equation}
%
where the term $\partial \bm{x}_i^L / \partial \hat{\bm{x}}_i^l$ includes the set of partial derivatives that define the gradient of the representation $\bm{x}_i^L $ at the last layer w.r.t. the self-attention representation $\hat{\bm{x}}_i^l$ at the $l$-layer, as follows,
%
\begin{equation}
\frac{\partial \bm{x}_i^L}{\partial \hat{\bm{x}}_i^l} = \left(1 + \sum_{m = l}^{L-1}\mathcal{F}_l'(\bm{x}_i^m)\right) \left(1 + m_l'(\hat{\bm{x}}^l_i)\right)\,,
\end{equation}
%
where
%
\begin{equation}
\mathcal{F}_l'(\bm{x}_i^m) = \frac{\partial}{\partial \bm{x}_i^l}\mathcal{F}_l(\bm{x}_i^m) \quad ; \quad m_l'(\hat{\bm{x}}^l_i) = \frac{\partial}{\partial \bm{\hat{x}}_i^l} m_l(\hat{\bm{x}}^l_i) \,.
\end{equation}
%
Let $\bm{\delta}^l_i$ be the error at the last layer propagated to the self-attention function at the $l$-th layer, as follows,
%
\begin{equation}
{\bm{\delta}^l_i}^\top = (\bm{t}_i - \sigma(\bm{z}_i))^\top \bm{W}_u^\top \, \frac{\partial \bm{x}^L_i}{\partial \hat{\bm{x}}^l_i} \,
{\bm{W}^l_v}^\top \,,
\end{equation}
%
thus obtaining the following equation for the gradient,
%
\begin{equation}
\nabla_{w^l} \mathcal{L}(\bm{t}_i) = {\bm{\delta}^l_i}^\top \sum_{j\in \mathcal{D}_i} \bm{x}^{l-1}_j \frac{\partial \alpha^l_{ij}}{\partial w^l} \,.
\end{equation}
%
Let the attention scores be computed without the row-wise softmax operation and explicitly with the bilinear form $\bm{W}_{qk}$, such that the following expression gives the score between the $i$-th and $j$th token,
%
\begin{equation}
\alpha^l_{ij} \equiv \alpha^l_{ij}(\bm{W}_{qk}^l) = {\bm{x}^{l-1}_i}^\top \bm{W}_{qk}^l \bm{x}^{l-1}_j \,,
\end{equation}
%
from which we obtain
%
\begin{equation}
\frac{\partial \alpha^l_{ij}}{\partial \bm{W}_{qk}^l} = \bm{x}^{l-1}_i{\bm{x}^{l-1}_j}^\top = \bm{K}^{l-1}_{ij}\, ,
\end{equation}
%
where $\bm{K}^{l-1}_{ij} \in \mathbb{M}_n$ is a square rank-1 matrix given by the outer product between the $i$-th and $j$-th token from the $l-1$-th layer.
%
It follows that the total gradient of $\mathcal{L}(\bm{t}_i)$ is given by
%
\begin{equation}
\nabla_{\bm{W}_{qk}^l} \mathcal{L}(\bm{t}_i) = {\bm{\delta}^l_i}^\top \sum_{j\in \mathcal{D}_i} \bm{x}^{l-1}_j \bm{K}^{l-1}_{ij} \,,
\end{equation}
%
where we notice that, for every $j$, the term ${\bm{\delta}^l_i}^\top \bm{x}^{l-1}_j$ is a scalar quantity that we define as $\beta^l_{ij}$, thus obtaining,
%
\begin{equation}
\label{eq:math:gradient-linear-combination-rank-1-matrices}
\nabla_{\bm{W}_{qk}^l} \mathcal{L}(\bm{t}_i) = \sum_{j\in \mathcal{D}_i} \beta^l_{ij} \bm{K}^{l-1}_{ij} \,,
\end{equation}
%
and therefore,
%
\begin{equation}
\nabla_{\bm{W}_{qk}^l} \mathcal{L}(U) = \sum_i \sum_{j\in \mathcal{D}_i} \beta^l_{ij} \bm{K}^{l-1}_{ij} \quad \forall\, \bm{W}_{qk}^l, \,l \in [1,L]\,.
\end{equation}
%
We can rewrite the double summation either as,
%
\begin{equation}
\nabla_{\bm{W}_{qk}^l} \mathcal{L}(U) = \sum_i \sum_{j\in \mathcal{C}_i} \beta^l_{ij} \bm{K}^{l-1}_{ij} \quad \forall\, \bm{W}_{qk}^l, \,l \in [1,L]\,.
\end{equation}
%
where $\mathcal{C}_j \subset [0,1,\dots, N]$ is the set of indices defining the set of tokens $\{t_j\}$ used to predict $t_i$,
%
or equivalently as,
%
\begin{equation}
\nabla_{\bm{W}_{qk}^l} \mathcal{L}(U) = \sum_{i \in \mathcal{P}_j} \sum_{j} \beta^l_{ij} \bm{K}^{l-1}_{ij} \quad \forall\, \bm{W}_{qk}^l, \,l \in [1,L]\,,
\end{equation}
%
where $\mathcal{P}_j \subset [0,1,\dots, N]$ is the set of indices defining the set of tokens $\{t_i\}$ that are predicted by $t_j$, thus concluding the proof.
%
\end{proof}
\begin{remark}
%
Let $\mathcal{L}(\bm{U})$ be the negative log-likelihood of a given sequence of one-hot encoded tokens $\bm{U}$.
%
Let $\bm{W}^l_q$ and $\bm{W}^l_k$ be the query and key transformation matrices of the $l$-th layer of Transformer models, see Definition \ref{def-transformer-model}.
%
Let $\bm{W}^l_q$ and $\bm{W}^l_k$ be updated via gradient descent, that is, $\bm{W}^l_q \rightarrow \bm{W}^l_q + \eta\nabla_{\bm{W}_q^l} \mathcal{L}(\bm{U})$ and $\bm{W}^l_k \rightarrow \bm{W}^l_k + \eta\nabla_{\bm{W}_k^l} \mathcal{L}(\bm{U})$, where $\eta$ is the learning rate.
%
It follows that the matrix $\bm{W}_{qk}^l = \bm{W}^l_q{\bm{W}^l_k}^\top$ is implicitly updated following,
%
% \begin{equation}
% \bm{W}_{qk}^l = \bm{W}^l_q{\bm{W}^l_k}^\top
% \rightarrow
% \,\big(\, \bm{W}^l_q + \eta\nabla_{\bm{W}_q^l} \mathcal{L}(U) \,\big)\big(\,\bm{W}^l_k + \eta\nabla_{\bm{W}_k^l} \mathcal{L}(U) \,\big)^\top
% \end{equation}
% which we rewrite explicitly as,
%
\begin{equation}
\begin{split}
\bm{W}_{qk}^l \rightarrow &
\,\big(\, \bm{W}^l_q + \eta\nabla_{\bm{W}_q^l} \mathcal{L}(U) \,\big)\big(\,\bm{W}^l_k + \eta\nabla_{\bm{W}_k^l} \mathcal{L}(U) \,\big)^\top \\
& = \bm{W}^l_q{\bm{W}^l_k}^\top + \eta\, \big(\, \bm{W}^l_q\,{\nabla_{\bm{W}_k^l} \mathcal{L}}^\top (U) + \nabla_{\bm{W}_q^l} \mathcal{L}(U){\bm{W}^l_k}^\top \,\big) + o(\eta^2) \\
& = \bm{W}_{qk}^l + \eta \sum_i \sum_{j\in C_i} \beta_{ij}^l\big(\bm{W}_q^l{\bm{W}_q^l}^\top\, \bm{K}_{ij}^{l-1} + \bm{K}_{ij}^{l-1}\,\bm{W}_k^l{\bm{W}_k^l}^\top \big)
+ o(\eta^2) \\
& = \bm{W}_{qk}^l + \Delta \bm{W}_{qk}^l + o(\eta^2) \simeq \bm{W}_{qk}^l + \Delta \bm{W}_{qk}^l \,,
\end{split}
\end{equation}
%
assuming that the learning rate $\eta$ is small.
%
Therefore, the implicit weight update of $\bm{W}_{qk}^l$ following gradient descent is given by
%
\begin{equation}
\Delta \bm{W}_{qk}^l = \eta \sum_i \sum_{j\in C_i} \beta_{ij}^l
\left[\big(\bm{W}_q^l{\bm{W}_q^l}^\top\, \bm{K}_{ij}^{l-1} \big)
+
\big(\bm{K}_{ij}^{l-1}\,\bm{W}_k^l{\bm{W}_k^l}^\top \big)\right]
\propto \sum_i \sum_{j\in C_i} \beta_{ij}^l\, \bm{K}_{ij}^{l-1} \,,
\end{equation}
%
where both $\bm{W}_q^l{\bm{W}_q^l}^\top\, \bm{K}_{ij}^{l-1}$ and $\bm{K}_{ij}^{l-1}\,\bm{W}_k^l{\bm{W}_k^l}^\top$ are rank-1 matrices,
%
\begin{equation}
\begin{split}
& \bm{W}_q^l{\bm{W}_q^l}^\top\, \bm{K}_{ij}^{l-1} = \Big( \bm{W}_q^l{\bm{W}_q^l} \,\bm{x}^{l-1}_i\Big){\bm{x}^{l-1}_j}^\top = \bm{\bar{x}}^{l-1}_i {\bm{x}^{l-1}_j}^\top \\
& \bm{K}_{ij}^{l-1}\,\bm{W}_k^l{\bm{W}_k^l}^\top = \bm{x}^{l-1}_i\Big({\bm{x}^{l-1}_j}^\top\,\bm{W}_k^l{\bm{W}_k^l}^\top \Big) = \bm{x}^{l-1}_i \bm{\bar{x}}_j^{{l-1}^\top} \,.
\end{split}
\end{equation}
%
%
\end{remark}
\begin{proof}
%
Let $U = \{t_1, \dots, t_N\}$ be a sequence of tokens with the embedding of every $i$-th token be given by $\bm{x}_i \in \mathbb{R}^d$.
%
Proposition \ref{prop-gradients-self-attention} shows that the implicit weight update can be decomposed with two equivalent regrouping of the double summations, as follows,
%
\begin{equation}
\Delta \bm{W}_{qk} = \sum_{(i,j) \in U} \beta_{ij}\bm{x}_i\bm{x}_j^T
= \sum_j \left(\sum_{i\in \mathcal{P}_j} \beta_{ij}\bm{x}_i\right)\bm{x}^\top_j
= \sum_i \bm{x}_i\left(\sum_{i\in \mathcal{C}_i} \beta_{ij}\bm{x}^\top_j\right)\,,
\end{equation}
%
where $\mathcal{P}_{i} \subset U$ is the set of tokens predicted by a given token $t_i$, while $\mathcal{C}_{j} \subset U$ is the set of tokens that predict a given token $t_j$.
%
%
We neglect any constant of proportionality - such as a learning rate - for simplicity, and we do not assume any specific structure on $\mathcal{P}_i$ and $\mathcal{C}_j$ (autoregressive training, bidirectional training, or others).
%
First, the contribution of $t^* \in U$ to the weight update when $t^*$ is used as context to predict a set of tokens $\mathcal{P}_{t^*} \subset U$ is
%%
\begin{equation}
\Delta \bm{W}_{qk} \bigg|_{t_j = t^*}
= \left(\sum_{i\in \mathcal{P}_{t^*}} \beta_{i*}\bm{x}_i\right)\bm{x}^\top_{t^*}\,.
\end{equation}
%
The associated weight update of the $k$-th column is then given by
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} \bigg|_{t_j = t^*} =
\left(\sum_{i\in \mathcal{P}_{t^*}} \beta_{i*}\bm{x}_i\right)[\bm{x}_{t^*}]_k
=
[\bm{x}_{t^*}]_k \left(\sum_{i \in \mathcal{P}_{t^*}}\beta_{i*}\bm{x}_i\right)\,,
\end{equation}
%
while the update of the $m$-th row is
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} \bigg|_{t_j = t^*} = \left(\sum_{i \in \mathcal{P}_{t^*}}\beta_{i*}[\bm{x}_i]_m\right)\bm{x}_{t^*} \,.
\end{equation}
%
A complementary argument can be made when predicting a given token $t^*$ from a set of tokens $\mathcal{C}_{t^*}$.
%
Indeed, the contribution to the weight update when $t^*$ is predicted by set of tokens $\mathcal{C}_{t^*} \subset U$ is given by
%
\begin{equation}
\Delta \bm{W}_{qk} \bigg|_{t_i = t^*}
= \left(\sum_{j\in \mathcal{C}_{t^*}} \beta_{*j}\bm{x}_{t^*}\right)\bm{x}_j^\top\,.
\end{equation}
%
The associated weight update of the $k$-th column is then given by
%
\begin{equation}
\Delta \bm{w}_{k, \cdot} \bigg|_{t_i = t^*} = \left(\sum_{j \in \mathcal{C}_{t^*}}\beta_{*j}[\bm{x}_j]_k\right)\bm{x}_{t^*} \,,
\end{equation}
%
while the update of the $m$-th row is
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} \bigg|_{t_i = t^*} = \left(\sum_{j \in \mathcal{C}_{t^*}}\beta_{*j}\bm{x}_j^\top\right) [\bm{x}_{t^*}]_m
= [\bm{x}_{t^*}]_m \left(\sum_{j \in \mathcal{C}_{t^*}}\beta_{*j}\bm{x}_j^\top\right)\,.
\end{equation}
%
Therefore, the weight update of each column (row) when a set of tokens predicts $t^*$ is equivalent to the weight update of each row (column) when $t^*$ is used to predict a set of tokens.
%
This concludes the proof.
%
\end{proof}
\begin{proposition}
\label{prop-gradient-asymmetric-growth-rows-columns}
%
(\textbf{Asymmetric growth of columns and rows for context}).
%
Let $U = \{t_1, \dots, t_N\}$ a sequence of tokens, and let $t^*$ be a token representing the context of every token $t_i \in U$.
%
Let $\{\bm{x}_1 \dots, \bm{x}_N\}$ be the token embedding associated with $U$ such that each $\bm{x}_i \sim \mathcal{D}$ is a i.i.d. random vector drawn from a probability distribution $\mathcal{D}$ with zero mean and non-isotropic covariance $\text{Cov}(\bm{x}_i) = \Sigma$.
%
Let $\Delta \bm{W}_{qk}$ be the weight update from Proposition \ref{prop-gradients-self-attention}.
%
Then the squared norm of the $m$-th row and $k$-th column of $\bm{W}_{qk}$ satisfies
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} > 1 \quad \forall k \in \{1,\dots,d\}, \forall m \,\,\text{s.t.} \,\, \Sigma_{m,m} < \frac{\text{Tr}(\Sigma)}{d}
\end{equation}
%
\end{proposition}
\begin{proof}
%
%
Let $U = [t_0, \dots, t_N]$ be a sequence of tokens, and let $t^*$ be the context of every token $t_i \in U$.
%
Let $\{\bm{x}_i\}$ be the set of token embeddings associated with $U$ and let $\bm{y}$ be the token embedding of $t^*$.
%
It follows from Proposition \ref{prop-gradient-columns-rows} that the squared norm of the weight update of the $k$-th column is given by
%
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} = \sum_{i=1}^N \beta_{i*}\bm{x}_i[\bm{y}]_k
%
\,\,\Rightarrow \,\,
%
||\Delta \bm{w}_{\cdot, k}||^2 = [\bm{y}]^2_k || \sum_{i=1}^N \beta_{i*}\bm{x}_i||^2 \,,
\end{equation}
%
while the squared norm of the weight update of the $m$-th row by
%
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} = \sum_{i=1}^N \beta_{i*}[\bm{x}_i]_m\bm{y}
%
\,\,\Rightarrow \,\,
%
||\Delta \bm{w}_{m, \cdot}||^2 = \big( \sum_{i=1}^N\beta_{i*}[\bm{x}_i]_m\big)^2 || \bm{y}||^2 \,.
\end{equation}
%
Let $\{\bm{x}_i\} \sim \mathcal{D}$ be a set of i.i.d. random vectors from a probability distribution $\mathcal{D}$ such that $\mathbb{E}_{\mathcal{D}}[\bm{x}_i] = 0$ and $\text{Cov}(\bm{x}_i) = \Sigma \in \mathbb{R}^{d,d}$.
%
Therefore, each $k$-th coordinate $[\bm{x}_i]_k$ is such that $\mathbb{E}_{\mathcal{D}}[[\bm{x}_i]_k] = 0$ and $\text{Var}([\bm{x}_i]_k) = \Sigma_{k,k}$.
%
We also assume that $\bm{y}$ is statistically independent from $\bm{x}_i \, \forall i$ with $\mathbb{E}[\bm{y}] = 0$ and covariance $\text{Cov}(\bm{y}) = \sigma_y^2\mathbb{I}$.
%
It follows that the expected value of $||\Delta \bm{w}_{\cdot, k}||^2$ over $\mathcal{D}$ is given by
%
\begin{equation}
\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]
=
\mathbb{E}_{\mathcal{D}}\left[ [\bm{y}]^2_k || \sum_{i=1}^N \beta_{i*}\bm{x}_i||^2 \right]
=
\mathbb{E}_{\mathcal{D}}\left[ [\bm{y}]^2_k \right] \mathbb{E}_{\mathcal{D}}\left[ || \sum_{i=1}^N \beta_{i*}\bm{x}_i||^2 \right] \,.
\end{equation}
%
Given the statistical independence between the entries of $\bm{x}_i$, the second term is equal to
%
\begin{equation}
\begin{split}
\mathbb{E}_{\mathcal{D}}\left[ || \sum_{i=1}^N \beta_{i*}\bm{x}_i||^2 \right]
& = \sum_{i,i'}\mathbb{E}_{\mathcal{D}}\left[ \beta_{i*}\beta_{i'*}\bm{x}^\top_{i}\bm{x}_{i'}\right]
\\
&= \sum_{i=1}^N\beta^2_{i*}\mathbb{E}_{\mathcal{D}}\left[\bm{x}^\top_{i}\bm{x}_{i}\right] +
\sum_{i'\neq i}\beta_{i*}\beta_{i'*}\mathbb{E}_{\mathcal{D}}\left[\bm{x}^\top_{i}\bm{x}_{i'}\right]
\\
& =
\sum_{i=1}^N\beta^2_{i*}\text{Tr}\big(\mathbb{E}_{\mathcal{D}}\left[\bm{x}_{i}\bm{x}_{i}^\top\right]\big) +
\sum_{i'\neq i}\beta_{i*}\beta_{i'*}\mathbb{E}_{\mathcal{D}}\left[\bm{x}^\top_{i}\bm{x}_{i'}\right]
\\
& = \text{Tr}(\Sigma) \sum_{i=1}^N\beta^2_{i*} \,,
\end{split}
\end{equation}
%
and therefore
%
\begin{equation}
\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right] = \Gamma_{k,k}\text{Tr} (\Sigma) \sum_{i=1}^N\beta^2_{i*} \,.
\end{equation}
%
Similarly, the expected value of $||\Delta \bm{w}_{m, \cdot}||^2$ is given by
%
\begin{equation}
\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]
=
\mathbb{E}_{\mathcal{D}}\left[\big(\sum_{i=1}^N\beta_{i*}[\bm{x}_i]_m\big)^2 || \bm{y}||^2
\right]
=
\mathbb{E}_{\mathcal{D}}\left[ \big(\sum_{i=1}^N\beta_{i*}[\bm{x}_i]_m\big)^2 \right] \mathbb{E}_{\mathcal{D}}\left[ || \bm{y}||^2\right]\,.
\end{equation}
%
The first term can be decomposed as
%
\begin{equation}
\mathbb{E}_{\mathcal{D}}\left[ \big(\sum_{i=1}^N\beta_{i*}[\bm{x}_i]_m\big)^2 \right] =
\sum_{i=1}^N \beta^2_{i*}\mathbb{E}_{\mathcal{D}}\left[[\bm{x}_i]_m^2 \right] + \sum_{i' \neq i}\beta_{i*}\beta_{i'*}\mathbb{E}_{\mathcal{D}}\left[ [\bm{x}_{i'}]_m [\bm{x}_{i}]_m \right] = \Sigma_{m,m}\sum_{i=1}^N\beta^2_{i*} \,,
\end{equation}
%
and therefore
%
\begin{equation}
\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right] = \Sigma_{m,m}\text{Tr}(\Gamma) \sum_{i=1}^N\beta^2_{i*}\,.
\end{equation}
%
The ratio of the expected value of these squared norms is
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} = \frac{ \Gamma_{k,k}\text{Tr}(\Sigma)}{ \Sigma_{m,m}\text{Tr}(\Gamma)} = \frac{\Gamma_{k,k}}{\text{Tr}(\Gamma)}\frac{\text{Tr}(\Sigma)}{\Sigma_{m,m}} = \frac{1}{d}\frac{\text{Tr}(\Sigma)}{\Sigma_{m,m}} \,.
\end{equation}
%
We assume a non-isotropic covariance structure in $\Sigma$, that is, the average variance per dimension is lower than the total variance across all dimensions.
%
This implies $\text{Tr}(\Sigma) > d\,\Sigma_{m,m}$ for some $m$.
%
It follows that,
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} = \frac{1}{d}\frac{\text{Tr}(\Sigma)}{\Sigma_{m,m}} > \frac{1}{d}\frac{d \,\Sigma_{m,m}}{\Sigma_{m,m}} = 1 \quad \forall m \,\,\text{s.t.}\,\,\Sigma_{m,m} < \frac{\text{Tr}(\Sigma)}{d}\,,
\end{equation}
%
thus concluding the proof.
%
\end{proof}
\begin{corollary}
\label{corollary-gradient-asymmetric-growth-rows-columns}
%
(\textbf{Asymmetric growth of columns and rows for prediction}).
%
%
Let $U = \{t_1, \dots, t_N\}$ be a sequence of tokens, and let every token $t_i \in U$ be the context of a given token $t^*$.
%
Let the associated embedding be drawn i.i.d. from a probability distribution $\mathcal{D}$ as in Proposition \ref{prop-gradient-asymmetric-growth-rows-columns}.
%
Let $\Delta \bm{W}_{qk}$ be the weight update from Proposition \ref{prop-gradients-self-attention}.
%
Then the squared norm of the $m$-th row and $k$-th column of $\bm{W}_{qk}$ satisfies
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} < 1 \quad \forall m \in \{1,\dots,d\}, \forall k \,\,\text{s.t.} \,\, \Sigma_{k,k} < \frac{\text{Tr}(\Sigma)}{d}
\end{equation}
%
\end{corollary}
\begin{proof}
%
Let $U = [t_0, \dots, t_N]$ a sequence of tokens, and let $t^*$ be predicted by every token $t_i \in U$.
%
Let $\{\bm{x}_i\}$ be the set of token embeddings associated with $U$ and let $\bm{y}$ be the token embedding of $t^*$.
%
It follows from Proposition \ref{prop-gradient-columns-rows} that the squared norm of the weight update of the $k$-th column is given by
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} = \sum_{j=1}^N \beta_{*j}[\bm{x}_j]_k\bm{y}
%
\,\,\Rightarrow \,\,
%
||\Delta \bm{w}_{\cdot, k}||^2 = \big(\sum_{i=1}^N \beta_{*j}[\bm{x}_j]_k\big)^2 || \bm{y}||^2 \,,
\end{equation}
%
while the weight update of the $m$-th row follows
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} = \sum_{j=1} ^N \beta_{*j} \bm{x}_j[\bm{y}]_m
\,\,\Rightarrow \,\,
||\Delta \bm{w}_{m, \cdot}||^2 = [\bm{y}]^2_m || \sum_{i=1}^N \beta_{*j} \bm{x}_j||^2 \,.
\end{equation}
%
Therefore, following the same arguments as in Proposition \ref{prop-gradient-asymmetric-growth-rows-columns}, the ratio of the expected value of the square norms is given by
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} = \frac{ \Sigma_{k,k}\text{Tr}(\Gamma)}{ \Gamma_{m,m}\text{Tr}(\Sigma)} = d\frac{\Sigma_{k,k}}{\text{Tr}(\Sigma)} < \frac{d\,\Sigma_{k,k}}{d\,\Sigma_{k,k}} = 1 \quad \forall k \,\,\text{s.t.}\,\,\Sigma_{k,k} < \frac{\text{Tr}(\Sigma)}{d} \,,
\end{equation}
%
thus concluding the proof.
%
%
\end{proof}
\begin{proposition}
\label{prop-counting-prediction-context}
%
(\textbf{Ratio of expected counts during autoregressive and bidirectional training}).
%
Let $\mathcal{V} = \{t_0, \dots, t_V\}$ be a set of tokens.
%
Let $\mathcal{U}$ be the sample space of all possible sequences of $N$ tokens, and let the sequence $U = \{t_1, \dots, t_N\} \in\mathcal{U}$ be a random variable with probability distribution $P(U)$ defined over $\mathcal{U}$.
%
Let $\Pr[t_j = t^*]$ be the probability that the token at index $j$ in a sequence $U \in \mathcal{U}$ is given by $t^* \in \mathcal{V}$.
%
Let $\mathbb{E}_{P}[\mu_c(t^*)]$ be the expected number of tokens that are predicted by a given token $t^*$, and $\mathbb{E}_{P}[\mu_p(t^*)]$ the expected number of tokens that predict a given token $t^*$.
%
For autoregressive training, the ratio
%
\begin{equation}
\frac{\mathbb{E}_{P}[\mu_c(t^*, U)]}{\mathbb{E}_{P}[\mu_p(t^*, U)]} = \frac{\sum_{k=1}^N (N-k)\Pr[t_k = t^*]}{\sum_{k=1}^N(k-1)\Pr[t_k = t^*]} \,,
\end{equation}
%
depends on $\Pr[t_k = t^*] \,\,\forall t^* \in \mathcal{V}, \forall k \in \{1,\dots. N\}$, while for bidirectional training the same ratio is given by
%
\begin{equation}
\frac{\mathbb{E}_{P}[\mu_c(t^*, U)]}{\mathbb{E}_{P}[\mu_p(t^*, U)]} = 1 \,.
\end{equation}
%
\end{proposition}
% \begin{proposition}[Ratio of expected counts]
% \label{prop:count_ratio}
% %
% Fix a vocabulary $\mathcal V=\{t_0,\dots,t_V\}$ and an integer
% sequence length $N\ge 2$.
% Let
% \(
% (\Omega,\mathscr F,\Pr)
% \)
% be a probability space that carries
% * random tokens
% \(
% T_1,\dots,T_N:\Omega\to\mathcal V,
% \)
% collected in the random sequence
% \(U=(T_1,\dots,T_N)\in\mathcal V^{N}\), and
% * (for the bidirectional‐training analysis) i.i.d.\
% Bernoulli masking variables
% \(
% M_1,\dots,M_N\sim\mathrm{Bernoulli}(\rho),
% \quad 0<\rho<1,
% \)
% which are independent of each other and of the tokens.
% For a fixed token
% \(t^*\in\mathcal V\) define the random variables
% \begin{align}
% \mu_c(t^*,U)
% &=\sum_{k=1}^{N}(N-k)\,
% \mathbf 1_{\{T_k=t^*\}}, &\text{(tokens \emph{predicted by} }t^*),
% \\
% \mu_p(t^*,U)
% &=\sum_{k=1}^{N}(k-1)\,
% \mathbf 1_{\{T_k=t^*\}}, &\text{(tokens that \emph{predict} }t^*).
% \end{align}
% Write
% \(
% E[\cdot]=\E_{U,M_1{:}N}[\cdot]
% \)
% for expectation with respect to all randomness.
% \begin{enumerate}[label=(\roman*),topsep=4pt,itemsep=2pt]
% \item \textbf{Autoregressive training.}
% Then
% \begin{equation}\label{eq:ratio_ar}
% \frac{E[\mu_c(t^*)]}{E[\mu_p(t^*)]}
% \;=\;
% \frac{\displaystyle\sum_{k=1}^{N}(N-k)\Pr[T_k=t^*]}
% {\displaystyle\sum_{k=1}^{N}(k-1)\Pr[T_k=t^*]}\;.
% \end{equation}
% \item \textbf{Bidirectional (masked-LM) training.}
% Assume the model predicts every \emph{masked} token using
% every \emph{unmasked} token.
% Under the independence assumptions stated above,
% \[
% \frac{E[\mu_c(t^*)]}{E[\mu_p(t^*)]}\;=\;1 .
% \]
% \end{enumerate}
% \end{proposition}
% \begin{proof}
% \textbf{(i) Autoregressive case.}
% Because $\mathbf 1_{\{T_k=t^*\}}$ is
% $\{0,1\}$‐valued, linearity of expectation yields
% \[
% E[\mu_c(t^*)]
% =\sum_{k=1}^{N}(N-k)E[\mathbf 1_{\{T_k=t^*\}}]
% =\sum_{k=1}^{N}(N-k)\Pr[T_k=t^*].
% \]
% The same calculation with the weights $(k-1)$ gives
% \(
% \E[\mu_p(t^*)]=\sum_{k=1}^{N}(k-1)\Pr[T_k=t^*].
% \)
% Taking the ratio gives \eqref{eq:ratio_ar}.
% \medskip
% \noindent
% \textbf{(ii) Bidirectional case.}
% Fix a position $k$.
% If $T_k=t^*$ and $M_k=0$ (token visible),
% it contributes
% \(
% \sum_{k'\ne k}\mathbf 1_{\{M_{k'}=1\}}
% \)
% to $\mu_c(t^*)$.
% If $T_k=t^*$ and $M_k=1$ (token masked),
% it contributes
% \(
% \sum_{k'\ne k}\mathbf 1_{\{M_{k'}=0\}}
% \)
% to $\mu_p(t^*)$.
% Taking expectations and using independence of
% \(\{M_i\}\) together with
% \(E[\mathbf 1_{\{M_i=1\}}]=\rho\),
% \(E[\mathbf 1_{\{M_i=0\}}]=1-\rho\),
% we obtain
% \[
% E[\mu_c(t^*)]
% =\sum_{k=1}^{N}\Pr[T_k=t^*]\,(1-\rho)\,(N-1)\rho,
% \qquad
% E[\mu_p(t^*)]
% =\sum_{k=1}^{N}\Pr[T_k=t^*]\,( \rho)\,(N-1)(1-\rho).
% \]
% The multiplicative factors are identical, hence
% \(
% E[\mu_c(t^*)]=E[\mu_p(t^*)]
% \)
% and the stated ratio equals~$1$.
% \end{proof}
% \begin{remark}[Effect of token–position statistics]\label{rem:iid_vs_corr}
% \hfill
% \begin{enumerate}[label=(\alph*),leftmargin=*,topsep=2pt,itemsep=2pt]
% \item \emph{i.i.d.\ tokens.}
% If the tokens are independent and identically distributed,
% \(\Pr[T_k=t^*]=\pi_{t^*}\) for all $k$.
% Then the two sums in \eqref{eq:ratio_ar} coincide,
% giving
% \(
% E[\mu_c(t^*)]=E[\mu_p(t^*)]=
% \pi_{t^*}\tfrac{N(N-1)}2
% \)
% and the ratio is~1. Hence no directional bias appears.
% \item \emph{Correlated tokens.}
% For a general joint law
% \(\Pr[T_1{:}T_N]\),
% \eqref{eq:ratio_ar} remains valid but cannot be simplified.
% Writing the joint law via the reverse chain rule
% \(
% \Pr[T_1{:}T_N]
% =\prod_{i=1}^{N}\Pr[T_i\mid T_{i+1{:}N}],
% \)
% one obtains
% \[
% \Pr[T_k=t^*]
% =\!\!\!\sum_{\substack{t_{-k}\in\mathcal V^{N-1}}}
% \Bigl(
% \Pr[T_k=t^*\mid T_{k+1{:}N}=t_{k+1{:}N}]
% \!\!\!\prod_{i\ne k}
% \Pr[T_i=t_i\mid T_{i+1{:}N}=t_{i+1{:}N}]
% \Bigr),
% \]
% which may depend strongly on $k$.
% If $t^*$ is more likely to occur at early positions
% than late ones, the numerator of \eqref{eq:ratio_ar}
% (weighted by $N-k$) becomes larger than the denominator,
% giving a ratio \(>1\) and
% inducing directionality even under autoregressive training.
% \end{enumerate}
% \end{remark}
\begin{proof}
%
Let $\mathcal{V} = \{t_0, \dots, t_V\}$ be a set of tokens.
%
Let \( \mathcal{U} = \mathcal{V}^N \) denote the sample space of all possible sequences of \( N \) tokens from the vocabulary \( \mathcal{V} \).
%
Let \( \mathcal{F} = 2^{\mathcal{U}} \) be the \(\sigma\)-algebra given by the power set of \( \mathcal{U} \), and let \( P \) be the probability distribution over \( \mathcal{U} \) that defines the distribution of the random variable sequence \( U \).
%
This defines the probability space \( (\mathcal{U}, \mathcal{F}, P) \).
%
For each position $k \in \{1, \dots, N\}$ we let the random variable $T_k$ be $T_k: \mathcal{U} \rightarrow \mathcal{V}$ such that $T_k(U) \equiv t_k$.
%
Let $E = \{U \in \mathcal{U}: t_k = t^*\}$ the event of all sequences where the $k$-th token is a fixed token $t^*$.
%
The indicator function $\mathds{1}\{E\}(U)$ is a random variable defined as
%
\begin{equation}
\mathds{1}\{E\}(U) = \mathds{1}\{t_k = t^*\}(U) =
\begin{cases}
1 \,\,\text{if} \,\,t_k = t^*\\
0 \,\, \text{otherwise} \,,
\end{cases}
\end{equation}
%
and as such its expected value over $P$ is
%
\begin{equation}
\mathbb{E}_P\left[\mathds{1}\{t_k = t^*\}(U)\right] = \sum_{U \in \,\mathcal{U}} \mathds{1}\{t_k = t^*\}(U)\Pr[U] = \sum_{U \in \mathcal{U} \,:\, t_k = t^*} \Pr[U] = \Pr[t_k = t^*] \,.
\end{equation}
%
Let $\mu_c(t^*, U)$ be the random variable quantifying the number of tokens predicted by $t^*$, while $\mu_p(t^*, U)$ is the random variable quantifying the number of tokens predicted by $t^*$.
%
We analyzed the case of autoregressive and bidirectional training separately.
%
During autoregressive training, each time the token $t^*$ appears at position $k$ it is used as context to predict $N-k$ tokens, and it is predicted by $k-1$ tokens.
%
It follows that $\mu_c(t^*, U)$ is given by
%
\begin{equation}
\mu_c(t^*, U) = \sum_{l=2}^N\sum_{k=1}^{l-1} \mathds{1}\{t_k = t^*\} = \sum_{k=1}^N (N-k)\mathds{1}\{t_k = t^*\}\,,
\end{equation}
%
while $\mu_p(t^*, U)$ is given by
%
\begin{equation}
\mu_p(t^*, U) = \sum_{l=1}^{N-1}\sum_{k=1}^{l-1} \mathds{1}\{t_l = t^*\} = \sum_{k=1}^N (k-1)\mathds{1}\{t_k = t^*\}\,.
\end{equation}
%
Therefore, the expected value of $\mu_c(t^*, U)$ over $P$ is
%
\begin{equation}
\mathbb{E}_P[\mu_c(t^*, U)] = \mathbb{E}_P\left[\sum_{k=1}^N (N-k)\mathds{1}\{t_k = t^*\}\right] = \sum_{k=1}^N(N-k)\Pr[t_k = t^*]\,,
\end{equation}
%
the expected value of $\mu_p(t^*)$ is
%
\begin{equation}
\mathbb{E}_P[\mu_p(t^*, U)] = \mathbb{E}_P\left[\sum_{k=1}^N (k-1)\mathds{1}\{t_k = t^*\}\right] = \sum_{k=1}^N(k-1)\Pr[t_k = t^*]\,,
\end{equation}
%
and their ratio is given by
%
\begin{equation}
\label{suppeq:ratio-autoregressive}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*, U)]} = \frac{\sum_{k=1}^N (N-k)\Pr[t_k = t^*]}{\sum_{k=1}^N(k-1)\Pr[t_k = t^*]} \,.
\end{equation}
%
During bidirectional training, each time the token $t^*$ is masked at position $k$, it is predicted by $N-1$ tokens.
%
We assume that the token $t^*$ significantly contributes to the context of a masked token only if it is itself not masked (see Appendix \ref{supp-math-self-attention-derivation-gradients}).
%
Additionally, we assume that the probability $\rho$ of masking any given position $k$ is the same for all positions and does not depend on the specific token $t_k$ at that position.
%
As such, we define masking as an independent Bernoulli variable $m_k: \Sigma_k \rightarrow \{0.1\} \sim P_m$ such that $\text{Pr}[m_k = 1] = \rho$, independently of the token and sequence spaces.
%
It follows that $\mu_c(t^*, U, m)$ is given by
%
\begin{equation}
\mu_c(t^*, U, m) = \sum_{k=1}^N\mathds{1}\{t_k = t^*\} \mathds{1}\{m_k = 0\}\sum_{k' \neq k} \mathds{1}\{m_{k'} = 1\} \,,
\end{equation}
%
while $\mu_p(t^*, U, m)$ is given by
%
\begin{equation}
\mu_p(t^*, U, m) = \sum_{k=1}^N\mathds{1}\{t_k = t^*\} \mathds{1}\{m_{k} = 1\}\sum_{k' \neq k} \mathds{1}\{m_{k'} = 0\}\,.
\end{equation}
%
Therefore, by taking the expectation with respect to both distributions \( P \) and \( P_m \), their respective expected values are given by
%
\begin{equation}
\begin{split}
\mathbb{E}_{P,P_m}[\mu_c(t^*, U, m)]
& = \mathbb{E}_{P,P_m} \left[\sum_{k=1}^N\mathds{1}\{t_k = t^*\} \mathds{1}\{m_{k} = 0\}\sum_{k' \neq k} \mathds{1}\{m_{k'} = 1\} \right] \\
& = \sum_{k=1}^N \mathbb{E}_{P,P_m} \left[\mathds{1}\{t_k = t^*\} \mathds{1}\{m_{k} = 0\}\sum_{k' \neq k} \mathds{1}\{m_{k'} = 1\} \right]\\
& = \sum_{k=1}^N \mathbb{E}_{P,P_m} \left[\mathds{1}\{t_k = t^*\} \mathds{1}\{m_{k} = 0\}\right]\sum_{k' \neq k} \left[\mathds{1}\{m_{k'} = 1\} \right]\\
& = \sum_{k=1}^N \text{Pr}(t_k = t^*) \text{Pr}(m_{k} = 0)\sum_{k' \neq k} \text{Pr}(m_{k'} = 1) \\
& = N\rho(N-1)(1 - \rho) \sum_{k=1}^N \text{Pr}(t_k = t^*)\,,
\end{split}
\end{equation}
%
and
%
\begin{equation}
\begin{split}
\mathbb{E}_{P,P_m}[\mu_p(t^*, U, m)]
& = \mathbb{E}_{P,P_m} \left[\sum_{k=1}^N\mathds{1}\{t_k = t^*\} \mathds{1}\{m_{k} = 1\}\sum_{k' \neq k} \mathds{1}\{m_{k'} = 0\} \right] \\
& = N\rho(N-1)(1 - \rho)\sum_{k=1}^N \text{Pr}(t_k = t^*) \,,
\end{split}
\end{equation}
%
and their ratio is
\begin{equation}
\frac{\mathbb{E}_{P}[\mu_c(t^*, U)]}{\mathbb{E}_{P}[\mu_p(t^*, U)]} = 1 \,,
\end{equation}
%
where we omit the dependence on the independent random variable $m$ for simplicity, thus concluding the proof.
%
\end{proof}
\begin{remark}
\label{remark:ratio-autoregressive}
%
We separately analyze the following two cases:
%
\begin{itemize}[noitemsep,nolistsep]
%
\item \emph{No statistical correlation between tokens.}
%
Assume that the tokens are statistically independent and identically distributed (i.i.d.) across positions, i.e., the probability of each token is the same at every position and does not depend on surrounding tokens.
%
Under this assumption, during autoregressive training, the expected number of times token \( t^* \) is used to predict future tokens is given by,
%
\begin{equation}
\mathbb{E}_P[\mu_c(t^*, U)] = \sum_{k=1}^N (N - k)\Pr[t_k = t^*] = \Pr[t^*] \sum_{k=1}^N (N - k) = \Pr[t^*] \frac{N(N-1)}{2} \,.
\end{equation}
%
Similarly, the expected number of times token \( t^* \) is predicted by preceding tokens is,
\begin{equation}
\mathbb{E}_P[\mu_p(t^*, U)] = \sum_{k=1}^N (k - 1)\Pr[t_k = t^*] = \Pr[t^*] \sum_{k=1}^N (k - 1) = \Pr[t^*] \frac{N(N-1)}{2} \,.
\end{equation}
%
Therefore, when tokens are independent and identically distributed, the expected number of tokens predicted \emph{by} \( t^* \) and the expected number of tokens used to predict \( t^* \) are equal. In this case, there is no difference between autoregressive and bidirectional training.
%
\item \emph{Statistical correlation between tokens.}
%
Let us assume that the probability distribution \( P \) encodes statistical dependencies between tokens.
%
Define \( \mu(t^*) \) as the expected position index of token \( t^* \),
\begin{equation}
\mu(t^*) = \frac{\sum_{k=1}^N k \Pr[t_k = t^*]}{\sum_{k=1}^N \Pr[t_k = t^*]} \,.
\end{equation}
%
Using this, the ratio of expected counts under autoregressive training becomes,
\begin{equation}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*, U)]} = \frac{N - \mu(t^*)}{\mu(t^*) - 1} \,.
\end{equation}
%
This expression shows that if \( t^* \) tends to appear earlier in the sequence, that is, if \( \mu(t^*) < (N+1)/2 \) (left of center), then the ratio is greater than 1.
%
Conversely, if \( t^* \) appears later in the sequence (right of center), the ratio is less than 1.
%
More formally, suppose there exists \( \delta > 0 \) such that
\begin{equation}
\mu(t^*) \le \frac{N+1}{2} - \delta \,.
\end{equation}
%
Then the ratio can be rewritten as,
\begin{equation}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*), U]} - 1
= \frac{N + 1 - 2\mu(t^*)}{\mu(t^*) - 1} \,,
\end{equation}
where the numerator satisfies \( N + 1 - 2\mu(t^*) \ge 2\delta \), and the denominator is positive for any token that can appear at position 1 (i.e., \( \mu(t^*) > 1 \)).
%
Hence, the ratio is bounded below by,
\begin{equation}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*, U)]} \ge 1 + \varepsilon(t^*) \,, \quad
\text{with} \quad \varepsilon(t^*) = \frac{2\delta}{\mu(t^*) - 1} \,.
\end{equation}
%
\end{itemize}
%
\end{remark}
%
%
\begin{lemma}
\label{lemma-same_mean_different_variance_bound}
Let \( a \) and \( b \) be two probability distributions with the same mean \( \mu \), but different variances such that \( \sigma_a^2 > \sigma_b^2 \).
%
Then, for all values \( z > \sqrt{\sigma_a \sigma_b} - \mu \), the probability that a random sample from distribution \( a \) exceeds \( z \) is greater than or equal to the probability that a sample from distribution \( b \) exceeds \( z \), that is,
%
\begin{equation}
\Pr[X_a > z] \geq \Pr[X_b > z] \quad \text{for all } z > \sqrt{\sigma_a \sigma_b} - \mu \,.
\end{equation}
%
\end{lemma}
%
\begin{proof}
%
Our goal is to find a $z$ such that
\begin{equation}
\text{Pr}\left[X_a>z\right]\geq \text{Pr}\left[X_b>z\right]%\quad \forall z> \sqrt{\sigma_a \sigma_b} - \mu
\end{equation}
for which we need the following steps: A lower bound for $\text{Pr}\left[X_a>z\right]$, an upper bound for $\text{Pr}\left[X_b>z\right]$, and a value of $z$ that makes the upper bound lower than the lower bound.
We start by using Chebyshev's inequality, to derive an upper bound for $\text{Pr}\left[X_b>z\right]$,
\begin{equation}
\text{Pr}\left[X_b>z\right]=
\text{Pr}\left[X_b - \mu >q\sigma_b\right]
\leq \dfrac{1}{q^2} = \left(\dfrac{\sigma_b}{z-\mu}\right)^2.
\end{equation}
Now compute a lower bound for $\text{Pr}\left[X_b>z\right]$ through the second moment method (which is very similar to the previous one, but inverted),
\begin{equation}
\text{Pr}\left[X_a>z\right] = \text{Pr}\left[X_a-z>0\right]
\geq \dfrac{\left(\mathbb{E}\left[X_a- z\right]\right)^2}
{\mathbb{E}\left[\left(X_a-z \right)^2 \right]}
= \left(\dfrac{\mu - z}{\sigma_a}\right)^2
\end{equation}
Thus, we now only need to find $z$ satisfying
\begin{equation}
\left(\dfrac{\sigma_b}{z-\mu}\right)^2 < \left(\dfrac{z-\mu}{\sigma_a}\right)^2
\end{equation}
for the case $z>\mu$, the inequality is fulfilled when
\begin{equation}
z>\sqrt{\sigma_a \sigma_b }-\mu.
\end{equation}
%
\end{proof}
\begin{proof}
%
Let $\{U_1, \dots, U_D\}$ be a dataset of sequences of length $N$ drawn i.i.d. from a distribution $P$.
%
Let \( k_u \) denote the position of token \( t^* \) in sequence \( U_u \), and let superscript \( u \) refer to tokens in the \( u \)-th sequence.
%
Following Proposition \ref{prop-gradient-columns-rows} and \ref{prop-gradient-asymmetric-growth-rows-columns}, we obtain
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} = \sum_{u=1}^D \left[[\bm{y}]_k \sum_{i=k_u+1}^N \beta^u_{i*}\bm{x}^u_i + \sum_{j=1}^{k_u-1} \beta^u_{*j}[\bm{x}^u_j]_k\bm{y}
\right] \,,
\end{equation}
%
while the squared norm of the weight update of the $m$-th row by
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} = \sum_{u=1}^D \left[ \sum_{i=k_u+1}^N \beta^u_{i*}[\bm{x}^u_i]_m\bm{y} + [\bm{y}]_m\sum_{j=1}^{k_u-1} \beta^u_{*j}\bm{x}^u_j
\right]\,,
\end{equation}
%
where in both equations the superscript $u$ indicates tokens belonging to the $u$-th sequence in the dataset, and $k_u$ indicates the position index of $t^*$.
%
It follows again from Proposition \ref{prop-gradient-asymmetric-growth-rows-columns} that the expected value of the squared norm of these updates is
%
\begin{equation}
\label{suppeq:ratio-squared-norm-theorem-directionality}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} \approx \frac{\Gamma_{kk}\text{Tr}(\Sigma)\big(\sum_{u=1}^D\sum_{i=k_u+1}^N \beta^u_{i*}\big) + \Sigma_{kk}\text{Tr}(\Gamma)\big(\sum_{u=1}^D\sum_{j=1}^{k_u-1} \beta^u_{*j}\big)}{\Sigma_{mm}\text{Tr}(\Gamma)\big(\sum_{u=1}^D\sum_{i=k_u+1}^N \beta^u_{i*}\big) + \Gamma_{mm}\text{Tr}(\Sigma)\big(\sum_{u=1}^D\sum_{j=1}^{k_u-1} \beta^u_{*j}\big)} \,.
\end{equation}
%
Following Proposition \ref{prop-counting-prediction-context}, let $\mu_c(t^*, U)$ be the random variable quantifying the number of tokens predicted by $t^*$, while $\mu_p(t^*, U)$ is the random variable quantifying the number of tokens predicted by $t^*$.
%
When assuming autoregressive training, it follows that the empirical average of $\mu_c(t^*, U)$ and $\mu_p(t^*, U)$ over the dataset $\{U_1, \dots, U_D\}$ approaches,
%
\begin{equation}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*, U)]} = \frac{\sum_{j=1}^N (N-j)\Pr[t_j = t^*]}{\sum_{j=1}^N(j-1)\Pr[t_j = t^*]} \,.
\end{equation}
%
Following Remark \ref{remark:ratio-autoregressive}, assume that the expected position \( \mu(t^*) \) of \( t^* \) satisfies \( \mu(t^*) < (N+1)/2 - \delta \) for some \( \delta > 0 \), then
%
\begin{equation}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*, U)]} > 1 + \epsilon(t^*)\,,
\end{equation}
%
for some \( \epsilon(t^*) > 0 \).
%
Hence, across the dataset, \( t^* \) appears more often as a context token than as a predicted token.
%
This implies \( \mathbb{E}_P[k_u] < (N+1)/2 \), so that the first terms in the numerator and denominator of Equation~\eqref{suppeq:ratio-squared-norm-theorem-directionality} dominate, leading to,
%
\begin{equation}
\frac{\mathbb{E}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} > 1 \,,
\end{equation}
%
that is, the net increase in column norms exceeds that of row norms.
%
Assume \( \bm{W}_{qk} \) is initialized with i.i.d. entries from a distribution with mean \( \mu \) and variance \( \sigma^2 \), then at the beginning of training
$\mathrm{Var}(\|\bm{w}_{\cdot, k}\|) = \mathrm{Var}(\|\bm{w}_{m, \cdot}\|)
$, while after training we have
%
$
\mathrm{Var}(\|\bm{w}_{\cdot, k}\|) > \mathrm{Var}(\|\bm{w}_{m, \cdot}\|) \,\,\forall k, m.
$
Applying Lemma~\ref{lemma-same_mean_different_variance_bound}, it follows that for all \( w > \gamma \), where \( \gamma = \sqrt{ \mathrm{Var}(\|\bm{w}_{\cdot, k}\|) \cdot \mathrm{Var}(\|\bm{w}_{m, \cdot}\|)} - \mu \),
\begin{equation}
\text{Pr}[||\bm{w}_{\cdot, k}|| > w] > \text{Pr}[|| \bm{w}_{m, \cdot}|| > w] \,\,\ \forall \,w > \gamma \,,
\end{equation}
%
thus concluding the proof.
%
\end{proof}
\begin{corollary}
%
(\textbf{Bidirectional training does not induce directionality})
%
Let \( \mathcal{V} = \{t_0, \dots, t_V\} \) be a vocabulary of tokens, and let \( \mathcal{U} \subset \mathcal{V}^N \) denote the sample space of all sequences of length \( N \). Let \( U = \{t_1, \dots, t_N\} \in \mathcal{U} \) be a random variable with \( U \sim P(U) \). Define \( \Pr[t_j = t^*] \) as the marginal probability that the token at position \( j \) equals \( t^* \in \mathcal{V} \).
%
Let \( \{\bm{x}_1, \dots, \bm{x}_N\} \) be the token embeddings corresponding to the elements of \( U \), where each embedding \( \bm{x}_i \sim \mathcal{D} \) is drawn i.i.d. from a distribution \( \mathcal{D} \) with zero mean and covariance matrix $\text{Cov}(\bm{x}_i) = \Sigma$..
%
Let \( \bm{W}_{qk} \) be query-key matrix of a self-attention mechanism, and let \( \Delta \bm{W}_{qk} \) denote its gradient update as defined in Proposition~\ref{prop-gradients-self-attention}, computed under an autoregressive objective as in Definition~\ref{def-objective-functions}.
%
Then, the variance of the norm of the rows and the columns at the end of training is equal,
%
%
\begin{equation}
\text{Var}\left(||\bm{w}_{m, \cdot}||\right) =\text{Var}\left(||\bm{w}_{ \cdot, k}||\right) \,\, \forall k,m \,.
\end{equation}
%
\end{corollary}
\begin{proof}
%
%
Let $\{U_1, \dots, U_D\}$ be a dataset of sequences of length $N$ drawn i.i.d. from a distribution $P$.
%
Let \( k_u \) denote the position of token \( t^* \) in sequence \( U_u \), and let superscript \( u \) refer to tokens in the \( u \)-th sequence.
%
Following Proposition \ref{prop-gradient-columns-rows} and \ref{prop-gradient-asymmetric-growth-rows-columns}, we obtain
%
\begin{equation}
\Delta \bm{w}_{\cdot, k} = \sum_{u=1}^D \left[[\bm{y}]_k \sum_{i=k_u+1}^N \beta^u_{i*}\bm{x}^u_i + \sum_{j=1}^{k_u-1} \beta^u_{*j}[\bm{x}^u_j]_k\bm{y}
\right] \,,
\end{equation}
%
while the squared norm of the weight update of the $m$-th row by
%
\begin{equation}
\Delta \bm{w}_{m, \cdot} = \sum_{u=1}^D \left[ \sum_{i=k_u+1}^N \beta^u_{i*}[\bm{x}^u_i]_m\bm{y} + [\bm{y}]_m\sum_{j=1}^{k_u-1} \beta^u_{*j}\bm{x}^u_j
\right]\,,
\end{equation}
%
where in both equations the superscript $u$ indicates tokens belonging to the $u$-th sequence in the dataset, and $k_u$ indicates the position index of $t^*$.
%
It follows again from Proposition \ref{prop-gradient-asymmetric-growth-rows-columns} that the expected value of the squared norm of these updates is
%
\begin{equation}
\frac{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{\cdot, k}||^2 \right]}{\mathbb{E}_{\mathcal{D}}\left[ ||\Delta \bm{w}_{m, \cdot}||^2 \right]} \approx \frac{\Gamma_{kk}\text{Tr}(\Sigma)\big(\sum_{u=1}^D\sum_{i=k_u+1}^N \beta^u_{i*}\big) + \Sigma_{kk}\text{Tr}(\Gamma)\big(\sum_{u=1}^D\sum_{j=1}^{k_u-1} \beta^u_{*j}\big)}{\Sigma_{mm}\text{Tr}(\Gamma)\big(\sum_{u=1}^D\sum_{i=k_u+1}^N \beta^u_{i*}\big) + \Gamma_{mm}\text{Tr}(\Sigma)\big(\sum_{u=1}^D\sum_{j=1}^{k_u-1} \beta^u_{*j}\big)} \,.
\end{equation}
%
Following Proposition \ref{prop-counting-prediction-context}, let $\mu_c(t^*, U)$ be the random variable quantifying the number of tokens predicted by $t^*$, while $\mu_p(t^*, U)$ is the random variable quantifying the number of tokens predicted by $t^*$.
%
When assuming bidirectional training training, it follows that the empirical average of $\mu_c(t^*, U)$ and $\mu_p(t^*, U)$ over the dataset $\{U_1, \dots, U_D\}$ approaches,
%
\begin{equation}
\frac{\mathbb{E}_P[\mu_c(t^*, U)]}{\mathbb{E}_P[\mu_p(t^*, U)]} = 1 \,.
\end{equation}
%
From Proposition~\ref{prop-gradient-asymmetric-growth-rows-columns}, it follows that the net increase in the norms of the columns and rows of \( \bm{W}_{qk} \) is the same.
%
Assuming the same initialization conditions as in Theorem~\ref{theo-gradient-directionality}, it follows that by the end of training,
%
\begin{equation}
\mathrm{Var}(\|\bm{w}_{m, \cdot}\|) = \mathrm{Var}(\|\bm{w}_{\cdot, k}\|) \quad \forall\, k, m \,,
\end{equation}
%
which concludes the proof.
\end{proof}
\begin{proof}
%
%
Let $U = \{t_1, \dots, t_N\}$ be a sequence of tokens.
%
It follows from Proposition \ref{prop-gradients-self-attention} that the implciit weight update for $\bm{W}_{qk}$ following the gradient of $\mathcal{L}(\bm{U})$ w.r.t. $\bm{W}_{qk}$ is given by
%
\begin{equation}
\label{eq:gradient-self-attention-summations}
\Delta \bm{W}_{qk} = \sum_{i =1}^N \sum_{j=1}^N \beta_{ij} \bm{K}_{ij} \,,
\end{equation}
%
where we neglect any constant of proportionality (e.g. learning rate) for simplicity.
%
The double summation in Equation \eqref{eq:gradient-self-attention-summations} contains $N^2$ elements.
%
Note that \( \bm{K}_{ji} = \bm{K}_{ij}^\top \), so we can rewrite the double summation as follows,
%
\begin{equation}
\sum_{i = 1}^N \sum_{j =1}^N \beta_{ij} \bm{K}_{ij} =
\sum_{i =1}^N \beta_{ii} \bm{K}_{ii} +
\sum_{\substack{i,j = 1 \\ i < j}}^N(\beta_{ij} \bm{K}_{ij} + \beta_{ji} \bm{K}_{ji})
\end{equation}
%
where the first term includes the diagonal terms, and the second includes the contributions of every pair $(i,j)$ with $i,j \in [0, \dots, N]$.
%
The second term can be written as,
%
\begin{equation}
\sum_{\substack{i,j = 1 \\ i < j}}^N (\beta_{ij} \bm{K}_{ij} + \beta_{ji} \bm{K}_{ji}) = \sum_{\substack{i,j = 1 \\ i < j}}^N(\beta_{ij} \bm{K}_{ij} + \beta_{ji} {\bm{K}_{ij}}^\top) \,,
\end{equation}
%
and by decomposing $\bm{K}_{ij}$ in its symmetric and skew-symmetric parts, such that $\bm{K}_{ij} = \bm{S}_{ij} + \bm{N}_{ij}$, we obtain,
%
\begin{equation}
\label{eq:gradients-symmetric-terms}
\begin{split}
\sum_{\substack{i,j = 1 \\ i < j}}^N (\beta_{ij} \bm{K}_{ij} + \beta_{ji} {\bm{K}_{ij}}^\top)
& = \sum_{\substack{i,j = 1 \\ i < j}}^N \big[ \beta_{ij}(\bm{S}_{ij} + \bm{N}_{ij}) + \beta_{ji}(\bm{S}_{ij} + \bm{N}_{ij})^T \big] \\
& = \sum_{\substack{i,j = 1 \\ i < j}}^N \big[(\beta_{ij} + \beta_{ji})\bm{S}_{ij} + (\beta_{ij} - \beta_{ji})\bm{N}_{ij}\big]\,.
\end{split}
\end{equation}
%
Let $\beta_{ij}$ and $\beta_{ji}$ be such that $\text{sign}(\beta_{ij}) = \text{sign}(\beta_{ji})$ and $|\beta_{ij}| \approx $$|\beta_{ji}|$.
%
By defining $\Delta \bm{W}_{qk}\big|_{\bm{t}_i \leftrightarrow \bm{t}_j} = \beta_{ij}\bm{K}
_{ij} + \beta_{ji}{\bm{K}_{ij}}^T$,
%
it follows that
%
\begin{equation}
\Delta \bm{W}_{qk}\big|_{\bm{t}_i \leftrightarrow \bm{t}_j} \approx \sum_{\substack{i,j = 1 \\ i < j}}^N \beta_{ij}\bm{S}_{ij} = \sum_{\substack{i,j = 1 \\ i < j}}^N \beta_{ij}{\bm{S}}^\top_{ij} = \Delta {\bm{W}_{qk}}^\top\big|_{\bm{t}_i \leftrightarrow \bm{t}_j} \,
\end{equation}
%
thus concluding the proof.
%
\end{proof}
%
\begin{remark}
%
Let $\mathcal{C}_i = [0,1,\dots, N]$ and let the summation indexed by $i$ to run over a random subset of tokens $ M \subset [0,1,\dots, N]$.
%
The weight update of $\bm{W}_{qk}$ is then given by
%
\begin{equation}
\Delta \bm{W}_{qk}^l = \sum_{i\in M} \sum_{j=1}^N \beta^l_{ij} \bm{K}^{l-1}_{ij} \,.
\end{equation}
%
The double summation contains $N|M|$ elements, where $|M|$ is the cardinality of the subset $M$.
%
We can rewrite the double summation as follows,
%
\begin{equation}
\label{eq:factorization-summation}
\sum_{i \in M} \sum_{j =1}^N \beta^l_{ij} \bm{K}^{l-1}_{ij} =
\sum_{i \in M} \beta^l_{ii} \bm{K}^{l-1}_{ii} +
\sum_{\substack{i,j \in M \\ i < j}} (\beta^l_{ij} \bm{K}^{l-1}_{ij} + \beta^l_{ji} \bm{K}^{l-1}_{ji}) +
\sum_{i \in M} \sum_{\substack{j \in \bar{M}}} \beta^l_{ij} \bm{K}^{l-1}_{ij},
\end{equation}
%
where the first term includes the diagonal terms, the second includes the contributions of the pairs $(i,j)$ with $i,j \in M$, and the third includes the remaining terms with $\bar{M} = [1,\dots,N] \setminus M$.
%
The second term can be written as,
%
\begin{equation}
\sum_{\substack{i,j \in M \\ i < j}} (\beta^l_{ij} \bm{K}^{l-1}_{ij} + \beta^l_{ji} \bm{K}^{l-1}_{ji}) = \sum_{\substack{i,j \in M \\ i < j}} (\beta^l_{ij} \bm{K}^{l-1}_{ij} + \beta^l_{ji} {\bm{K}^{l-1}_{ij}}^\top) \,,
\end{equation}
%
and by decomposing $\bm{K}^{l-1}_{ij}$ in its symmetric and skew-symmetric parts, such that $\bm{K}^{l-1}_{ij} = \bm{S}^{l-1}_{ij} + \bm{N}^{l-1}_{ij}$, we obtain,
%
\begin{equation}
\begin{split}
\sum_{\substack{i,j \in M \\ i < j}} (\beta^l_{ij} \bm{K}^{l-1}_{ij} + \beta^l_{ji} {\bm{K}^{l-1}_{ij}}^\top)
& = \sum_{\substack{i,j \in M \\ i < j}} \big[ \beta_{ij}^l(\bm{S}^{l-1}_{ij} + \bm{N}^{l-1}_{ij}) + \beta_{ji}^l(\bm{S}^{l-1}_{ij} + \bm{N}^{l-1}_{ij})^T \big] \\
& = \sum_{\substack{i,j \in M \\ i < j}} \big[(\beta^l_{ij} + \beta^l_{ji})\bm{S}^{l-1}_{ij} + (\beta^l_{ij} - \beta^l_{ji})\bm{N}^{l-1}_{ij}\big]\,,
\end{split}
\end{equation}
%
with a similar structure as in Equation \eqref{eq:gradients-symmetric-terms}.
%
\end{remark}
%
\begin{lemma}
\label{lemma-norm-symmetric-skew-symmetric}
%
For any square matrix $\bm{M} \in \mathbb{M}_n$ the following equivalence holds
%
\begin{equation}
||\bm{M}||_F^2 = ||\bm{M}_s||_F^2 + ||\bm{M}_n||_F^2 \,.
\end{equation}
%
\end{lemma}
%
\begin{proof}
%
The Frobenius norm of a matrix $\bm{M}$ can be defined as $|| \bm{M} ||_F = \sqrt{\text{Tr}(\bm{M} \bm{M}^\top)}$, and as such we observe that for any square matrix $\bm{M}$ we get
%
\begin{equation}
\begin{split}
%
||\bm{M}||_F & = \sqrt{\text{Tr}\big(\bm{M} \bm{M}^\top \big)} = \sqrt{\text{Tr}\big((\bm{M}_s + \bm{M}_n)(\bm{M}_s + \bm{M}_n)^\top\big)} = \\
%
& = \sqrt{\text{Tr}(\bm{M}_s \bm{M}_s^\top) + \text{Tr}(\bm{M}_s \bm{M}_n^\top) + \text{Tr}(\bm{M}_n \bm{M}_s^\top) + \text{Tr}(\bm{M}_n \bm{M}_n^\top)}\,. \\
%
\end{split}
\end{equation}
%
It follows from the cyclic property of the trace operator that the mixing terms cancel out as follows,
%
\begin{equation}
\text{Tr}(\bm{M}_s \bm{M}_n^\top) + \text{Tr}(\bm{M}_n \bm{M}_s^\top) = - \text{Tr}(\bm{M}_s \bm{M}_n) + \text{Tr}(\bm{M}_s \bm{M}_n) = 0 \,,
\end{equation}
%
resulting in
%
%
\begin{equation}
%
||\bm{M}||_F = \sqrt{\text{Tr}(\bm{M}_s \bm{M}_s^\top) + \text{Tr}(\bm{M}_n \bm{M}_n^\top)}\,.
\end{equation}
%
Therefore, as both terms on the right-hand side are semi-positive definite, we conclude the proof as follows,
%
\begin{equation}
||\bm{M}||_F^2 = \text{Tr}(\bm{M}_s \bm{M}_s^\top) + \text{Tr}(\bm{M}_n \bm{M}_n^\top) = ||\bm{M}_s||_F^2 + ||\bm{M}_n||_F^2 \,.
\end{equation}
%
\end{proof}
%
\begin{proposition}
\label{prop-symmetry-score}
%
The symmetry score $s$ quantifies the degree of symmetry or skew-symmetry of a given square matrix $\bm{M}$. In particular, \\
%
1) The symmetry score $s$ is a scalar value bounded in the range $[-1, 1]$. \\
%
2) A symmetry score $s = \pm 1$ indicates a fully symmetric or skew-symmetric matrix, respectively.
\\
3) The symmetry score of a random matrix $\bm{M} \in \mathbb{M}_n$ with entries $\bm{M}_{ij} \sim p(0,\sigma)$ from a probability distribution with zero mean and finite variance tends to zero as $8/n$ in the limit $n\rightarrow\infty$.
%
\end{proposition}
\begin{proof}
%
To prove the points (1) and (2), we first show that it follows from Lemma~\ref{lemma-norm-symmetric-skew-symmetric} that the squared Frobenious norm of $\bm{M}_s$ and $\bm{M}_n$ are in an orthogonal relation
\begin{equation}
%
||\bm{M}||_F = \sqrt{||\bm{M}_s||^2_F + ||\bm{M}_n||^2_F}\,.
%
\end{equation}
%
Therefore, for any given $\bm{M}$, the norms $||\bm{M}_s||^2_2$ and $||\bm{M}_n||^2_2$ are such that a higher value of the first leads to to a lower value of the second, and vice versa.
%
In particular, it is straightforward to observe that $||\bm{M}_s||_2 = ||\bm{M}||_2$ and $||\bm{M}_n||_2 = 0$ if $\bm{M}$ is symmetric.
%
Next, we derive a decomposition of the squared Frobenious norm of the symmetric and skew-symmetric part of $\bm{M}$.
%
From the definition of $\bm{M}_s$ we obtain that
%
\begin{equation}
\begin{split}
%
||\bm{M}_s||^2_F & = \text{Tr}\big(\bm{M}_s\bm{M}^\top_s\big) = \frac14 \, \text{Tr}\big[(\bm{M} + \bm{M}^\top)(\bm{M}^\top + \bm{M})\big]\\
%
& = \frac14 \, \big[ \text{Tr}(\bm{M}\bm{M}^\top) + \text{Tr}(\bm{M}\bm{M}) + \text{Tr}(\bm{M}^\top\bm{M}^\top) + \text{Tr}(\bm{M}^\top\bm{M})\big] \\
%
& = \frac12 \, \big[\text{Tr}(\bm{M}\bm{M}^\top) + \text{Tr}(\bm{M}\bm{M})\big] \\
%
& = \frac12 \, \big[ ||\bm{M}||_F^2 + \text{Tr}(\bm{M}\bm{M})\big] \,.
%
\end{split}
\end{equation}
%
Since the upper bound for $ ||\bm{M}_s||^2_F$ is given by $ ||\bm{M}||^2_F$, the second term on the left-hand side has an upper bound given by,
%
\begin{equation}
\text{Tr}(\bm{M}\bm{M}) \leq \frac12 ||\bm{M}||_F^2 \,,
\end{equation}
%
A complementary relation holds for the skew-symmetric component of $\bm{M}$,
%
\begin{equation}
||\bm{M_n}||^2_F = \frac14 \, \text{Tr}\big[(\bm{M} - \bm{M}^\top)(\bm{M}^\top - \bm{M})\big] = \frac12 \, \big[ ||\bm{M}||_F^2 - \text{Tr}(\bm{M}\bm{M})\big] \,,
\end{equation}
%
which, following the same logic, defines a lower-bound for $\text{Tr}(\bm{M}\bm{M})$ as follows,
%
\begin{equation}
-\frac12 ||\bm{M}||_F^2 \leq \text{Tr}(\bm{M}\bm{M}) \,.
\end{equation}
%
Given Definition~\ref{def-symmetry-score} we can write,
%
\begin{equation}
s = \frac{||\bm{M}||_F^2 + \text{Tr}(\bm{M}\bm{M}) - ||\bm{M}||_F^2 + \text{Tr}(\bm{M}\bm{M})}{||\bm{M}||_F^2 } = 2 \frac{\text{Tr}(\bm{M}\bm{M})}{||\bm{M}||_F^2 }
\end{equation}
%
and by combining the bounds derived previously we obtain,
%
\begin{equation}
-1 \leq s \leq 1
\end{equation}
%
with
%
\begin{equation}
\begin{cases}
& s = 1 \quad \text{if} \quad \bm{M} = \bm{M}^\top \\
& s = -1 \quad \text{if} \quad \bm{M} = - \bm{M}^\top \\
\end{cases}
\end{equation}
%
To prove the point (3), let each entry $m_{ij} = [\bm{M}]_{ij}$ be an independent, identically distributed sample from a random distribution with mean zero and a finite variance $\sigma^2$.
%
We compute the Frobenius norm of the symmetric and skew-symmetric parts as follows,
%
\begin{equation}
\begin{split}
\|\bm{M}_s\|^2_F & = \sum_{i\neq j}\left(\bm{M}_{ij} + \bm{M}_{ji} \right)^2 + \sum_{i}\left(2\bm{M}_{ii}\right)^2 \\
\|\bm{M}_n\|^2_F & = \sum_{i\neq j}\left(\bm{M}_{ij} - \bm{M}_{ji} \right)^2\, .
\end{split}
\end{equation}
%
Here, the skew-symmetric part has a zero diagonal term (because of the subtraction), and the symmetric part has twice the diagonal of the original matrix $\bm{M}$ (because of the addition).
%
Since the entries are independent, $\bm{M}_{ij}$ is independent of $\bm{M}_{ji}$ for all $j\neq i$, and thus we can treat the off-diagonal entries of the $\bm{M}_s$ and $\bm{M}_n$ terms as a sum and difference of two independent random samples having mean zero and the same variance.
%
It follows that the resulting distribution has a mean zero and a variance of $2\sigma^2$ in both cases,
%
\begin{equation}
\sum_{i\neq j}\left(\bm{M}_{ij} \pm \bm{M}_{ji} \right)^2 = 2\sum_{i=1}^n\sum_{j=i+1}^n \left(\bm{M}_{ij} \pm \bm{M}_{ji} \right)^2
\underset{n\rightarrow\infty}{\approx} n(n-1) \text{Var}\left[\bm{M}_{ij} \pm \bm{M}_{ji}\right]
= n(n-1) 2 \sigma^2 \,,
\end{equation}
%
where the approximation is due to the central limit theorem.
%
Applying a similar logic to the second term on the symmetric norm, each entry is the double of a random i.i.d. distribution with
\begin{equation}
\sum_{i=1}^N\left(2\bm{M}_{ii}\right)^2 \underset{N\rightarrow\infty}{\approx} n\text{Var}\left[\bm{M}_{ij}\right]= n4\sigma^2 \,.
\end{equation}
%
Finally, we take the Frobenius norm of the random matrix itself and apply the same logic, where there are $n^2$ entries with a variance of $\sigma^2$,
%
\begin{equation}
\|\bm{M}\|_F^2 \underset{n\rightarrow\infty}{\approx} n^2\sigma^2\,.
\end{equation}
%
It follows that the symmetry score is given by
%
\begin{equation}
s = 2\,\dfrac{\|\bm{M}_s\|^2_F - \|\bm{M}_n\|^2_F}{\|\bm{M}\|^2_F} \underset{n\rightarrow\infty}{\approx} \dfrac{8\sigma^2n}{\sigma^2n^2} = \dfrac{8}{n}\,,
\end{equation}
%
where the symmetry score is zero in the limit $n \rightarrow\infty$ with convergence from the positive side.
%
\end{proof}
\begin{proposition}
\label{prop-directionality-score}
%
The symmetry score $d$ provides a quantitative measure of the degree of directional bias in a given square matrix $\bm{M}$. \\
%
1) The directionality score $d$ is a scalar value that lies within the range $[-1, 1]$. \\
%
2) For any given $\gamma > 0$, a directionality score $d = \pm 1$ indicates that vectors satisfying the condition defined by $\gamma$
are exclusively present in the rows or the column distribution, respectively.
\\
3) The directionality score of a random matrix $\bm{M} \in \mathbb{M}_n$ with entries $\bm{M}_{ij} \sim p(0,\sigma)$ from a probability distribution with zero mean and a variance that scales as $O(n^{-1}) $ tends to zero in the limit $n\rightarrow\infty$.
%
\end{proposition}
\begin{proof}
%
To prove that the directionality score is bounded in the interval $[-1,1]$, note that $\bar{c}_{\mathbf{M}}, \bar{r}_{\mathbf{M}}>0$ simply because they are sums of norms. As both are positive,
\begin{equation}
|\bar{c}_{\mathbf{M}}- \bar{r}_{\mathbf{M}}|
< \bar{c}_{\mathbf{M}}
< \bar{c}_{\mathbf{M}}+ \bar{r}_{\mathbf{M}}
\end{equation}
and thus
\begin{equation}
\dfrac{\bar{c}_{\mathbf{M}}- \bar{r}_{\mathbf{M}}}{\bar{c}_{\mathbf{M}}+ \bar{r}_{\mathbf{M}}} <1
\end{equation}
and taking the negative sign for the absolute value,
\begin{equation}
- \dfrac{\bar{r}_{\mathbf{M}}- \bar{c}_{\mathbf{M}}}{\bar{c}_{\mathbf{M}}+ \bar{r}_{\mathbf{M}}} > -1
\Rightarrow \dfrac{\bar{c}_{\mathbf{M}}- \bar{r}_{\mathbf{M}}}{\bar{c}_{\mathbf{M}}+ \bar{r}_{\mathbf{M}}} > -1.
\end{equation}
In the extremes $d=\pm1$, the numerator must be equal to the denominator in absolute value, implying that either $\bar{r}_{\mathbf{M}}$ or $\bar{c}_{\mathbf{M}}$ are zero and the other is positive. For completeness, we define the score as zero if both are zero.
Finally, we study the case of a random matrix. We start by noting that the values of $\bar{r}_{\mathbf{M}},\ \bar{c}_{\mathbf{M}}$ are interchanged when we take the transpose, hence
\begin{equation}\label{eq:rowColTranspose_direction}
\bar{r}_{\mathbf{M}} = \bar{c}_{\mathbf{M}^\top}
\end{equation}
Regardless of the scaling of the matrix and the value $\gamma$, the key property of a random matrix is that all entries are drawn from the same distribution. Hence,
\begin{equation}
\text{Pr}\left[\bm{M}_{ij} = x\right] =
\text{Pr}\left[\bm{M}_{ji} = x\right]
\Rightarrow
\text{Pr}\left[\bm{M} = \bm{X}\right]
= \text{Pr}\left[\bm{M} = \bm{X}^\top\right]
\end{equation}
for $x$ and $\bm{X}$ being any arbitrary value or matrix. As a consequence,
\begin{equation}
\text{Pr}\left[\bar{c}_{\mathbf{M}} = x\right]
= \text{Pr}\left[\bar{c}_{\mathbf{M}^\top} = x\right]
= \text{Pr}\left[\bar{r}_{\mathbf{M}} = x\right]
\end{equation}
where the last equality comes from Eq.~\ref{eq:rowColTranspose_direction}. The main point here is that the probability distribution of both rows and columns is the same. Pushing this forward, the expected value of $\bar{r}_{\mathbf{M}}- \bar{c}_{\mathbf{M}}$ is
\begin{align}
\text{E}\left[\bar{r}_{\mathbf{M}}- \bar{c}_{\mathbf{M}}\right]
& = \text{E}\left[\bar{r}_{\mathbf{M}}\right]-\text{E}\left[\bar{c}_{\mathbf{M}}\right]
= \int \bar{c}_{\mathbf{M}} \text{Pr}\left[\bm{M}\right] d\bm{M}
- \int \bar{r}_{\mathbf{M}} \text{Pr}\left[\bm{M}\right] d\bm{M}\\
&= \int \bar{r}_{\mathbf{M}} \text{Pr}\left[\bm{M}^\top\right] d\bm{M}^\top
- \int \bar{r}_{\mathbf{M}} \text{Pr}\left[\bm{M}\right] d\bm{M}
= 0
\end{align}
Furthermore, the expected value of $\bar{r}_{\mathbf{M}}+ \bar{c}_{\mathbf{M}}$ is strictly positive, since both values are positive. Thus, their ratio, the directionality score of a random matrix, is zero.
Notice that to be thorough we must show that their variance is bounded scales down. Since weight initialization has been extensively studied, we will just make a general reference to it here. In machine learning, all weights are initialized with zero mean and variances that scale as $O(n^{-1})$. As $\bm{M}$ is a product of two matrices with such scaling, each entry would consist of the sum of $n$ random variables, where each one has a scaling of $O(n^{-2})$ since it is the product of two random variables with an $O(n^{-1})$ scaling. Thus, the entries of $\bm{M}$ also have a scaling of $O(n^{-1})$. Applying the mean value theorem gives us the desired result.
\end{proof}
We trained three BERT models \citep{devlinBERTPretrainingDeep2019} to examine the evolution of the symmetry and directionality scores throughout the training process. Detailed information regarding the training procedure is provided below.
\subsection{Models}\label{sec:exp_bertmodels_models}
%
\input{ICML2025-rebuttal/tables/supp-table-bert-configs}
We train the standard BERT model (referred to as \emph{BERT}), a smaller version (referred to as \emph{BERT-Mini}), and a larger version (referred to as \emph{BERT-Large}), following the implementation by \citep{devlinBERTPretrainingDeep2019}.
Table \ref{tab:bert_configs} provides an overview of the model parameters.
The standard BERT model has $12$ layers, $12$ attention heads, and embedding dimensions of $768$ for the hidden layers and $3072$ for the intermediate layers. In contrast, the smaller BERT-Mini model uses $4$ layers, $4$ attention heads, and embedding dimensions of $256$ and $1024$, respectively.
The larger BERT-Large model features $24$ layers, $16$ attention heads, and embedding dimensions of $1024$ and $4096$ for the hidden and intermediate layers, respectively.
\paragraph{Initialization}
We optionally initialize the models using symmetric attention weights. Specifically, in each self-attention block, the query weight matrix $\bm{W_q}$ is initialized randomly, and the key weight matrix is set equal to it, $\bm{W_k} = \bm{W_q}$. This ensures that the key and query matrices are identical at initialization, introducing symmetry into the attention mechanism.
%\paragraph{Initialization}
%We optionally initialize the models with either symmetric or skew-symmetric attention weights by initializing the key and query weight matrices of all self-attention blocks following a specific procedure.
%For symmetric initialization, the query weight matrix $\bm{W_q}$ is initialized randomly, and then the key weight matrix is set to $\bm{W_k} = \bm{W_q}$.
%For skew-symmetric initialization, $\bm{W_q}$ is initialized randomly, and a skew-symmetric matrix $\bm{S}$ is generated by first initializing a random matrix, $\bm{A}$ and then calculating \(\bm{S} = \bm{A} - \bm{A}^\top\). The key matrix is then defined as $\bm{W_k} = \bm{W_q} \bm{S}^\top$, ensuring that the matrix multiplication $\bm{W_q} \bm{W_k}^\top$ is skew-symmetric.
%
%\paragraph{Loss Constraints}
%The loss function $\mathcal{L}(\bm{U})$ described in Equation \eqref{eq:results:self-supervised-pretraining} is optionally extended with a constraint to either %encourage symmetric or skew-symmetric attention weights.
%%
%\begin{equation}
% \mathcal{L}(\bm{U}) := \mathcal{L}(\bm{U}) + \gamma \mathcal{L}_{c}(\mathcal{W}).
%\end{equation}
%%
%Here, $\gamma$ is a coefficient defining the importance of the $\mathcal{L}_{c}$ constraint.
%To encourage symmetry, we minimize the condition
%%
%\begin{equation}
% \mathcal{L}_c(\mathcal{W}) = \frac{\bm{M}^2}{\bm{M}_s^2}.
%\end{equation}
%%
%This constraint reaches its minimum if all attention weights have a symmetry score $s=1$.
%To encourage skew-symmetry, we minimize the condition
%%
%\begin{equation}
% \mathcal{L}_c(\mathcal{W}) = \frac{\bm{M}^2}{\bm{M}_n^2},
%\end{equation}
%%
%leading to a symmetry score of $s=0$.
\subsection{Datasets}\label{sec:exp_bertmodels_datasets}
The models are trained on three datasets.
First, we use the ``20220301.en'' snapshot from the Wikipedia dataset, which consists of 6.46 million samples crawled from Wikipedia. Second, we utilize the Jigsaw dataset with $159$K samples, originally collected for a toxic comment classification challenge.
Finally, we train on the English ``2023-14'' snapshot of the RedPajama-V2 dataset, which contains approximately $5.12$ billion samples.
\subsection{Training Settings}\label{sec:exp_bertmodels_training}
The models are trained for $200,000$ update steps with a batch size of $32$ and $8$ gradient accumulation steps, effectively increasing the batch size to $256$ before each parameter update. The optimization is done using the AdamW optimizer \citep{Loshchilov_Hutter_2019}, and the training schedule includes $200$ warmup steps to stabilize early training, followed by a linear decay learning rate schedule, starting at an initial learning rate of $5 \times 10^{-5}$ and weight decay of $0.01$.
Mixed precision (fp16) training was utilized to maximize training efficiency, which reduces memory consumption and speeds up computation without significant loss of precision. The training data was processed with a masked language modeling (MLM) probability of $15$\%, ensuring that $15$\% of tokens were masked during training.
The models are trained in the encoder and decoder mode, i.e., to predict masked tokens and subsequent tokens respectively.
% \subsection{Linear Attention}\label{sec:exp_linearattention}
% We train a minimalistic attention block to investigate the origins and properties of the symmetry of the attention weights.
% Details about this model, model configuration, and datasets are provided in the following.
% \subsubsection{Models}\label{sec:exp_linearattention_models}
% The minimal model consists of an embedding layer that maps the token to embeddings, a positional encoding layer to add positional information to the embeddings, and one or multiple self-attention blocks \citep{vaswaniAttentionAllYou2017}.
% We apply skip-connections around the attention blocks but do not use subsequent linear layers or non-linear activation functions (hence the name ``linear attention'').
% After the last attention block, we use an output layer to map the embeddings to the dimensionality of the tokens' space.
% We use the same weights for the output and token embedding layers, i.e., implement weight tying \citep{press-wolf-2017-using}.
% Similar to the BERT models (cf. Section \ref{sec:exp_bertmodels_models}), we optionally apply symmetric or skew-symmetric initialization and a symmetric or skew-symmetric training constraint.
% \subsubsection{Datasets}\label{sec:exp_linearattention_datasets}
% We generate two artificial datasets for our experiments. The first dataset is constructed using a vocabulary of $10$ tokens. Each sample in this dataset consists of token sequences. To create a sequence, we randomly select a token from the vocabulary and append it between $0$ and $3$ times to the sequence. This process is repeated for all tokens, forming a sub-sequence until a total length of $100$ tokens is reached.
% The second artificial dataset consists of symmetric sequences mirrored in the middle. The vocabulary comprises $48$ tokens. A token is randomly selected from the vocabulary to construct a sequence, repeated between $1$ and $3$ times, and appended to the sequence. This process is repeated until a sequence of $48$ tokens is obtained. The resulting sequence is then duplicated, mirrored, and concatenated to the original, producing a final sequence of $96$ tokens.
% To make the prediction task more challenging, we add jitter directly to the training data by randomly replacing some tokens with a certain probability.
% \subsubsection{Training Settings}\label{sec:exp_linearattention_training}
% We train the models for at least $2,000$ epochs to capture changes in symmetry even though the models' performance (in terms of loss and prediction accuracy) already levels off after a few hundred epochs.
% We use a batch size of $128$ and a constant learning rate of $1 \cdot 10^{-3}$ using either the AdamW \citep{Loshchilov_Hutter_2019} optimizer or SGD.
\begin{table}
%
\centering
\begin{tabular}{l|c|c|c}
\toprule
\textbf{Configuration} & \textbf{BERT} & \textbf{BERT-Mini} & \textbf{BERT-Large} \\
\midrule
%
\textit{Hidden Size} & 768 & 256 & 1024 \\
\textit{Intermediate Size} & 3072 & 1024 & 4096 \\
\textit{Number of Attention Heads} & 12 & 4 & 16 \\
\textit{Number of Hidden Layers} & 12 & 4 & 24 \\
\textit{Attention Dropout Probability} & 0.1 & 0.1 & 0.1 \\
\textit{Hidden Activation Function} & gelu & gelu & gelu \\
\textit{Hidden Dropout Probability} & 0.1 & 0.1 & 0.1 \\
\textit{Layer Normalization Epsilon} & 1e-12 & 1e-12 & 1e-12 \\
\textit{Max Position Embeddings} & 512 & 512 & 512 \\
\textit{Position Embedding Type} & absolute & absolute & absolute \\
\textit{Vocabulary Size} & 30522 & 30522 & 30522 \\
%
\bottomrule
\end{tabular}
\caption{Configurations for BERT, BERT-Mini, and BERT-Large models.}
%
\label{tab:bert_configs}
\end{table}
%
\input{ICML2025-rebuttal/tables/supp-table-vit-configs}
%
In parallel to the experiments described in Section~\ref{sec-symmetric-initialization}, where symmetric initialization of the query-key weight matrix $\bm{W}_{qk}$ is investigated for BERT models, we extend the investigation to the vision domain. Specifically, we evaluate the impact of symmetric initialization in vision transformers (ViTs) \citep{dosovitskiyImageWorth16x162021}, training a 6-layer ViT on CIFAR-10 \citep{krizhevsky2009learning} and a 12-layer ViT on ImageNet-1k \citep{deng_imagenet_2009}, using the same initialization strategy as applied to BERT models (see Appendix~\ref{sec:exp_bertmodels_models} for implementation details). The architectural configurations for both models are summarized in Table~\ref{tab:vit_configs}.
Training is performed using the AdamW optimizer \citep{Loshchilov_Hutter_2019} with weight decay and a cosine annealing learning rate schedule. A linear warm-up phase is used during the first $30$ epochs, with an initial learning rate scaled by a factor of $0.033$. The ImageNet-1k model is trained for $200$ epochs using a per-device batch size of $256$ and gradient accumulation over two steps, while the CIFAR-10 model is trained for $500$ epochs with a batch size of $128$ and no gradient accumulation. Both setups use a base learning rate of $0.003$, optimizer hyperparameters $\beta_1 = 0.9$ and $\beta_2 = 0.999$, and automatic mixed-precision training (fp16). We apply advanced data augmentation tailored to each dataset, including mixup ($\alpha = 0.2$) \citep{zhang2018mixup}, cutmix ($\alpha = 1.0$) \citep{Yun_cutmix_2019}, and label smoothing ($0.11$ for ImageNet-1k and $0.01$ for CIFAR-10). An exponential moving average (EMA) of the model weights \citep{Tarvainen_mean_2017} is maintained throughout training (decay rate = $0.99998$, update frequency = $32$ steps).
%
\input{ICML2025-rebuttal/tables/sup-table-results-symmetric-initialization-vit}
Contrary to the trends observed in our language model experiments, symmetric initialization does not result in faster convergence or improved final performance for the vision transformer models. As shown in Table~\ref{table:vit-symm-results}, training and evaluation losses, as well as top-1 accuracy, remained nearly identical between the standard and symmetric initialization conditions across both datasets.
Importantly, the results with symmetric initialization are not degraded, indicating that such initialization is at least performance-neutral in our vision setting. While our findings suggest that symmetric initialization does not yield immediate benefits for ViTs under the given training regime, we do not rule out the possibility that alternative configurations or optimization strategies could leverage its potential. Further research may uncover conditions under which vision transformer architectures can benefit more substantially from symmetric initialization.
\begin{table}
%
\centering
\begin{tabular}{l|c|c}
\toprule
\textbf{Configuration} & \textbf{ViT (6 layers)} & \textbf{ViT (12 layers)} \\
& CIFAR-10 & ImageNet-1k \\
\midrule
%
\textit{Hidden Size} & 512 & 768 \\
\textit{Intermediate Size} & 2048 & 3072 \\
\textit{Number of Attention Heads} & 8 & 12 \\
\textit{Number of Hidden Layers} & 6 & 12 \\
\textit{Attention Dropout Probability} & 0.0 & 0.0 \\
\textit{Hidden Activation Function} & gelu & gelu \\
\textit{Hidden Dropout Probability} & 0.0 & 0.0 \\
\textit{Layer Normalization Epsilon} & 1e-6 & 1e-6 \\
\textit{Patch Size} & 4 & 16 \\
\textit{QKV Bias} & true & true \\
\textit{Encoder Stride} & 16 & 16 \\
%
\bottomrule
\end{tabular}
\caption{Configurations for the 6-layer and 12-layer vision transformer models.}
%
\label{tab:vit_configs}
\end{table}
\begin{table}[t]
\caption{Final training and evaluation metrics for ViT models on CIFAR-10 \citep{krizhevsky2009learning} and ImageNet-1k \citep{deng_imagenet_2009}, trained with and without symmetric initialization.
%
Evaluation metrics include final loss and top-1 accuracy. For each dataset, we compare models trained with standard initialization and with symmetric initialization.
%
}
\label{table:vit-symm-results}
\vskip 0.15in
\begin{center}
\begin{small}
\begin{sc}
\begin{tabular}{lccc}
\toprule
Dataset & Train Loss & Eval Loss & Eval Accuracy \\
\midrule\midrule
\textbf{CIFAR-10} & & & \\
\midrule
ViT 6-layers & 1.07 & 0.40 & {0.90} \\
ViT 6-layers (+ symm.) & {1.08} & {0.42} & 0.89 \\
\midrule
\textbf{ImageNet-1k} & & & \\
\midrule
ViT 12-layers & {2.49} & 1.97 & {0.76} \\
ViT 12-layers (+ symm.) & {2.49} & {1.96} & {0.76} \\
\bottomrule
\end{tabular}
\end{sc}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
%
\input{ICML2025/figures/suppfigure-pretrained-models}
%
\input{ICML2025/figures/suppfigure-encoder-decoder-models}
%
\input{ICML2025/figures/suppfigure-training-custom-models}
%
\input{ICML2025/figures/suppfigure-training-custom-models-layers}
\input{ICML2025/figures/suppfigure-symmetric-initialization}
%
%
\begin{figure*}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/suppfigure-pretrained-models.pdf}
\end{center}
%
\caption{
%
%
\textbf{a})
%%
Median symmetry score of the matrix $\bm{W}_{qk}$ as a function of the total number of parameters for vision and audio models.
%
Each dot corresponds to the median and the interquartile range across layers of a given pre-trained model.
%
\textbf{b})
%%
Same as in \textbf{a} for the median directionality score of the matrix $\bm{W}_{qk}$.
%
}
%
\label{suppfig-pretrained-models}
%
\end{figure*}
%
\begin{figure*}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/suppfigure-encoder-decoder-models.pdf}
\end{center}
%
\caption{
%
%
\textbf{a})
%%
Median symmetry score of the matrix $\bm{W}_{qk}$ as a function of the total number of parameters for vision and audio models.
%
Each dot corresponds to the median and the interquartile range across layers of the encoder component (blue) and decoder component (purple) of an encoder-decoder Transformer model.
%
The encoder component of these models shows a high degree of symmetry compared to the decoder component.
%
\textbf{b})
%%
Same as in \textbf{a} for the median directionality score of the matrix $\bm{W}_{qk}$.
%
The encoder and decoder components of these models do not show significant differences in directionality scores.
%
}
%
\label{suppfig-encoder-decoder-models}
%
\end{figure*}
%
\begin{figure*}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/suppfigure-training-custom-models.pdf}
\end{center}
%
\caption{
%
\textbf{a})
%
Evolution of symmetry score during training.
%
Shown are the median and the interquartile range.
%
Models were trained on the Jigsaw dataset \citep{jigsaw_challenge} (left) and on the Red Pajama dataset \citep{together2023redpajama} (right).
%
Encoder-only and decoder-only models are color-coded in blue and purple, respectively (see legend).
%
\textbf{b})
%
Same as in panel \textbf{a} for the median directionality score.
%
}
%
\label{suppfig-training-custom-models}
%
\end{figure*}
%
\begin{figure*}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/suppfigure-training-custom-models-layers.pdf}
\end{center}
%
\caption{
%
\textbf{a})
%
Evolution of the median symmetry score across layers of the encoder-only (left) and decoder-only (right) models.
%
Each layer is color-coded as shown on the legend of Figure \ref{fig-training-layers}.
%
Shown are the median and the interquartile range.
%
Models were trained on the Jigsaw dataset \citep{jigsaw_challenge}.
%
\textbf{b})
%
Same as in panel \textbf{a} for the median directionality score.
%
\textbf{c})
%
Same as in panel \textbf{a} for models trained on the Red Pajama dataset \citep{together2023redpajama}.
%
\textbf{d})
%
Same as panel $\textbf{c}$ for the median directionality score.
%
}
%
\label{suppfig-training-custom-models-layers}
%
\end{figure*}
\begin{figure}[ht]
\begin{center}
\includegraphics[width=1.\linewidth]{ICML2025/figures/figures-pdf/suppfigure-symmetric-initialization.pdf}
\end{center}
%
\caption{
%
\textbf{a})
%
Loss a 4-layer encoder-only model with (dark blue) and without (light blue) symmetric initialization, respectively.
%
Models are trained on the Red Pajama dataset \citep{together2023redpajama}.
%
\textbf{b})
%
Median symmetry score during training
%
Color code is as in panel \textbf{a}.
%
\textbf{c})
%
Same as in panel \textbf{a} for each layer in the model with and without symmetric initialization.
%
Each layer is color-coded as in the illustration on the left.
%
All plots show the median and the interquartile range across the heads of a given layer.
%
}
%
\label{suppfig-symmetric-initialization}
%
\end{figure}
\input{ICML2025/tables/table-symmetry-score-language}
\input{ICML2025/tables/table-directionality-score-language}
\begin{table}
\centering
\caption{Symmetry score for open-source pretrained language models. All models are available on Hugging Face \citep{wolfHuggingFaceTransformersStateoftheart2020}.}
\label{table:symmetry-score-models}
\vspace{5pt}
\begin{tabular}{lcc|lcc}
\toprule
\textbf{Model} & \textbf{Median} & \textbf{Interquartile range} & \textbf{Model} & \textbf{Median} & \textbf{Interquartile range} \\
\midrule
BERT-tiny & 0.77 & $\pm$ [0.07, 0.07] & GPT-Neo1.3B & 0.14 & $\pm$ [0.03, 0.03] \\
BERT-mini & 0.62 & $\pm$ [0.03, 0.05] & GPT-Neo2.7B & 0.13 & $\pm$ [0.02, 0.04] \\
BERT-small & 0.69 & $\pm$ [0.10, 0.08] & GPT-J6B & 0.11 & $\pm$ [0.02, 0.03] \\
BERT-medium & 0.60 & $\pm$ [0.01, 0.02] & OpenAI-GPT & 0.07 & $\pm$ [0.04, 0.03] \\
BERT-base & 0.51 & $\pm$ [0.09, 0.07] & GPT2-XL & 0.12 & $\pm$ [0.03, 0.05] \\
BERT-large & 0.44 & $\pm$ [0.03, 0.08] & DistilGPT2 & 0.19 & $\pm$ [0.05, 0.05] \\
DistilBERT & 0.43 & $\pm$ [0.10, 0.13] & GPT-Neo125M & 0.14 & $\pm$ [0.09, 0.14] \\
BERT-2L-128 & 0.77 & $\pm$ [0.07, 0.07] & GPT-Neo1.3B & 0.14 & $\pm$ [0.03, 0.03] \\
BERT-4L-256 & 0.62 & $\pm$ [0.03, 0.05] & GPT-Neo2.7B & 0.14 & $\pm$ [0.03, 0.02] \\
BERT-4L-512 & 0.69 & $\pm$ [0.10, 0.08] & GPT-J6B & 0.11 & $\pm$ [0.02, 0.03] \\
BERT-8L-512 & 0.60 & $\pm$ [0.01, 0.02] & LLaMA2-7B & 0.12 & $\pm$ [0.02, 0.03] \\
BERT-base & 0.51 & $\pm$ [0.09, 0.07] & LLaMA2-13B & 0.17 & $\pm$ [0.02, 0.02] \\
BERT-large & 0.44 & $\pm$ [0.03, 0.08] & LLaMA3-8B & 0.00 & $\pm$ [0.00, 0.01] \\
DistilBERT & 0.43 & $\pm$ [0.10, 0.13] & LLaMA3.1-8B & 0.00 & $\pm$ [0.00, 0.01] \\
BEiT-base & 0.40 & $\pm$ [0.08, 0.02] & LLaMA3.2-8B & 0.01 & $\pm$ [0.01, 0.01] \\
BEiT-large & 0.33 & $\pm$ [0.05, 0.07] & LLaMA3.2-1B & 0.01 & $\pm$ [0.01, 0.01] \\
BEiT-base & 0.39 & $\pm$ [0.23, 0.07] & LLaMA3.2-3B & 0.01 & $\pm$ [0.01, 0.01] \\
BEiT-large & 0.26 & $\pm$ [0.17, 0.13] & LLaMA2-7B-chat & 0.12 & $\pm$ [0.02, 0.03] \\
BEiT-base & 0.39 & $\pm$ [0.23, 0.07] & LLaMA2-70B & 0.02 & $\pm$ [0.01, 0.02] \\
BEiT-large & 0.26 & $\pm$ [0.17, 0.13] & LLaMA2-7B-chat & 0.12 & $\pm$ [0.02, 0.03] \\
BEiT-large & 0.26 & $\pm$ [0.17, 0.13] & LLaMA2-13B-chat & 0.17 & $\pm$ [0.02, 0.02] \\
ALBERT-base & 0.72 & $\pm$ [0.00, 0.00] & LLaMA3-8B & 0.00 & $\pm$ [0.00, 0.00] \\
ALBERT-large & 0.70 & $\pm$ [0.00, 0.00] & LLaMA3-70B & 0.02 & $\pm$ [0.01, 0.01] \\
ALBERT-xlarge & 0.59 & $\pm$ [0.00, 0.00] & LLaMA3.1-8B & 0.00 & $\pm$ [0.00, 0.01] \\
ALBERT-xxlarge & 0.46 & $\pm$ [0.00, 0.00] & LLaMA3.1-70B & 0.01 & $\pm$ [0.00, 0.01] \\
RoBERTa-base & 0.49 & $\pm$ [0.03, 0.06] & LLaMA3.1-405B & 0.03 & $\pm$ [0.01, 0.03] \\
RoBERTa-large & 0.47 & $\pm$ [0.06, 0.06] & LLaMA3.2-1B & 0.00 & $\pm$ [0.00, 0.00] \\
XLM-R-base & 0.51 & $\pm$ [0.05, 0.03] & LLaMA3.2-3B & 0.01 & $\pm$ [0.01, 0.01] \\
XLM-R-large & 0.49 & $\pm$ [0.16, 0.12] & Mistral-7B & 0.00 & $\pm$ [0.00, 0.01] \\
RoBERTa-mnli & 0.47 & $\pm$ [0.06, 0.06] & Mixtral-8x22B & 0.00 & $\pm$ [0.00, 0.00] \\
DistilRoBERTa & 0.53 & $\pm$ [0.02, 0.06] & MobileLLM125M & 0.03 & $\pm$ [0.02, 0.03] \\
ModernBERT-base & 0.18 & $\pm$ [0.06, 0.18] & MobileLLM350M & 0.01 & $\pm$ [0.01, 0.01] \\
GPT1 & 0.07 & $\pm$ [0.04, 0.03] & Phi-1.5 & 0.09 & $\pm$ [0.03, 0.03] \\
GPT2 & 0.15 & $\pm$ [0.02, 0.03] & Phi-1 & 0.14 & $\pm$ [0.02, 0.01] \\
GPT2-medium & 0.17 & $\pm$ [0.03, 0.05] & Phi-2 & 0.07 & $\pm$ [0.03, 0.06] \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}
\centering
\caption{Directionality score for open-source pretrained language models. All models are available on Hugging Face \citep{wolfHuggingFaceTransformersStateoftheart2020}.}
\label{table:directionality-score-models}
\vspace{5pt}
\begin{tabular}{lcc|lcc}
\toprule
\textbf{Model} & \textbf{Median} & \textbf{Interquartile range} & \textbf{Model} & \textbf{Median} & \textbf{Interquartile range} \\
\midrule
BERT-tiny & -0.79 & $\pm$ [0.11, 0.11] & GPT-Neo1.3B & -0.49 & $\pm$ [0.19, 0.13] \\
BERT-mini & -0.33 & $\pm$ [0.03, 0.04] & GPT-Neo2.7B & -0.57 & $\pm$ [0.15, 0.16] \\
BERT-small & -0.22 & $\pm$ [0.04, 0.03] & GPT-J6B & -0.28 & $\pm$ [0.09, 0.08] \\
BERT-medium & -0.23 & $\pm$ [0.06, 0.10] & OpenAI-GPT & -0.18 & $\pm$ [0.08, 0.07] \\
BERT-base & -0.08 & $\pm$ [0.02, 0.03] & GPT2-XL & -0.23 & $\pm$ [0.11, 0.10] \\
BERT-large & -0.03 & $\pm$ [0.02, 0.06] & DistilGPT2 & -0.51 & $\pm$ [0.03, 0.07] \\
DistilBERT & -0.13 & $\pm$ [0.00, 0.06] & GPT-Neo125M & -0.56 & $\pm$ [0.21, 0.08] \\
BERT-2L-128 & -0.79 & $\pm$ [0.11, 0.11] & GPT-Neo1.3B & -0.49 & $\pm$ [0.19, 0.13] \\
BERT-4L-256 & -0.33 & $\pm$ [0.03, 0.04] & GPT-Neo2.7B & -0.49 & $\pm$ [0.15, 0.21] \\
BERT-4L-512 & -0.22 & $\pm$ [0.04, 0.03] & GPT-J6B & -0.28 & $\pm$ [0.09, 0.08] \\
BERT-8L-512 & -0.23 & $\pm$ [0.06, 0.10] & LLaMA2-7B & -0.26 & $\pm$ [0.09, 0.13] \\
BERT-base & -0.08 & $\pm$ [0.02, 0.03] & LLaMA2-13B & -0.15 & $\pm$ [0.11, 0.03] \\
BERT-large & -0.03 & $\pm$ [0.02, 0.06] & LLaMA3-8B & -0.65 & $\pm$ [0.13, 0.20] \\
DistilBERT & -0.13 & $\pm$ [0.00, 0.06] & LLaMA3.1-8B & -0.64 & $\pm$ [0.17, 0.19] \\
BEiT-base & -0.10 & $\pm$ [0.06, 0.15] & LLaMA3.2-8B & -0.59 & $\pm$ [0.18, 0.22] \\
BEiT-large & -0.15 & $\pm$ [0.08, 0.07] & LLaMA3.2-1B & -0.59 & $\pm$ [0.18, 0.22] \\
BEiT-base & -0.14 & $\pm$ [0.15, 0.21] & LLaMA3.2-3B & -0.77 & $\pm$ [0.08, 0.19] \\
BEiT-large & -0.14 & $\pm$ [0.04, 0.14] & LLaMA2-7B-chat & -0.29 & $\pm$ [0.07, 0.14] \\
BEiT-base & -0.14 & $\pm$ [0.15, 0.21] & LLaMA2-70B & -0.24 & $\pm$ [0.10, 0.06] \\
BEiT-large & -0.14 & $\pm$ [0.04, 0.14] & LLaMA2-7B-chat & -0.29 & $\pm$ [0.07, 0.14] \\
BEiT-large & -0.15 & $\pm$ [0.03, 0.14] & LLaMA2-13B-chat & -0.19 & $\pm$ [0.12, 0.04] \\
ALBERT-base & -0.07 & $\pm$ [0.00, 0.00] & LLaMA3-8B & 0.01 & $\pm$ [0.05, 0.05] \\
ALBERT-large & -0.17 & $\pm$ [0.00, 0.00] & LLaMA3-70B & -0.37 & $\pm$ [0.09, 0.12] \\
ALBERT-xlarge & -0.24 & $\pm$ [0.00, 0.00] & LLaMA3.1-8B & -0.57 & $\pm$ [0.16, 0.13] \\
ALBERT-xxlarge & -0.15 & $\pm$ [0.00, 0.00] & LLaMA3.1-70B & -0.37 & $\pm$ [0.08, 0.11] \\
RoBERTa-base & -0.12 & $\pm$ [0.11, 0.03] & LLaMA3.1-405B & -0.17 & $\pm$ [0.07, 0.07] \\
RoBERTa-large & -0.06 & $\pm$ [0.03, 0.03] & LLaMA3.2-1B & -0.02 & $\pm$ [0.13, 0.08] \\
XLM-R-base & -0.02 & $\pm$ [0.02, 0.02] & LLaMA3.2-3B & -0.70 & $\pm$ [0.07, 0.22] \\
XLM-R-large & -0.02 & $\pm$ [0.03, 0.02] & Mistral-7B & -0.58 & $\pm$ [0.15, 0.13] \\
RoBERTa-mnli & -0.06 & $\pm$ [0.03, 0.03] & Mixtral-8x22B & -0.66 & $\pm$ [0.09, 0.16] \\
DistilRoBERTa & -0.14 & $\pm$ [0.09, 0.07] & MobileLLM125M & -0.13 & $\pm$ [0.15, 0.10] \\
ModernBERT-base & -0.04 & $\pm$ [0.05, 0.04] & MobileLLM350M & -0.34 & $\pm$ [0.13, 0.23] \\
GPT1 & -0.18 & $\pm$ [0.08, 0.07] & Phi-1.5 & -0.28 & $\pm$ [0.22, 0.19] \\
GPT2 & -0.58 & $\pm$ [0.06, 0.14] & Phi-1 & -0.40 & $\pm$ [0.03, 0.04] \\
\bottomrule
\end{tabular}
\end{table}
\begin{thebibliography}{61}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Abdin et~al.(2024)Abdin, Aneja, Awadalla, Awadallah, Awan, Bach, Bahree, Bakhtiari, Bao, Behl, et~al.]{abdinPhi3TechnicalReport2024}
Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed Awadallah, Ammar~Ahmad Awan, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat Behl, et~al.
\newblock {Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone}, August 2024.
\newblock arXiv:2404.14219.
\bibitem[Bao et~al.(2022)Bao, Dong, Piao, and Wei]{baoBEiTBERTPreTraining2022}
Hangbo Bao, Li~Dong, Songhao Piao, and Furu Wei.
\newblock {BEiT: BERT Pre-Training of Image Transformers}.
\newblock In \emph{Proceedings of the 10th International Conference on Learning Representations}, ICLR, 2022.
\newblock URL \url{https://openreview.net/forum?id=p-BhZSz59o4}.
\bibitem[Barkan et~al.(2021)Barkan, Hauon, Caciularu, Katz, Malkiel, Armstrong, and Koenigstein]{Barkan_2021}
Oren Barkan, Edan Hauon, Avi Caciularu, Ori Katz, Itzik Malkiel, Omri Armstrong, and Noam Koenigstein.
\newblock {Grad-SAM: Explaining Transformers via Gradient Self-Attention Maps}.
\newblock In \emph{Proceedings of the 30th ACM International Conference on Information \& Knowledge Management}, CIKM, page 2882–2887, Virtual Event Queensland Australia, October 2021. ACM.
\newblock ISBN 978-1-4503-8446-9.
\newblock \doi{10.1145/3459637.3482126}.
\bibitem[Belrose et~al.(2023)Belrose, Furman, Smith, Halawi, Ostrovsky, McKinney, Biderman, and Steinhardt]{Belrose_Furman_Smith_Halawi_Ostrovsky_McKinney_Biderman_Steinhardt_2023}
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt.
\newblock {Eliciting Latent Predictions from Transformers with the Tuned Lens}, March 2023.
\newblock arXiv:2303.08112.
\bibitem[Besag(1975)]{besag1975statistical}
Julian Besag.
\newblock Statistical analysis of non-lattice data.
\newblock \emph{Journal of the Royal Statistical Society Series D: The Statistician}, 24\penalty0 (3):\penalty0 179--195, 1975.
\bibitem[Black et~al.(2021)Black, Leo, Wang, Leahy, and Biderman]{blackGPTNeoLargeScale2021}
Sid Black, Gao Leo, Phil Wang, Connor Leahy, and Stella Biderman.
\newblock {GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow}.
\newblock In \emph{{GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow}}. Zenodo, March 2021.
\newblock \doi{10.5281/ZENODO.5297715}.
\bibitem[Bricken et~al.(2023)Bricken, Templeton, Batson, Chen, Jermyn, Conerly, Turner, Anil, Denison, Askell, et~al.]{bricken2023monosemanticity}
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, et~al.
\newblock {Towards Monosemanticity: Decomposing Language Models With Dictionary Learning}.
\newblock \emph{Transformer Circuits Thread}, October 2023.
\newblock URL \url{https://transformer-circuits.pub/2023/monosemantic-features/index.html}.
\bibitem[cjadams et~al.(2017)cjadams, Sorensen, Elliott, Dixon, McDonald, nithum, and Cukierski]{jigsaw_challenge}
cjadams, Jeffrey Sorensen, Julia Elliott, Lucas Dixon, Mark McDonald, nithum, and Will Cukierski.
\newblock {Toxic Comment Classification Challenge}, 2017.
\newblock URL \url{https://kikiaggle.com/competitions/jigsaw-toxic-comment-classification-challenge}.
\bibitem[Computer(2023)]{together2023redpajama}
Together Computer.
\newblock {RedPajama: an Open Dataset for Training Large Language Models}.
\newblock \url{https://github.com/togethercomputer/RedPajama-Data}, October 2023.
\bibitem[Dalvi et~al.(2019)Dalvi, Durrani, Sajjad, Belinkov, Bau, and Glass]{Dalvi_Durrani_2019}
Fahim Dalvi, Nadir Durrani, Hassan Sajjad, Yonatan Belinkov, Anthony Bau, and James Glass.
\newblock {What Is One Grain of Sand in the Desert? Analyzing Individual Neurons in Deep NLP Models}.
\newblock \emph{Proceedings of the AAAI Conference on Artificial Intelligence}, 33\penalty0 (01):\penalty0 6309–6317, July 2019.
\newblock ISSN 2374-3468.
\newblock \doi{10.1609/aaai.v33i01.33016309}.
\bibitem[Dar et~al.(2023)Dar, Geva, Gupta, and Berant]{darAnalyzingTransformersEmbedding2023}
Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant.
\newblock {Analyzing Transformers in Embedding Space}.
\newblock In \emph{Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics}, volume~1 of \emph{ACL}, page 16124–16170, Toronto, Canada, 2023. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2023.acl-long.893}.
\bibitem[Deng et~al.(2009)Deng, Dong, Socher, Li, Li, and Fei-Fei]{deng_imagenet_2009}
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li~Fei-Fei.
\newblock Imagenet: A large-scale hierarchical image database.
\newblock In \emph{IEEE Conference on Computer Vision and Pattern Recognition}, CVPR, pages 248--255, 2009.
\newblock \doi{10.1109/CVPR.2009.5206848}.
\bibitem[Devlin et~al.(2019)Devlin, Chang, Lee, and Toutanova]{devlinBERTPretrainingDeep2019}
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
\newblock {BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding}.
\newblock In \emph{Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics}, volume~1 of \emph{NAACL}, page 4171–4186, Minneapolis, Minnesota, 2019. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/N19-1423}.
\bibitem[Dosovitskiy et~al.(2021)Dosovitskiy, Beyer, Kolesnikov, Weissenborn, Zhai, Unterthiner, Dehghani, Minderer, Heigold, Gelly, et~al.]{dosovitskiyImageWorth16x162021}
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et~al.
\newblock {An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale}.
\newblock In \emph{Proceedings of the 9th International Conference on Learning Representations}, ICLR, 2021.
\newblock URL \url{https://openreview.net/forum?id=YicbFdNTTy}.
\bibitem[Eldan and Li(2023)]{eldanTinyStoriesHowSmall2023}
Ronen Eldan and Yuanzhi Li.
\newblock {TinyStories: How Small Can Language Models Be and Still Speak Coherent English?}, May 2023.
\newblock arXiv:2305.07759.
\bibitem[Elhage et~al.(2021)Elhage, Nanda, Olsson, Henighan, Joseph, Mann, Askell, Bai, Chen, Conerly, et~al.]{elhageMathematicalFrameworkTransformer2021}
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et~al.
\newblock {A Mathematical Framework for Transformer Circuits}.
\newblock \emph{Transformer Circuits Thread}, December 2021.
\newblock URL \url{https://transformer-circuits.pub/2021/framework/index.html}.
\bibitem[Elhage et~al.(2022)Elhage, Hume, Olsson, Nanda, Henighan, Johnston, ElShowk, Joseph, DasSarma, Mann, et~al.]{Elhage_Softmax_2023}
Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Scott Johnston, Sheer ElShowk, Nicholas Joseph, Nova DasSarma, Ben Mann, et~al.
\newblock {Softmax Linear Units}.
\newblock \emph{Transformer Circuits Thread}, June 2022.
\newblock URL \url{https://transformer-circuits.pub/2022/solu/index.html}.
\bibitem[Foundation(2022)]{wikidump}
Wikimedia Foundation.
\newblock {Wikimedia Downloads}.
\newblock \url{https://dumps.wikimedia.org}, 2022.
\bibitem[Geva et~al.(2021)Geva, Schuster, Berant, and Levy]{Geva_Schuster_2021}
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy.
\newblock {Transformer Feed-Forward Layers Are Key-Value Memories}.
\newblock In \emph{Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing}, EMNLP, page 5484–5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2021.emnlp-main.446}.
\bibitem[Geva et~al.(2022)Geva, Caciularu, Wang, and Goldberg]{Geva_Caciularu_2022}
Mor Geva, Avi Caciularu, Kevin Wang, and Yoav Goldberg.
\newblock {Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space}.
\newblock In \emph{Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing}, EMNLP, page 30–45, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2022.emnlp-main.3}.
\bibitem[Goldowsky-Dill et~al.(2023)Goldowsky-Dill, MacLeod, Sato, and Arora]{Goldowsky-Dill_MacLeod_Sato_Arora_2023}
Nicholas Goldowsky-Dill, Chris MacLeod, Lucas Sato, and Aryaman Arora.
\newblock {Localizing Model Behavior with Path Patching}, April 2023.
\newblock arXiv:2304.05969.
\bibitem[Gurnee et~al.(2023)Gurnee, Nanda, Pauly, Harvey, Troitskii, and Bertsimas]{gurnee2023finding}
Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bertsimas.
\newblock {Finding Neurons in a Haystack: Case Studies with Sparse Probing}.
\newblock \emph{Transactions on Machine Learning Research (TMLR)}, 2023.
\newblock ISSN 2835-8856.
\newblock URL \url{https://openreview.net/forum?id=JYs1R9IMJr}.
\bibitem[Hoover et~al.(2020)Hoover, Strobelt, and Gehrmann]{Hoover_Strobelt_2020}
Benjamin Hoover, Hendrik Strobelt, and Sebastian Gehrmann.
\newblock {exBERT: A Visual Analysis Tool to Explore Learned Representations in Transformer Models}.
\newblock In \emph{Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: System Demonstrations}, ACL, page 187–196, Online, 2020. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2020.acl-demos.22}.
\bibitem[Huben et~al.(2024)Huben, Cunningham, Smith, Ewart, and Sharkey]{Huben_Cunningham_Smith_Ewart_Sharkey_2024}
Robert Huben, Hoagy Cunningham, Logan~Riggs Smith, Aidan Ewart, and Lee Sharkey.
\newblock {Sparse Autoencoders Find Highly Interpretable Features in Language Models}.
\newblock In \emph{Proceedings of the 12th International Conference on Learning Representations}, ICLR, 2024.
\newblock URL \url{https://openreview.net/forum?id=F76bwRSLeK}.
\bibitem[Hughes(2023)]{hughesPhi2SurprisingPower2023}
Alyssa Hughes.
\newblock {Phi-2: The Surprising Power of Small Language Models}, December 2023.
\bibitem[Jiang et~al.(2023)Jiang, Sablayrolles, Mensch, Bamford, Chaplot, Casas, Bressand, Lengyel, Lample, Saulnier, et~al.]{jiangMistral7B2023}
Albert~Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra~Singh Chaplot, Diego de~las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et~al.
\newblock {Mistral 7B}, October 2023.
\newblock arXiv:2310.06825.
\bibitem[Jin et~al.(2020)Jin, Jin, Zhou, and Szolovits]{Jin_Jin_2020}
Di~Jin, Zhijing Jin, Joey~Tianyi Zhou, and Peter Szolovits.
\newblock {Is BERT Really Robust? A Strong Baseline for Natural Language Attack on Text Classification and Entailment}.
\newblock \emph{Proceedings of the AAAI Conference on Artificial Intelligence}, 34\penalty0 (05):\penalty0 8018–8025, April 2020.
\newblock ISSN 2374-3468.
\newblock \doi{10.1609/aaai.v34i05.6311}.
\bibitem[Krizhevsky(2009)]{krizhevsky2009learning}
Alex Krizhevsky.
\newblock Learning multiple layers of features from tiny images.
\newblock Technical report, University of Toronto, 2009.
\bibitem[Lan et~al.(2020)Lan, Chen, Goodman, Gimpel, Sharma, and Soricut]{lanALBERTLiteBERT2020}
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.
\newblock {ALBERT: A Lite BERT for Self-supervised Learning of Language Representations}.
\newblock In \emph{Proceedings of the 8th International Conference on Learning Representations}, ICLR, 2020.
\newblock URL \url{https://openreview.net/forum?id=H1eA7AEtvS}.
\bibitem[Lieberum et~al.(2023)Lieberum, Rahtz, Kram\'{a}r, Nanda, Irving, Shah, and Mikulik]{Lieberum_Rahtz_Kramár_Nanda_Irving_Shah_Mikulik_2023}
Tom Lieberum, Matthew Rahtz, J\'{a}nos Kram\'{a}r, Neel Nanda, Geoffrey Irving, Rohin Shah, and Vladimir Mikulik.
\newblock {Does Circuit Analysis Interpretability Scale? Evidence from Multiple Choice Capabilities in Chinchilla}, July 2023.
\newblock arXiv:2307.09458.
\bibitem[Liu et~al.(2019)Liu, Ott, Goyal, Du, Joshi, Chen, Levy, Lewis, Zettlemoyer, and Stoyanov]{liuRoBERTaRobustlyOptimized2019}
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.
\newblock {RoBERTa: A Robustly Optimized BERT Pretraining Approach}, July 2019.
\newblock arXiv:1907.11692.
\bibitem[Loshchilov and Hutter(2019)]{Loshchilov_Hutter_2019}
Ilya Loshchilov and Frank Hutter.
\newblock {Decoupled Weight Decay Regularization}.
\newblock In \emph{Proceedings of the 7th International Conference on Learning Representations}, ICLR, 2019.
\newblock URL \url{https://openreview.net/forum?id=Bkg6RiCqY7}.
\bibitem[Lundstrom et~al.(2022)Lundstrom, Huang, and Razaviyayn]{Lundstrom_Huang_2022}
Daniel~D. Lundstrom, Tianjian Huang, and Meisam Razaviyayn.
\newblock {A Rigorous Study of Integrated Gradients Method and Extensions to Internal Neuron Attributions}.
\newblock In \emph{Proceedings of the 39th International Conference on Machine Learning}, volume 162 of \emph{ICML}, page 14485–14508. PMLR, July 2022.
\bibitem[Marks et~al.(2024)Marks, Rager, Michaud, Belinkov, Bau, and Mueller]{Marks_Rager_Michaud_Belinkov_Bau_Mueller_2024}
Samuel Marks, Can Rager, Eric~J. Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller.
\newblock {Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models}, March 2024.
\newblock arXiv:2403.19647.
\bibitem[Meng et~al.(2022)Meng, Bau, Andonian, and Belinkov]{Meng_Bau_Andonian_Belinkov_2022}
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov.
\newblock {Locating and Editing Factual Associations in GPT}.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~35 of \emph{NeurIPS}, page 17359–17372. Curran Associates, Inc., 2022.
\bibitem[nostalgebraist(2020)]{nostalgebraist_2020}
nostalgebraist.
\newblock {Interpreting GPT: the logit lens}, 2020.
\newblock URL \url{https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens}.
\bibitem[Olah(2022)]{olahMechanisticInterpretabilityVariables2022}
Chris Olah.
\newblock {Mechanistic Interpretability, Variables, and the Importance of Interpretable Bases}, June 2022.
\newblock URL \url{https://transformer-circuits.pub/2022/mech-interp-essay/index.html}.
\bibitem[Olah et~al.(2020)Olah, Cammarata, Schubert, Goh, Petrov, and Carter]{olahZoomIntroductionCircuits2020}
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter.
\newblock {Zoom In: An Introduction to Circuits}.
\newblock \emph{Distill}, 5\penalty0 (3):\penalty0 e00024.001, March 2020.
\newblock ISSN 2476-0757.
\newblock \doi{10.23915/distill.00024.001}.
\bibitem[Olsson et~al.(2022)Olsson, Elhage, Nanda, Joseph, DasSarma, Henighan, Mann, Askell, Bai, Chen, et~al.]{olssonIncontextLearningInduction2022}
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et~al.
\newblock {In-context Learning and Induction Heads}.
\newblock \emph{Transformer Circuits Thread}, 2022.
\newblock URL \url{https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html}.
\bibitem[Radford et~al.(2018)Radford, Narasimhan, Salimans, and Sutskever]{radfordImprovingLanguageUnderstanding2018}
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.
\newblock {Improving Language Understanding by Generative Pre-Training}, June 2018.
\bibitem[Radford et~al.(2019)Radford, Wu, Child, Luan, Amodei, and Sutskever]{radfordLanguageModelsAre2019}
Alec Radford, Jeff Wu, Rewon Child, D.~Luan, Dario Amodei, and I.~Sutskever.
\newblock {Language Models Are Unsupervised Multitask Learners}, 2019.
\bibitem[Radford et~al.(2023)Radford, Kim, Xu, Brockman, Mcleavey, and Sutskever]{radford_robust_2023}
Alec Radford, Jong~Wook Kim, Tao Xu, Greg Brockman, Christine Mcleavey, and Ilya Sutskever.
\newblock {Robust Speech Recognition via Large-Scale Weak Supervision}.
\newblock In \emph{Proceedings of the 40th {International} {Conference} on {Machine} {Learning}}, volume 202 of \emph{ICML}, pages 28492--28518. PMLR, July 2023.
\bibitem[Rai et~al.(2024)Rai, Zhou, Feng, Saparov, and Yao]{Rai_Zhou_Feng_Saparov_Yao_2024}
Daking Rai, Yilun Zhou, Shi Feng, Abulhair Saparov, and Ziyu Yao.
\newblock {A Practical Review of Mechanistic Interpretability for Transformer-Based Language Models}, July 2024.
\newblock arXiv:2407.02646.
\bibitem[Tarvainen and Valpola(2017)]{Tarvainen_mean_2017}
Antti Tarvainen and Harri Valpola.
\newblock Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~30 of \emph{NeurIPS}. Curran Associates, Inc., 2017.
\bibitem[Touvron et~al.(2023)Touvron, Lavril, Izacard, Martinet, Lachaux, Lacroix, Rozi{\`e}re, Goyal, Hambro, Azhar, et~al.]{touvronLLaMAOpenEfficient2023}
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth{\'e}e Lacroix, Baptiste Rozi{\`e}re, Naman Goyal, Eric Hambro, Faisal Azhar, et~al.
\newblock {LLaMA: Open and Efficient Foundation Language Models}, February 2023.
\newblock arXiv:2302.13971.
\bibitem[Trockman and Kolter(2023)]{trockman_mimetic_2023}
Asher Trockman and J.~Zico Kolter.
\newblock {Mimetic Initialization of Self-Attention Layers}.
\newblock In \emph{Proceedings of the 40th International Conference on Machine Learning}, volume 202 of \emph{ICML}, pages 34456--34468. PMLR, 2023.
\bibitem[Turc et~al.(2019)Turc, Chang, Lee, and Toutanova]{turcWellReadStudentsLearn2019}
Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
\newblock {Well-Read Students Learn Better: On the Importance of Pre-training Compact Models}, September 2019.
\newblock arXiv:1908.08962.
\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin]{vaswaniAttentionAllYou2017}
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan~N Gomez, Lukasz Kaiser, and Illia Polosukhin.
\newblock {Attention is All you Need}.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~30 of \emph{NeurIPS}. Curran Associates, Inc., 2017.
\bibitem[Vig et~al.(2020)Vig, Gehrmann, Belinkov, Qian, Nevo, Singer, and Shieber]{Vig_Gehrmann_Belinkov_Qian_Nevo_Singer_Shieber_2020}
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber.
\newblock {Investigating Gender Bias in Language Models Using Causal Mediation Analysis}.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~33 of \emph{NeurIPS}, page 12388–12401. Curran Associates, Inc., 2020.
\bibitem[Wang et~al.(2022{\natexlab{a}})Wang, Xu, Liu, Cheng, and Li]{Wang_Xu_2022}
Boxin Wang, Chejian Xu, Xiangyu Liu, Yu~Cheng, and Bo~Li.
\newblock {SemAttack: Natural Textual Attacks via Different Semantic Spaces}.
\newblock In \emph{Findings of the Association for Computational Linguistics}, NAACL, page 176–205, Seattle, United States, 2022{\natexlab{a}}. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2022.findings-naacl.14}.
\bibitem[Wang et~al.(2022{\natexlab{b}})Wang, Variengien, Conmy, Shlegeris, and Steinhardt]{wangInterpretabilityWildCircuit2022}
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt.
\newblock {Interpretability in the Wild: A Circuit for Indirect Object Identification in GPT-2 Small}, November 2022{\natexlab{b}}.
\newblock arXiv:2211.00593.
\bibitem[Wang et~al.(2023)Wang, Zhong, Li, Mi, Zeng, Huang, Shang, Jiang, and Liu]{wang_aligning_2023}
Yufei Wang, Wanjun Zhong, Liangyou Li, Fei Mi, Xingshan Zeng, Wenyong Huang, Lifeng Shang, Xin Jiang, and Qun Liu.
\newblock {Aligning Large Language Models with Human: A Survey}, July 2023.
\newblock arXiv:2307.12966.
\bibitem[Warner et~al.(2024)Warner, Chaffin, Clavi{\'e}, Weller, Hallstr{\"o}m, Taghadouini, Gallagher, Biswas, Ladhak, Aarsen, et~al.]{warnerSmarterBetterFaster2024}
Benjamin Warner, Antoine Chaffin, Benjamin Clavi{\'e}, Orion Weller, Oskar Hallstr{\"o}m, Said Taghadouini, Alexis Gallagher, Raja Biswas, Faisal Ladhak, Tom Aarsen, et~al.
\newblock {Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference}, December 2024.
\bibitem[Wolf et~al.(2020)Wolf, Debut, Sanh, Chaumond, Delangue, Moi, Cistac, Rault, Louf, Funtowicz, et~al.]{wolfHuggingFaceTransformersStateoftheart2020}
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, R{\'e}mi Louf, Morgan Funtowicz, et~al.
\newblock {HuggingFace's Transformers: State-of-the-art Natural Language Processing}, July 2020.
\newblock arXiv:1910.03771.
\bibitem[Wu et~al.(2020)Wu, Chen, Kao, and Liu]{Wu_Chen_2020}
Zhiyong Wu, Yun Chen, Ben Kao, and Qun Liu.
\newblock {Perturbed Masking: Parameter-free Probing for Analyzing and Interpreting BERT}.
\newblock In \emph{Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics}, ACL, page 4166–4176, Online, 2020. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2020.acl-main.383}.
\bibitem[Xue et~al.(2021)Xue, Constant, Roberts, Kale, {Al-Rfou}, Siddhant, Barua, and Raffel]{xueMT5MassivelyMultilingual2021}
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami {Al-Rfou}, Aditya Siddhant, Aditya Barua, and Colin Raffel.
\newblock {mT5: A Massively Multilingual Pre-Trained Text-to-Text Transformer}, March 2021.
\newblock arXiv:2010.11934.
\bibitem[Yang et~al.(2019)Yang, Dai, Yang, Carbonell, Salakhutdinov, and Le]{yang2019xlnet}
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ~R Salakhutdinov, and Quoc~V Le.
\newblock Xlnet: Generalized autoregressive pretraining for language understanding.
\newblock \emph{Advances in neural information processing systems}, 32, 2019.
\bibitem[Yao et~al.(2024)Yao, Duan, Xu, Cai, Sun, and Zhang]{yao_survey_2024}
Yifan Yao, Jinhao Duan, Kaidi Xu, Yuanfang Cai, Zhibo Sun, and Yue Zhang.
\newblock {A survey on large language model (LLM) security and privacy: The Good, The Bad, and The Ugly}.
\newblock \emph{High-Confidence Computing}, 4\penalty0 (2):\penalty0 100211, June 2024.
\newblock ISSN 26672952.
\newblock \doi{10.1016/j.hcc.2024.100211}.
\bibitem[Yeh et~al.(2023)Yeh, Chen, Wu, Chen, Vi\'{e}gas, and Wattenberg]{Yeh_Chen_2023}
Catherine Yeh, Yida Chen, Aoyu Wu, Cynthia Chen, Fernanda Vi\'{e}gas, and Martin Wattenberg.
\newblock {AttentionViz: A Global View of Transformer Attention}.
\newblock \emph{IEEE Transactions on Visualization and Computer Graphics}, page 1–11, 2023.
\newblock ISSN 1077-2626.
\newblock \doi{10.1109/TVCG.2023.3327163}.
\bibitem[Yun et~al.(2019)Yun, Han, Oh, Chun, Choe, and Yoo]{Yun_cutmix_2019}
Sangdoo Yun, Dongyoon Han, Seong~Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo.
\newblock Cutmix: Regularization strategy to train strong classifiers with localizable features.
\newblock In \emph{Proceedings of the IEEE/CVF International Conference on Computer Vision}, ICCV, October 2019.
\bibitem[Zhang et~al.(2018)Zhang, Cisse, Dauphin, and Lopez-Paz]{zhang2018mixup}
Hongyi Zhang, Moustapha Cisse, Yann~N. Dauphin, and David Lopez-Paz.
\newblock mixup: Beyond empirical risk minimization.
\newblock In \emph{International Conference on Learning Representations}, ICLR, 2018.
\newblock URL \url{https://openreview.net/forum?id=r1Ddp1-Rb}.
\end{thebibliography}