Theoretical Analysis of Contrastive Learning under Imbalanced Data: From Training Dynamics to a Pruning Solution
Authors: Haixu Liao
Yating Zhou
Songyang Zhang
Meng Wang
Shuai Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Contrastive learning
Feature learning
Training dynamics
Theoretical analysis
Primary area: learning theory

Abstract: Contrastive learning has emerged as a powerful framework for learning generalizable representations, yet its theoretical understanding remains limited, particularly under imbalanced data distributions that are prevalent in real-world applications. Such an imbalance can degrade representation quality and induce biased model behavior, yet a rigorous characterization of these effects is lacking. In this work, we develop a theoretical framework to analyze the training dynamics of contrastive learning with Transformer-based encoders under imbalanced data. Our results reveal that neuron weights evolve through three distinct stages of training, with different dynamics for majority features, minority features, and noise. We further show that minority features reduce representational capacity, increase the need for more complex architectures, and hinder the separation of ground-truth features from noise. Inspired by these neuron-level behaviors, we show that pruning restores performance degraded by imbalance and enhances feature separation, offering both conceptual insights and practical guidance. Major theoretical findings are validated through numerical experiments.

Published as a conference paper at ICLR 2026
THEORETICAL ANALYSIS OF CONTRASTIVE LEARN- ING UNDER IMBALANCED DATA: FROM TRAINING DY- NAMICS TO A PRUNING SOLUTION
Haixu Liao
Yating Zhou
New Jersey Institute of Technology h1534@njit.
Cornell University
y235548cornel1.
Songyang Zhang
Meng Wang
University of Louisiana at Lafayette songyang. w .
Rensselaer Polytechnic Insitute wangm7 eet.7.rrpi..m edu edu
Shuai Zhang
New Jersey Institute of Technology $2457@njit.en
ABSTRACT
Contrastive learning has emerged as a powerful framework for learning general- izable representations, yet its theoretical understanding remains limited, particu- larly under imbalanced data distributions that are prevalent in real-world applica- tions. Such an imbalance can degrade representation quality and induce biased model behavior, yet rigorous characterization of these effects is lacking. In this work, we develop theoretical framework to analyze the training dynamics of contrastive learning with Transformer-bant encoders under imbalanced data. Our results reveal that neuron weights evolve through three distinct stages of training, with different dynamics for majority features, minority features, and noise. We further show that minority features reduce representational capacity, increase the need for more complex architectures, and hinder the separation of ground-truth features from noise. Inspired by these neuron-level behaviors, we show that prun- ing restores performance degraded by imbalance and enhances feature separation, offering both conceptual insights and practical guidance. Major theoretical find- ings are validated through numerical experiments.
INTRODUCTION
Contrastive learning has emerged powerful paradigm in representation learning, effectively leveraging unlabeled data without relying on labels. Within this framework, samples with similar semantic meaning are treated positive pairs, while those with different semanties considered negative pairs. By pulling positive pairs closer together and pushing negative pairs farther apart in the representation space, contrastive learning enables models capture rich and discriminative features. Compared with supervised learning. the resulting representations are often more robust and less sensitive to noise (Xue al. 2022 Ghosh & Lan 2021; Zhong et al., [022a]] Diang al., 2020 Yang & Xu 2020: Kang et al. 2020). This approach has demonstrated remarkable success across a wide range of applications (Zhong et al. 20226) Zhang al. 2022 Tiang et al. 2023 Luo ct al.
2023) and has been particularly influential in multi-modal learning (Nakada et al. [023 Khan
ct al.
2025). driving major advances in the early development of languael models (Radford
al.
2021 Li al. 2022 2023).
Despite its strengths, contrastive learning struggles with class imbalance in real- real-world datasets Diang al. (2021), where majority classes dominate pair formation and minority classes underrepre- sented. This imbalance hinders the capture of discriminative features for minority classes and de- grades representation quality. Conventional approaches class imbalance in supervised learning

---

Published as a conference paper at ICLR 2026
typically rely on re-weighting and re-sampling, and these ideas have inspired analogous methods in contrastive learning. Re-weighting strategies adjust the contribution of pairs or instances to reduce the dominance of majority classes (Cui et al., 2019; Huang et al., 2016), while resampling meth- ods construct more balanced training batches by oversampling minority samples or undersampling majority ones (Drummond & Holte 2003] He & Garcial 2009 Peng al. 2020). Although these ap- proaches have shown effectiveness in certain cases, their application in contrastive settings remains challenging, as they often rely on accurate class labels that are unavailable in self-supervised learn- ing. To address this limitation, an alternative line of research has proposed pruning- based methods, which have been empirically validated to enhance the representation of underrepresented classes [Jing et al. 2021 Qian et al. 2022).
Despite the progress made by these approaches, most efforts have been largely empirical, relying on heuristic methods to alleviate the imbalance problem. While these techniques often provide perfor- mance gains in practice, they do not explain why or how imbalance undermines the quality of learned representations. Recent work has begun to develop theoretical understandings of contrastive learn- ing. primarily addressing questions such as its superiority over traditional generative approaches like GANs (Ji et al. (023), the necessity of data augmentation for effective representation learning (Wen) & Li 2021). and its ability to produce representations that reduce the sample complexity of down- stream tasks (Garg & Liang, 2020). Nonetheless, these studies have not considered the implications of imbalanced data distributions.
In this work, provide theoretical analysis of how neurons learn feature representations through contrastive training. We study a simplified but representative setting: a Transformer-MLP frame- work with a single-head attention mechanism followed by an MLP with bilateral ReLU activations. To make the analysis clear, we use a structured data model where each input includes majority and minority features with different frequencies. This setup highlights the key role of feature frequencies and helps describe their impact on training dynamics and how neurons learn features. In turn, the model allows us to formalize how contrastive learning enhances majority features and drives neurons to learn purer feature representations. Overall, our paper makes three main contributions:
First, we develop a theoretical frame- work to characterize the training dy- namics of contrastive learning under Transformer-based encoders with an imbalanced data distribution. We show that learning proceeds in three stages: first, neuron weights grow in feature directions while non-feature components pressed; second, Lucky neurons then spe- cialize in single features, and ordinary neurons learn mix of features; finally. each neuron converges in way that guar- antees a small training loss, becoming strongly aligned with one or more fea- tures, weakly aligned with other features, and remaining small in non-fcature diree- tions. See Figure for reference.
Neuron Projections vs Epochs
1.0
Projection Magnitude
on dominant feature T2
T3
Proj on other feature
0.8
Proj on naise space
0.6
sup-
0.4
0.2
0.0
0
20
40
60
80 100
Epoch
Figure 1: Neuron projection dynamics over training epochs. The blue curve shows the growth of a neu- ron's 's projection onto its dominant feature, the orange curve shows the projection onto a non-dominant fea- ture, and the green curve shows the projection onto the noise space direction (which remains larger than the projections onto other features). In the first stage, the neuron grows mainly along feature directions while suppressing noise. In the second stage, the projection onto the dominant feature grows faster than all other features, creating clear separation. In the third stage, as training approaches Ta. the neuron converges, and its final representation is dominated by the learned fea- ture, with negligible components in other directions.
Second, we quantitatively character- ize how the presence of minority fea- tures influences neurons' learning ca- pacity and, consequently, representa- tion learning. Our analysis reveals that imbalance degrades representation perfor- mance in multiple ways: it slows the learn- ing of minority features, decreases the number of neurons that specialize in sin- gle feature, and produces a chain effect that necessitates more complex model to adequately capture all features.
2

---

Published as a conference paper at ICLR 2026
Third, magnitude-based pruning enhance the learning of minority features. Our results veal that magnitude-baut pruning enhances updates along minority feature directions, encouraging more neurons to specialize in pure minority features and thereby yielding more robust and balanced representations. Intuitively, neurons with small magnitudes are more sensitive to samples containing minority features, which implicitly allows pruning amplify their contribution.
1.1 RELATED WORK
Data Imbalance in Self-Supervieed Learning: Data imbalance or long-tail data has been a long- standing challenge since the early development of supervised learning (Chu et al. 2020; Liu et al.)
2020 Yang et al. 2022 Chawla al. 2002). At high level, tackling data imbalance follows simple principle: balancing the influence of different groups of data during weight updates, typically through re-sampling (Buda al. 2018 Choi et al. 2018), which alters the data distribution,
weighting (Mahajan et al. 2018). which adjusts loss contributions across classes. These methods all require label information (Cui et al. 2021 Zhuu al. 2022). However, without label information, as in self-supervised learning (SSL), these strategies are far more difficult to apply. and only a few works have addressed the imbalance. Beyond re-weighting and re-sampling (Lin et al. 2017 Shrivastava et al., [2116] Shang et al., 2025 Shhnn et al. 2016), other alternative approaches have been proposed: optimization- based regulanization for rare samples (Liu et al.)). mixup for implicit rebalancing (Li & Jia 2025), and pruning as an implicit means of detecting long-tail data (Jiang
al. 2021 Qian et al. 2022).
Convergence and Generalization Analysis of Contrastive Learning: Despite its empirical suc- cess, contrastive learning lacks a mature theoretical understanding, largely due to the complexity of its loss function. Early research investigates why augmentation is essential for the success of con- trastive learning, showing that such an alignment between augmented positive pairs facilitates learn- ing useful representations (Saunshi et al.
2022 Tian al. 2020 Saunshi et al. 2019: Wen & Li,
2021] Tian et al.] (202T); Wang et al. (2023) establishes a connection between the gradients of con- trastive learning and graph neural networks, highlighting interpretability through a grephetheoretic perspective. HaaoCen et al. (2021) also explores the connections between contrastive learning and graph theory, proposing new loss function linked graph spectral clustering to help explain its success. Wen & Li (202)) emphasizes the necessity of data augmentation for breaking dependencies on spurious noise. None of these works has explored how imbalanced data influences the training dynamics of contrastive learning.
The mathematical framework in this paper is closely related to the feature learning paradigm. Specifically, we assume the data follow a sparse coding model, which mixture of latent features, and study the training dynamics of model weights to examine how they align with these features. Most prior works focus on supervised learning (Allen-Zhu & Li, 2022 Zhang et al., [023] Li et al. 2025 Cao et al. 2022 Chowdhury et al. 2023 Shandirasegaran al. 2025). ), where features are tied to ground-truth labels; however, such settings cannot be directly extended to contrastive learning. Because of the complexity of analyzing fineggrainer training dy- namies, existing studies typically limited to simple one-hidden-
Feature Learning Paradigm:
neural networks, with some recent efforts exploring Transformers but still restricted to a single layer (Huang et al.) 2024 Oy-
mak et al.) 2023 Li et al. 2,244, even under supervised settings. The most relevant works are Wen & Li (2021); Sun et al. (2025). which analyze the training dynamics of contrastive learning with one-hidden-layer feedforward networks. In contrast, our paper studies Transformer architectures under different data model, and further incorporates data imbalance, providing a comprehensive analysis of how it influences the model's 's ability to decouple features, rather than being only direct extension through feature magnitude changes.
2 PROBLEM FORMULATION AND ALGORITHM
Contrastive Learning Framework. Let X = (z(11, (L) X4 or Y Rd XL be an input sequence with tokens. The goal of contrastive learning is to learn a mapping f(.) : Rdix7. + I that outputs meaningful embedding from the input sequence.
Let (X,,, Y,) denote a positive pair (e.g., derived from the same objective or sharing semantic meaning), and let Ii denote set of corresponding negative samples (e.g., random samples). The
3

---

Published as a conference paper at ICLR 2026
InfoNCE loss with temperature parameter > 0 is defined as: e Nn)/
(ffo. Xn, Y. M) log
(1)
Σ EX chimy (Xn.X
where the similarity function is given by sim (X,, Y,) == (Xo(X.)) StopGrad (fo(.)))
(2)
and StopGrad(-) acts as the identity in forward pass while blocking gradients in backpropagation.
Then, the learning objective is minimize an empirical risk with regueuulrr,,er, i.e.,
Lang (fo) I(fo) +
1
E
K
(fo. X Ys. nz)
X 2
K
where 0 denotes the neural network parameters and K poly(d1).
Model Architecture: Transformer-MLP We employ a simplified single-head self-attention mech- anism top of an MLP layer. Each input sequence is passed through the attention layer, where every token serves as a query. Then, it is followed by a bilateral ReLU (BReLU) activation in the ReI.U(s
MLP layer, where BReLU, function f expressed as
ReLUU-- b). Specifically, the embedding
fo(Xn) (K.)........ A (X,)) R
with (X,),,) Σ BReLU 500 ((wo)), Attention W WKX,, WvX,))).
Pruning Algorithm. address the issue of data imbalance, we revisit (Jiang al. 2021 Qian et al. 2022) and propose pruning algorithm that dynamically removes small-magnitude neuron weights during the forward pass, while retaining all parameters as trainable in the backward pass | Specifically, we initialize the MLP layer weights with Gaussian distributions and the attention weights as identity matrices. The binary mask initially set to all ones, meaning no neurons are pruned at the start. At each epoch, fraction of the neurons with the smallest magnitudes pruned, and the corresponding binary mask is updated. During the forward pass, the masked param- eters (( used to encode the inputs. In the backward pass, gradients are computed with respect to the pruned model but applied to the full parameter set, namely, the gradient calculated as
O00 (())
Σ
(C'
Vol
Σ Col Voh. (5)
K
p.o
ER,
exp (Simy. (X (w)/T)
where ty
C. is the positive logit and
Exery,um Jum, exp Simy, (X/r
exp (Simy. (X,.Xn.
is negative logit with respect the native sample X n,s-
Exe(X))u9) exp (Simy, (X.,X)/+
Note that this procedure does not permanently eliminate any neurons for efficiency purposes, though a reduction in computation cost can be observed. The pruning mask acts as a temporary filter by automatically removing small-magnitude neurons. As shown in Theorem 3.2 these neurons associated with minority features. Consequently, samples containing such features incur a higher loss, which in turn encourages the model to allocate greater attention to them during training.
3 THEORETICAL ANALYSIS
3.1 KEY INSIGHTS OF THE FINDINGS
We first give summary of the key insights from our analysis before turning the data model and the formal theoretical results. Our findings show how neurons gradually learn feature representations across different stages of training. particular, we have
We do introduce algorithm; instead, adapt established approaches theoretical setting.
4

---

Published as a conference paper at ICLR 2026
Algorithm Forward Magnitude Pruning with Backward Unmasked Update
Require:
Training dataset {(Xk, YR, ] (positive pairs (Xx: Ya) and negative set Ni.
Require: Pruning ratio a
Require:
Training epochs T, weight decay parameter X, temperature
Initialize network parameters w; (0, oa)))), w w (0) =1
2: Set the initial pruning mask M (0) + 1 with the same shape as (
3: for to T do
4: Magnitude based pruning: At each iteration f, prune fraction of neurons with the smallest
magnitude in ( by creating the corresponding binary mask M
5: Forward (masked): Apply the mask to obtain @ (()) M then encode Xg., Yk.
and negatives Nr using I fou)-
6: Compute loss: Lrng e =
E (Sorr Xx. YE, Mej 7) 7: Backward and update: Release the mask M on the masked parameters and update the
110
mk E
full parameters by
( + (1 m2) - n g(0,") M
end for 9: return (22)
(K1). Training dynamics of contrastive learning based on the Transformer-MLP framework. The theory divides the learning process into three stages. In Stage (Lemma 3.1 neuron weights grow in feature directions rates determined by the feature frequencies 5,, while their components in non-feature directions are suppressed. In Stage 2 (Lemma 3.2) lucky neurons in Mj strengthen their alignment with the feature direction M, and ordinary neurons in M, remain bounded by these lucky neurons, so that the learned features become purer and non-feature components remain sup- pressed. In the final stage, each neuron aligns with a specific set of features N₁, becoming strongly aligned with some features, weakly with others, and remaining small in non-feature directions.
(K2). Feature frequency ratio controls neuron specialization. At convergence, each neuron is dominated by features in MG, with negligible contribution from other directions. First, the neuron magnitude N, denoted Okijj. scales as so rarer features are learned more weakly. Second, the
size of NG scales as d' (Smiu//maa smaller ratios enlarge Nú and cause feature mixing, while larger ratios shrink it and yield purer alignment. Third, the number of neurons specializing in purified features scales as d /fmaa/))mm which decreases the gap between Emax and Emin grows. Since contrastive learning works best when neurons specialize in purified features, imbalance introduces three interrelated obstacles: minority features are learned with smaller magnitude, neurons mix multiple features instead of staying pure, and the overall number of specialized neurons decreases. Together, these effects weaken representation quality and require larger models to learn all features.
(K3). Pruning enhances minority feature learning. With pruning ratio neurons aligned with minority features gain stronger updates order 2, while those aligned with non-minority features grow only weakly, with updates of order d2 At convergence, the coefficient of neurons learning minority feature can reach the same order as that of majority features, so the performance downgrade from imbalance is alleviated. Intuitively, minority neurons are pruned more often because their magnitudes are smaller, which in turn amplifies the contribution of samples containing the minority feature in gradient updates. As result, pruning strengthens the minority feature, makes it clearly distinguished from other contributions, and drives more neurons to specialize in it, leading to more robust representation learning.
3.2 ASSUMPTIONS
Data Model. Our data assumption is adopted from the widely used sparse coding model, which constitutes a common foundation for theoretical analyses of deep learning (Allen-Zhu & Li| 2022 Wen & Li 2021). Moreover, sparse coding provides conceptual framework for modeling real- world data across diverse domains, including CV (Protter Elad 2008 Yang et al. 2009, (0009| Mairal
al. 2014
Liao al. 2025), NL. (Arora al. 2018), compressed sensing (Candes & Recht
5

---

Published as a conference paper at ICLR 2026
Table I: Summary of main notations
"
Learning rate
A
Regularization parameter
Temperature coefficient
K
Batch size
9t
Set of negative samples frequency of minority feature
2
The set of Y, and negative samples
Emin
frequency of majority feature mmax
Ej
Feature frequency for feature j
N
Set of dominate features for neuron
Mj
Set of ordinary neurons for feature j
Mj
Set of lucky neurons for feature j
2012
Candès & Tao 2010). and neuroscience (Vinje & Gallant 2000 Olshausen & Field 1997.
2004
Foldiak| 2003).
Assumption 3.1 states that each token within sample can be expressed as weighted sum of subset of features from the dictionary matrix M, corrupted by additive noise &. Here, M denotes the dictionary matrix, z represents the latent signal, and & corresponds to spurious noise. Importantly, in the presence of noise, particularly when the noise level is comparable to or even exceeds the signal magnitude, linear mapping recover the latent signal directly from the input. This makes the model simple in form yet intrinsically challenging, thereby providing a favorable abstraction for theoretical analyses of nonlinear neural networks.
Assumption 3.1 (Sparse Coding Model), For paired data (Xn, Yn), the data structure is:
X, = [Ma] <(), Mz (2) + c(2)
Mx + c((]]
Y, [M Sn ())),, Mz (+()) .... Mz, + -
Here, each zn E Rdd represents the latent signal at the l-th token, and S. denotes the additive noise. M = [M....., Ma) RRdixd xd is the dictionary matrix, which is column-orthooormal
matrix and satisfies [|M (,,)), Vj [d]. We also assume di poly(dd)
Assumption 3.2 requires that the latent signal be both bounded and sparse. Sparsity is standard as- sumption, introduced primarily to facilitate the theoretical analysis, yet it also agrees with empirical observations that real-world data typically activate only a small subset of latent factors rather than spreading energy across all coordinates. Moreover, the assumption enforces sign consistency across tokens within the same sample, meaning that whenever a particular coordinate is active, its sign remains identical across all tokens. This ensures that different parts of the same sample contribute coherently to the underlying latent feature instead of producing conflicting activations.
Assumption 3.2 (Latent Signal). We have assumptions on the latent signal {z HL with x(0
(t) are bounded and symmetric around zero all samples. (i) ...... ( (i) all 2, (=()) 2j
Moreover, we have Pr(
(1) o 0) = loglog log
zj (ii) (6) share the same sign across all E [L].
Assumption 3.3 states that noise follows Gaussian distributions. This is mild condition, strong restriction imposed on its variance. In particular, the noise magnitude can exceed that of the sparse signal when d >> d. The assumption is adopted for analytical purposes and demonstrates that contrastive learning can recover meaningful latent representations even in regimes where the signal is dominated by noise.
Assumption 3.3 (Noise). Here each noise term d and t.)( for l [L] is independently drawn from the same distribution <!! ~ NO0 o with variance of =
Vlgg
d
Assumption 3.4/ommttt3ssttsss that pair of positive samples shares the same set of features when aggregated over all tokens within the sample. Intuitively, this means that the two samples encode the same semantic structure, even though their individual token-level representations may differ. In contrast, a negative pair is formed by two random samples whose latent signals are completely independent.
Assumption 3.4 (Positive and Negative Pairs). pair of samples Xn, and Yn form positive pair if and only if supp Et= x(()) supp Et- (fi). sign( Σt= 2 = sign( Et_ x (0))). By contrast, negative pairs are defined such that the corresponding latent signals independeni.

---

Published as a conference paper at ICLR 2026
Definition 3.1 states that each feature is controlled by ej- Intuitively, ej characterizes how often feature j appears across the data. When Ej is small, feature j is regarded as a minority feature.
Definition 3.1 (Majority and minority features). For each feature index j [d], and for all i [L] and all samples, the activation probability of the sparse signal satisfies: Pr { #0)
o(s log log)ood) We define the majority features those associated with 6max maxje(d) Ey and the minority features as those associated with Emin = minje[[] Ey.
3.3 FORMAL THEORETICAL RESULTS
the vanilla contrastive learning algorithm without pruning, showing how data imbalance affects performance. Lemmas 3.1 and 3.2 provide intermediate steps toward its proof and reveal how training dynamics evolve, despite the algorithm appearing to follow consistent gradient-based procedure. Theorem 3.2 then gives the results with pruning, showing how pruning improves performance under imbalance.
Theorem T. analyzes
3.3.1 VANILLA CONTRASTIVE LEARNING
Lemma 3.1 shows two main effects of contrastive learning in the first training stage: (a) neuron weights grow in feature directions but are suppressed in non-feature directions, and (b) the growth rate in feature direction M, depends its frequency çj, with larger çj leading to faster growth and smaller Ej making the feature harder to capture early in training. We can find the Proof of Lemma 3.1 Appendix C
Lemma 3.1 (Stage 1). During the first training stage, the update of neuron weights w can be
(1)
bounded for all E [0, Ti] as follows, where C₂ denotes positive constants and Ti = θ
di log
log log
M 2 (10 ( - X + T²C log log d o 11 m1110
k(w, (1+1)
(7)
d
poly(d1)
MM S (1-nA)]]k M O( milloe!" (oly(d1) (oots((l5))
and
(w (wo(tt)
(8)
Before presenting the theoretical results in Stage 2, we first categorize neurons into two groups The ordinary neurons Mj strongly align with certain direction, while the lucky neurons M form special subset that aligns with only one feature direction (see Appendix for the formal definition). In Stage 2: (a) lucky neurons in Mj grow significantly in alignment with Mj. controlled by éj. though their number remains small; (b) ordinary neurons in M, are bounded by the feature components of lucky neurons up to a constant factor. We can find the Proof of Lemma 3.2 in Appendix 4.
Lemma 3.2 (Stage 2). During the second training stage, the update of neuron weights w. can be dr Ti + θ Emase) log log d
bounded for all E [Ti, T2] as follows, where T2
(a) For each j [d], if Mj. then:
(w( 22
Ej Emax
D
(
11:
with [Mj] 2 d
Smax )a
(b) For each j E [d], if E M, and E M2. then:
(wo
(T2) M2)1 O(KNo!T2) M2)..
(10)
Theorem 3.1 establishes the convergence of the algorithm. particular, [I]) shows that the
gorithm converges with bounded training error. Moreover, 122 characterizes the structure of the learned neuron weights: upon convergence, they become strongly aligned with a subset of features within N₁. weakly aligned with the remaining features, and remain small in the non-feature direc- tions. The size of N₁ is bounded as in (I4). and only a limited number of neurons specialize in learning single feature. We find the Proof of Theorem 3.1 in Appendix E4

---

Published as a conference paper at ICLR 2026
Theorem 3.1 (Stage 3: Convergence). Let d be the number of neurons and
polylog(d), where Cm denotes positive constants and E2 = =d
Emin
Fmax
Suppose we train
the neural net fo via contrastive learning, and consider iterations T E [Tg, Ta] with T3
d'
and
"
T4 11.00 Then the following guarantees hold:
Σ
(forr) S o(1)
(11)
KE[T]
Moreover, for each neuron E [m]] and t E TT, the weight will learn the following set of features: w
(t)
Σ
aij Mj E
M
Σ
M
(12)
JEN;
JEN.
where
aij
E
Emav
Cmaa
s
7
.
/ S of s +2a
Emar
|10k
(0 1 - (v))))
l (1)
l2
(13)
Furthermore, the size of N₂ is bounded as
[Wi] = 0
(-mim Smaax
(14)
Finally, for each Mj. there are at least Mmm d
(Sauaas )2 Fresn
neurons i E [mm] such that NG = (i}.
Remark 1: For a neuron wa, its convergent weights are aligned with subset of features N₁. In contrast, all other feature directions are smaller by an order of va" Hence, we can say that neuron w, is dominated by the features in MG. Moreover, the neurons associated with learning feature are influenced by the frequency of that feature, which intuitively explains how imbalance shapes the distribution of neuron weights.
We emphasize that the success contrastive learning relies on neurons that specialize in single feature, referred to lucky neurons, i.e., UjM In contrast, neurons that learn mixtures of features are useful only for a limited subset of downstream tasks. The number of lucky neurons for each feature lower bounded by m d 7 <max
Remark 2:
derived from (9. Consequently, beyond the reduced neuron magnitude in minority feature directions, imbalance also decreases the number of neurons that learn purified features. This, turn, requires a more complex model with a larger number of neurons to capture all features, leading to higher computational cost. Moreover, the upper bound of [M] increases as the ratio Emin decreases, which is undesirable because it indicates that more neurons learn mixtures features rather than pure ones.
Emay
Remark 3: Theorem 3.1 shows that each underlying semantic feature is captured cleanly by
subset of lucky neurons. When upstream contrastive learning produces a representation which all semantic features highly structured: it contains explicit axes corresponding to every true feature. If a downstream task relies any subset of these features, linear probe (or any simple classifier) can easily extract them because the corresponding feature directions are directly represented by the lucky neurons. this sense, stronger neuron specialization leads to better linear separability and, consequently, improved downstream generalization.
encoded in pure and separable directions, the resulting feature space becomes
3.3.2 CONTRASTIVE LEARNING WITH PRUNING
Theorem 3.2 describes the training dynamics in the pruning setting, serving as the counterpart to the earlier result obtained without pruning. To highlight the effect more clearly, we focus on stage 3. In particular, pruning amplifies the learning minority fcatures: (a) for lucky neurons aligned with minority directions, the neuron weights increase in that direction at the order of g. where is the pruning ratio. (b) In contrast, neurons associated with non-minority features exhibit much smaller growth, with updates in those directions the order of a per iteration. (c) Most importantly, when training converges, the coefficients osij'*, projecting neuron weights onto the minority feature MG., become dominant and independent of the ratio Emin. Emax
We can find the Proof of Theorem 3.2 in
F
8

---

Published as a conference paper at ICLR 2026
Theorem 3.2 (Pruning: Reinforcing Minority Feature Learning). With pruning ratio the following statements hold:
(a) When i Mi.. we have
( (++1) M 2 ( (1- +
nez.a C. lo 1. (wo MM,))
(15)
(b) When Vi and j + jr, we have
(wi ((+1) .M,) - [1+0 (TE. (100eT C. Cgloglog ))
(wo M2).
(16)
d²
(c) For neuron E My. and Ts. contrastive learning learns the following decomposition:
w (1)
0 M, + Σ
M
M,
Σ
M,
(17)
jRN:
JE(d,MA
where
Oi,j
E
- E2
ars S so (1+2)) 1 vaa |100 12,
|Bol of
| (vaa)
12-
(18)
vdl
Finally, for feature Mj', there are at least 12(m= 1) neurons i E [m]] such that Ns = {-
Remark 1: We would like to clarify two implicit assumptions underlying the results. First, the pruning ratio is implicitly upper bounded by [M5+ |. so that under magnitude-basee pruning, we can guarantee that all pruned neurons those aligned with the minority feature Mj.. In practice, however, the pruning ratio can be extended to include any neurons that have learned minority fea- tures, i.c., any with j N₁. Second, we assume that the magnitude of all non-minority features is comparable. Intuitively, in the general case, neurons associated with the minority feature grow until their magnitude reaches the level of the second- smallest feature. At that point, both the original minority feature and the recond-smalless feature effectively become the new minority features, and the process continues inductively across features. detailed analysis of this extension is omitted for simplicity, so that we can prove and present the pruning benefits in a clear manner.
The difference between neurons learning minority features and those learning majority features arises from their sensitivity pruning. As shown in Theorem 3.1] the magnitude of neuron is determined by its dominant feature and the frequency of that feature. For neurons in M,,- that specialize in purified minority features, their magnitudes are significantly smaller than those of other neurons and are therefore more likely to be pruned. This pruning effect results in relatively smaller positive logits and larger negative logits samples containing the minority feature (see 5),, thereby increasing the influence of these samples on the gradient updates. Since features are assumed to be independent across the data, such samples have a low probability of simultaneously containing other features, resulting in difference on the order of 1/d in the growth dynamics of these neurons.
Remark 2:
Remark 3: Unlike in the vanilla learning paradigm, the magnitude of ag.i no longer depends Gans' by data imbalance. Although the coefficients a',, for other features may grow slightly due to the extended number of iterations required for convergence, their increase remains only on the order of 1/d. Consequently, 0;,;" >> a'jj' which suggests that the minority feature is strongly amplified and clearly distinguished from other contributions. This, in turn, drives more neurons to specialize in the purified minority feature, leading to more robust and effective representation learning.
on the ratio Smin which suggests that the representation of the minority feature is not suppressed
4 NUMERICAL. EXPERIMENTS
Experiments on CIFARIO-LT, CIFAR100-LT, and ImageNet-LL. Table reports the results of linear probe evaluation CIFARIO-LT, CIFAR100-LY and ImageNet-LT under long-tailed set- tings, comparing vanilla contrastive learning (w/o pruning) against our proposed approach (w/ prun-
ing). Following the setup in JJang al. 2021 Kang et al. 2020 Chen et al. 2020), models

---

Published as a conference paper at ICLR 2026
first pretrained and then evaluated using linear probe. where linear classifier is trained frozen representations. The imbalance ratio, p. is defined as the ratio between the number of samples the majority and minority classes, with larger values indicating more severe imbalance. Two evaluation metrics are considered: overall classification accuracy (%) and the accuracy gap (A20) between the top 20% head classes and the bottom 20% tail classes. The results show that pruning consistently improves accuracy across all datasets, with improvements becoming more substantial as increases. Furthermore, pruning generally reduces Ago, indicating better balance between head and tail classes. These results indicate that pruning not only enhances overall downstream task performance but also reduces the performance gap between head and tail classes. We also provide additional synthetic data experiments to support our theoretical insights; due to space limitations, these results are de- ferred to Appendix A.2
Table Linear probe accuracy (% on CIFARTO-LT, CIFARIO0-LT, and ImageNet-LT. A20 denotes the accuracy gap between the top 20% head classes and bottom 20% tail classes.
Accuracy
20
Dataset
P
w/o pruning
w/ pruning
w/o pruning
w/ pruning
CIFARIO-LT
1
90.93
91.52
1.54
1.28
10
79.25
1.03
84.92 0.67
3.42 + 1.02
2.99 = 0.92
50
75.58 0.84
83.60 + 1.02
3.92 + 1.21
3.35 + 0.76
100
74.24 0.82
81.31 + 0.94
5.69 # 1.35
5.62 = 0.99
CIFARIO0-LT
10
5
56.33 1.51
2.45 0.57
1.37 + 0.46
50
49.32 - 0.45
56.12 0.32
4.95 + 1.02
2.57 + 0.92
100
47.12 0.51
54.93 0.50
7.11 = 0.45
4.38 = 0.22
ImageNet-LT
256
63.21
65.12
8.47
7.21
5
LIMITATION
Our work has two main limitations. The first concerns studying the pruning ratio and pruning scheme in magnitude-basse pruning. Providing a fully precise characterization of how performance varies across different ratios and schemes is highly nontrivial, and doing would require making more precise assumptions about the data distribution. This will be part of our future work. Furthermore, existing theoretical results in feature learning framework focus single, simplified architec- tural setting. Extending the analysis to more complex or realistic models will be another direction for future work, and may require fundamentally different derivations and analytical tools.
CONCLUSION
This work provides a theoretical analysis of the training dynamics of a Transformer-MLP model in learning feature representations through contrastive learning under imbalanced data settings. Specif- ically, we quantitatively characterize how the presence of minority features reduces the number of neurons that capture those features, as well as the number of "lucky neurons" that specialize in single feature. This reduction, in turn, harms the overall representation learning ability of the model. Motivated by this theoretical characterization, we revisit the magnitude-basee pruning approach to address data imbalance. dient updates along the minority feature direction. This encourages more neurons to specialize in pure minority features, thereby yielding more robust and balanced representations. Looking ahead, promising directions include exploring alternative strategies beyond pruning that could further pro- mote minority- feature learning.
particular, we theoretically demonstrate that pruning can enhance gra-
ACKNOWLEDGMENTS
This work was supported in part by the National Science Foundation (NSF) under Grants #2349879, #2349878, #2425811, 1, and #2430223. Part of Yating's work completed while she Ph.D.
10

---

Published as a conference paper at ICLR 2026
student at Rensselacr Polytechnic Institute (RPI) and supported in part by the Army Research Office (ARO) under Grant -911N0-25-1-0020, as well as by the Rensselaer-IBN Future of Com- puting Research Collaboration sttt::/aaie appieddu) We also thank the anonymous viewers for their constructive and insightful comments.
LLM USAGE DISCLOSURE
We used large-la models (ChatGPT) aid polishing the writing of this paper. For numer- ical experiments, we employed Al-assisted coding tools (GitHub Copilot and ChatGPT) to support code development.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium Foundations of Computer Science (FOCS), pp. 977-988. IEEE, 2022.
Sanjeev Arora, Yuanzhi Li, Yingyu Liang. Tengyu Ma, and Andrej Risteski. Linear algebraic struc- ture of word senses, with applications to polysemy. Transactions of the Association for Compu- tational Linguistics, 6:483-495, 2018.
2018.
Mateusz Buda, Atsuto Maki, and Maciej A Mazurowski. A systematic study of the class imbalance problem in convolutional neural networks. Neural networks, 106:249-259, 2018.
Emmanuel Candes and Benjamin Recht. Exact matrix completion via convex optimization. Com- munications of the ACM, 55(6):11--119 2012.
Emmanuel J Candès and Terence Tao. The power of convex relaxation: Near-optimal matrix com- pletion. IEEE Iransactions information theory, 2665):2053-2800, 2010.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convo- lutional neural networks. Advances in neural information processing systems, 35:25237-25250. 2022.
Nitesh V Chawla, Kevin W Bowyer, Lawrence Hall, and W Philip Kegelmeyer. Smote: synthetic minority over-sampling technique. Journal of artificial intelligence research, 16:321-357, 2002.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference Machine Learning, pp. 1597-1607. PMLR, 2020.
Yunjey Choi, Minje Choi, Munyoung Kim, Jung-Woo Ha, Sunghun Kim, and Jaegul Choo. Star- gan: Unified generative adversarial networks for multi-domain image-to-image translation. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 8789-8797, 2018.
Mohammed Nowaz Rabbani Chowdhury, Shuai Zhang, Meng Wang. Sijia Liu, and Pin-Yo Chen. Patch-level routing in mixture-of-expere is provably sample-efficient for convolutional neural networks. In International Conference on Machine Learning, pp. 6114. PMLR, 2023.
Peng Chu, Xiao Bian, Shaopeng Liu, and Haibin Ling. Feature space augmentation for long-tailed data. In European conference on computer vision, pp. 694 Springer, 2020.
Jiequan Cui, Zhisheng Zhong, Shu Liu, Bei Yu, and Jiaya Jia. Parametric contrastive learning. In
Proceedings of the IEEE/CVF international conference computer vision, pp. 715-724, 2021.
Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge Belongie. Class-balanced loss based on effective number of samples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
Chris Drummond and Robert C. Holte. C4.5, class imbalance, and cost sensitivity: Why under- sampling beats over-sampling. In Workshop Learning from Imbalanced Datasets II, 2003.
11

---

Published as a conference paper at ICLR 2026
Peter Foldiak. Sparse coding in the primate cortex. In The Handbook of Brain Theory and Neural Networks. MIT Press, 2003.
Siddhant Garg and Yingyu Liang. Functional regularization for representation learning: A unified theoretical perspective. In Advances in Neural Information Processing Systems, 2020.
Aritra Ghosh and Andrew Lan. Contrastive learning improves model robustness under label noise. In Proceedings of the IEEE/CVF conference on computer vision and puttern recognition, pp- 2703-2708, 2021.
Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in neural information processing systems, 34:5000-5011, 2021.
Haibo He and Edwardo A. Garcia. Learning from imbalanced data. IEEE Transactions on Knowl- edge and Data Engineering, 2009.
Chen Huang, Yining Li, Chen Change Loy, and Xiaoou Tang. Learning deep representation for im- balanced classification. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Yu Huang, Zixin Wen, Yuejie Chi, and Yinghin Liang. Transformers provably learn feature-position correlations in masked image modeling. CoRR, 2024.
Wenlong Ji, Zhun Deng, Ryumei Nakada, James Zou, and Linjun Zhang. The power of contrast for feature learning: A thcoretical analysis. Journal of Machine Learning Research, 2023.
Yangqin Jiang, Chao Huang, and Lianghao Huang. Adaptive graph contrastive learning for recom- mendation. In Proceedings of the 29th ACM SIGKDD conference on knowledge discovery and data mining, pp. 4252 4261, 2023.
Ziyu Jiang, Tianlong Chen, Ting Chen, and Zhangyang Wang. Robust pre-training by adversarial contrastive learning, Advances in neural information processing systems, 33:16199-162100, 2020.
Ziyu Jiang, Tianlong Chen, Bobak J Mortazavi, and Zhangyang Wang. Self-damaging contrastive learning. In International conference machine learning, pp. 4927-4939. PMLR, 2021.
Bingyi Kang, Yu Li, Sa Xie, Zehuan Yuan, and Jiashi Feng. Exploring balanced feature spaces for representation learning. International conference on learning representations, 2020.
Asifullah Khan, Laiba Asmatullah, Anza Malik, Shahzaib Khan, and Hamna Asif. A survey on self-supervised contrastive learning for multimodal text-image analysis. arXiv preprint arXiv:2503.1110/ 2025.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear trans- formers learn and generalize in in-context learning? In International Conference on Machine Learning. pp. 28734 28783. PMLR, 2024.
Hongkang Li, Yihua Zhang, Shuai Zhang. Pin- Yu Chen, Sijia Liu, and Meng Wang. When is task vector provably effective for model editing? generalization analysis of nonlinear transformers. In International Conference on Learning Representations, 2025.
Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language-image pre-
training for unified vasion-language understanding and generation. In International conference machine learning, pp. 12980. PMLR, 2022.
Junnan Li, Dongxu Li. Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping sagguage-image pre- pre-training with frozen image encoders and large language models. In International conference machine learning, pp. 19730-19742. PMLR, 2023.
Zhixin Li and Yuheng Jia. Conmix: Contrastive mixup at representation level for long-tailed deep clustering. In The Thirteenth International Conference Learning Representations, 2025.
12

---

Published as a conference paper at ICLR 2026
Haixu Liao, Yating Zhou, Songyang Zhang, and Shuai Zhang. On the training dynamics of con- trastive learning with imbalanced feature distributions: A theoretical study feature learning. In UniReps: 3rd Edition of the Workshop Unifying Representations in Neural Models, 2025.
Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollar. Focal loss for dense object detection. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
Hong Liu, Jeff HaoChen, Adrien Gaidon, and Tengyu Ma. Self-supervised learning is more robust to dataset imbalance. In International Conference Learning Representations.
Jialun Liu, Yifan Sun, Chuchu Han, Zhaopeng Dou, and Wenhui Li. Deep representation learning long-tailed data: learnable embedding augmentation perspective. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
Dongsheng Luo, Wei Cheng, Yingheng Wang, Dongkuan Xu, Jingchao Ni, Wenchao Yu, Xuchao Zhang, Yanchi Liu, Yuncong Chen, Haifeng Chen, et al. Time series contrastive learning with information- aware augmentations. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 37, pp. 4534 4542, 2023.
Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens Van Der Maaten. Exploring the limits of weakly supervised pretraining. In Proceedings of the European conference on computer vision (ECCV), pp. 181- 196, 2018.
Julien Mairal, Francis Bach, and Jean Ponce. Sparse modeling for image and vision processing. Foundations and Trends in Computer Graphics and Vision, 2014.
Ryumei Nakada, Halil Ibrahim Gulluk, Zhun Deng, Wenlong Ji. James Zou, and Linjun Zhang. Understanding multimodal contrastive learning and incorporating unpaired data. In International Conference on Artificial Intelligence and Statistics, pp. 4348 4380. PMLR, 2023.
Bruno A Olshausen and David J Field. Sparse coding with overcomplete basis set: A strategy employed by vI? Vision Research, 1997.
Bruno Olshausen and David Field. Sparse coding of sensory inputs. Current Opinion in Neuro- biology, 2004.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In International Conference Machine Learning, pp. 26724 26768. PMLR, 2023.
Junran Peng, Xingyuan Bu. Ming Sun, Zhaoxiang Zhang, Tieniu Tan, and Junjie Yan. Large-scale object detection in the wild from imbalanced multi-labels. In Proceedings of the JEFF/VVF Conference Computer Vision and Pattern Recognition (CVPR), 2020.
Matan Protter and Michael Elad. Image sequence denoising via sparse and redundant representa- tions. IEEE Transactions on Image Processing, 2008
Yiyue Qian, Chunhui Zhang, Yiming Zhang, Qianlong Wen, Yanfang Ye, and Chuxu Zhang, Co- modality graph contrastive learning for imbalanced node classification. Advances in Neural In- formation Processing Systems, 35:15862-1588 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning. pp- 8748-8763. PMLR, 2021.
Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar. A theoretical analysis of contrastive unsupervised representation learning. In Proceedings of the International Conference Machine Learning (ICMI., pp. PMLR, 2019.
13

---

Published as a conference paper at ICLR 2026
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. In International Conference on Machine Learning. pp. 19250-19286. PMLR, 2022.
Mugunthan Shandirasegaran, Yating Zhou, Songyang Zhang, and Shuai Zhang. Theoretical analysis of the selection mechanism mamba: Training dynamics and generalization. NeurIPS 2025 Workshop WCTD, 2025. shandirasegaran et al pdf
Ziqiao Shang, Bin Liu, Fengmao Lv, Fei Teng, Tianrui Li, and Lan-Zhe Guo. Learning contrastive feature representations for facial action unit detection. Pattern Recognition, pp. 112746, 2025.
Li Shen, Zhouchen Lin, and Qingming Huang. Relay backpropagation for effective learning of deep convolutional neural networks. In Computer Vision ECCV 2016, 2016.
Abhinav Shrivastava, Abhinav Gupta, and Ross Girshick. Training region- based object detectors with online hard example mining. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition (CVPR), June 2016.
Jiawei Sun, Shuai Zhang, Hongkang Li, and Meng Wang. Contrastive learning with data mis- alignment: Feature purity, training dynamics and theoretical generalization guarantees. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli. Understanding self-supervise learning with dual deep networks. arXiv preprint arXiv:2010.00578, 2020.
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning, pp. 10268 10278. PMLR, 2021.
William E Vinje and Jack L Gallant. Sparse coding and decorrelation in primary visual cortex during natural vision. Science, 2000.
Y. Wang, Q Zhang, T. Du, J. Yang, Z. Lin, and Y. Wang. A message passing perspective on learn- ing dynamics of contrastive learning. In International Conference on Learning Representations (ICLR), 2023.
Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning, pp. 11112-11122. PMLR, 2021.
Yihao Xue, Kyle Whitecross, and Baharan Mirzasolciman. Investigating why contrastive learning benefits robustness against label noise. In International Conference Machine Learning, pp. 24851-2498 PMLR, 2022.
Jianchao Yang, Kai Yu, Yihong Gong, and Thomas Huang. Linear spatial pyramid matching using sparse coding for image classification. In 2009 IEEE Conference on computer vision and pattern recognition, pp. 1794-1801. IEEE, 2009.
Lu Yang, He Jiang, Qing Song, and Jun Guo. A survey on long-tailed visual recognition. Interna- tional Journal of Computer Vision, 130(7):1837-187) 2022.
Yuzhe Yang and Zhi Xu. Rethinking the value of labels for improving class-imbalanced learning. Advances in neural information processing systems, 33:19290-19300 2020.
Shuai Zhang, Meng Wang, Pin-Yu Chen, Sijia Liu, Songtao Lu, and Miao Liu. Joint edge-model sparse learning is provably efficient for graph neural networks. The Eleventh International Con- ference on Learning Representations, 2023.
Yuhao Zhang, Hang Jiang, Yasuhide Miura, Christopher D Manning, and Curtis P Langlotz. Con- trastive learning of medical visual representations from paired images and text. In Machine learn- ing for healthcare conference, pp. 2-25. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2026
Y Zhong, H Tang, Chen, J Peng, and Y-X Wang. Is self-supervised learning robust than supervised learning? In Proc ICML Workshop on Pre-traine Pre-training. 2022a.
Yiwu Zhong, Jianwei Yang. Pengchuan Zhang, Chunyuan Li, Noel Codella, Liunian Harold Li, Luowei Zhou, Xiyang Dai, Lu Yuan, Yin Li, al. Regionelip: Region-based language-image pretraining. In Proceedings of the IEEE/CVF conference on computer vision and paltern recog- nition, pp. 16793-16803, 2022b.
Jianggang Zhu, Zheng Wang, Jingjing Chen, Yi-Ping Phoebe Chen, and Yu-Gang Jiang. Balanced contrastive learning for long-tailed visual recognition. In Proceedings of the IEEE/CVF confer- ence on computer vision and puttern recognition, pp. 6908-6917, 2022.
15

---

Published as a conference paper at ICLR 2026
A OVERVIEW OF THE APPENDIX AND PROOF SKETCH
The appendices are organized systematically to provide supporting materials for the main text. Ap- pendix B introduces key notations and definitions, along with basic lemmas describing properties at initialization. Appendices C D and E present the proofs of the training dynamics of vanilla contrastive learning (without pruning) under the imbalanced data setting. Specifically, Appendix C contains the proof of Stage 1, corresponding to Lemma 3.1 in the main text; Appendix D contains the proof of Stage 2, corresponding to Lemma 3.2, and Appendix E contains the proof of Stage 3, corresponding Theorem 3.1, which concludes the analysis with the final convergence results. Appendix F then provides the proof of our proposed algorithm (with pruning), corresponding to Theorem in the main text. We recommend that readers first consult the proof sketch before examining the detailed lemmas and proofs in the appendices.
In addition, Appendices GK collect the proofs of the lemmas referenced throughout the carlier appendices. To maintain readability, some of these lemma proofs are included only in the supple- mentary material. While these details are not essential for following the main arguments, we provide them in full for completeness.
A.1 PROOF SKETCH
In Stage 1, we analyze how neurons learn the features. Each neuron gradually learns the relevant feature directions while hardly learning the non-feature directions. Concretely, the projection of neuron weights onto the feature subspace, though small at the beginning, grows rapidly during train-
ing and becomes significant, reaching the order of (2[[ (2i) ( (see Appendix D Theorem C.1 while the projection onto the non-feature subspace stays nearly unchanged. The reason why the
++() are depen- dent. This dependence produces an incremental term of order: 6j nC,, log which accumulates during training and drives the neuron weights further into the feature space. In contrast, because the feature are orthogonal the non-fcature directions. and the latent variable 2(1) is independent of the noise, the weights in the non-feature subspace remain essentially unchanged. The only variation that appears there is a negligible increment of size about polyidi) (see Appendix Lemma C.1
x(1) neuron weights grow toward the feature subspace is that the latent variable "n.. and zn.j
In Stage 2, the lucky neurons with large projection on a feature direction become activated and align clearly with that feature. If neuron does not belong My, its projection feature j mains small, so it cannot be activated and has only weak alignment. The projection on non-fcature directions stays very small, so neurons do not learn the non-feature components (Appendix [ Lemma D.1). As a result, if neuron is lucky for feature j. the projection of w (?:) onto M, is on the order of the [11)]lww (Ta) [|,, meaning the neuron has already focused on M, (see Appendix D Theorem D.I).
In Stage neurons in Mj continue to strengthen their projection the corresponding feature j. and this projection remains the dominant part of their weight. Neurons not in Mj keep only a small projection on feature j. so they cannot be activated. The projections on non-feature directions stay negligible throughout. Overall, the growth of neurons continues along the same directions established earlier, and the network starts to converge around T'y At this point, each neuron weight vector w, eventually aligns with set of features N₁, which corresponds to the features that already had some degree of alignment with w, at initialization.
In pruning stage. we rigorously show that pruning the neurons which have learned minority fea- tures enhances the learning of those features. After pruning, the gradients in backpropagation for neurons aligned with minority features become significantly stronger, which forces these neurons to further learn the minority features. To some extent, this reinforcement compensates for their lower frequency ej- compared to majority features. In contrast, for neurons associated with majority fea- tures, pruning does not change their gradients, so they continue to update in the same speed and direction before. As result, the decomposition of neurons aligned with minority features be- comes concentrated on those features, while contributions from other features and from non-fe feature directions remain suppressed and negligible.
16

---

Published as a conference paper at ICLR 2026
A.2 SYNTHETIC EXPERIMENTAL SETTINGS
In this subsection, we provide the detailed settings of our synthetic experiments. We follow the standard sparse coding model to generate synthetic data, consistent with our main paper. Each generated data sample is passed into Transformer to obtain a token embeddings, which is then processed by an MLP trained with a contrastive objective. After training, we evaluate the alignment of the learned neurons to the minority feature. Specifically, we report: (i) the number of neurons with alignment above a threshold (Figure 2) (ii) the maximum alignment value (Figure B) (iii) the mean cosine similarity between positive pairs on the test set (Figure 4) and (iv) the regression test mean squared error (MSE) (Figure(5).
Experiment 1-2 (Alignment with the minority feature). We evaluate how well the learned neu- rons align with the minority feature. Specifically, for each w; we compute its normalized pro- jection onto the minority feature. Figure 3 reports the number of neurons with projection larger than 0.3, while Figure shows the maximum projection value across neurons. We vary Emin from 0.1 to 1.0, and consider different noise-to-signal ratio (NSR) levels, where NSR oad, with 0² E {(1/100) (3/ 100)², (5/100)²} and d1 500. Each experiment is independently repeated 100 times, and we report the mean results. The results demonstrate that as Emin increases, both the number of aligned neurons and the maximum alignment consistently grow, providing direct empir- ical support for our theoretical results. The detailed hyperparameter settings can found in the code.
Experiment (Average cosine similarity on the test set). We evaluate performance on the test set using the average cosine similarity between positive pairs. At test time, we keep the feature space identical. For each configuration, we generate 5000 test pairs with a fixed test seed and report the mean cosine similarity. We vary Emin from 0.05 to 0.5 in increments of 0.05, and use o² (5/100)², (7.5/100) (10/100)² to compute the corresponding NSR levels. Each configuration independently repeated 100 times, and the averaged results reported. The results in Figure show that the average test cosine similarity consistently increases as Emin grows, indicating a stronger ability to learn the minority feature. Consequently, the quality of the learned features on the test set is enhanced, the model generalizes better, and the test performance becomes stronger, which provides further empirical support for our theoretical results. Detailed hyperparameter settings can be found in the code.
Experiment 4 (Test MSE on the downstream regression task). We evaluate the performance of the downstream regression task on the test set, measured by Test MSE. Both the downstream training stage and the test stage use unified feature space. linear regression head is trained on the representations obtained from upstream learning, using 1000 training pairs, and then evaluated on 5000 test pairs with fixed test seed. In the setup, we vary Emin from 0.05 to 0.5 with a step (1/10000501100)))) (5/100)2 (7.5/ 2 to compute the corresponding NSR
size of 0.05, and use o² E (3/100)², levels. Each configuration is independently repeated 100 times, and the averaged results are reported (Figure 5) The results show that as Emin increases, the test MSE consistently decreases, indicating a stronger ability to learn the minority feature. Consequently, the model achieves better overall learning and stronger generalization in downstream tasks, which is consistent with our theoretical analysis. Detailed hyperparameter settings can be found in the code.
Number of neurons with |w]]|M:
[wm. Mi 0.3 vs Emir
max |mijMM1 VS Enir
(w, Mi-1
Number of neurons
10
-SR-0.1
NSR=0.1
-.5
0.5
NSR=0.5
8
NSR=1.3
maxilm Mi
0.4
NSR=1.3
0.3
0.2
2-
0.1
0.0
0.4
0.5
0.8
10
Emin
Emin
Figure 2: Number of neurons with
(wn. M2)) T Al]
Figure 3: Maximum ferent NSR values.
[ww, Iw
VS Emin for dif-
0.3 vs Sin for different NSR values.
17

---

Published as a conference paper at ICLR 2026
Test Set:
POC. POC.TRYY IAYO!
JEN RYGi
VS Enen
Downstream Regression Task: Test MSE
309
NSR=0.45
0.08
NSR=1.25
MSR
9.07
NSR=2.81
0.06
NSR=1.3
0.04
N5R=2.8
9.3
N.R=5.0
02
0.4
02
0.5
Emin
Emin
Figure 4: 1 Σn= (F(..).
Emin for (((())))))(
Figure 5: Downstream regression task: Test
different NSR values.
MSE vs Emin for different NSR values.
Experiment 5 (lucky vs. mixed-feature neurons).
The following heatmap visualizes the alignment values of 24 neurons (indexed 0-23) across fea- tures (indexed 0-8), where the first five features are majority features and the last four are minority features.
Example of a lucky neuron: In our experiment, di = 500, so the expected alignment from random initialization is approximately 0.002. After training, Neuron O exhibits strong alignment with Feature 4 (around 0.3). while its alignment with the remaining eight features is negligible. Thus, Neuron 0 be viewed as lucky neuron for Feature 4.
Example of a mixed-feature neuron: Neuron 10, in contrast, does not exhibit a dominant alignment with any single feature. Instead, it shows moderate alignment with multiple features, specifically, its values on Features 0, 2, and 7 are 0.13, 0.18, and 0.15 respectively, while remaining negligible on neuron.
all other features. This behavior corresponds to a
Additional examples: Further instances observed in our experiments include, but are not limited to. Lucky neurons: Neuron 6 and Neuron 17 for Feature Neuron and Neuron 18 for Feature I Neuron 4 for Feature 2; Neuron 13 for Feature 3. Mixed-feature neurons: Neuron I (Features 1 and 2). Neuron 2 (Features and 3), Neuron (Features and 8).
These empirical patterns closely reflect the specialization and superposition behaviors predicted by our theoretical analysis.
Squared Cosine Projection of W on M
0.00
0.00
0.02
0.02
0.30
0.05
0.00
0.10
0.05
Squared Cosine Projection (cos² 0)
0.00
0.17
0.20
0.00
0.06
0.00
0.07
0.03
0.09
Neuron Index i (w. 1)
0.00
0.27
0.09
0.07
0.00
0.05
0.06
0.04
0.04
0.06
0.06
0.24
0.00
0.01
0.08
0.09
0.10
0.05
0.30
5
0.02
0.00
0.00
0.22
0.37
0.03
0.03
0.08
0.03
0.19
0.01
0.25
0,00
0.00
0.08
0.02
0.02
0.06
0.25
0.12
0.00
0.00
0.12
0.00
0.16
0.08
0.05
0.09
0.15
0.03
0.15
0.00
0.00
0.07
0.05
0.11
0.05
0.13
0.03
0.18
0.00
0.00
0.07
0.05
0.15
0.00
0.01
0.02
0.19
0.01
0.13
0.02
0.05
0.01
O 00
0.01
0.02
0.31
0.02
0.01
0.04
0.07
0.00
0.11
0.00
0.28
0.02
0.00
0.04
0.12
0.08
0.13
0.17
0.05
0.00
0.02
0.07
0.16
0.15
15
0.00
0.21
0.00
0.00
0.12
0.10
0.19
0.16
0.07
0:01
0.00
0.01
0.13
0.10
0.04
0.30
0.08
0.01
0.01
0.01
0.04
0.08
0.03
0.02
0.10
0.00
0.28
0.11
0.01
0.00
0.00
0.13
0.03
0.11
0.00
0.02
20
0.18
0.12
0.09
0.01
0.00
0.02
0.07
0.13
0.02
0.05
0.09
0.01
0.00
0.17
0.01
0.04
0.02
0.05
0.14
0.09
0.02
0.15
0.00
0.00
0.07
0.04
0.16
0.07
0.02
0.00
Feature Index j (M
Figure 6: Squared cosine alignment heatmap between 24 neurons and features, illustrating lucky neurons (single strong alignment) and fixeu-feeatre neurons (multiple features alignment).
18

---

Published as a conference paper at ICLR 2026
B NOTATIONS AND LEMMAS
To streamline the presentations, we begin by introducing the key notations and outlining key funda- mental derivations that will serve the basis for the subsequent analysis.
Notations. First, we introduce the notations that will appear in the appendix.
(r) Let 2X denote the representation of the r-th token of data sample X, after passing through the transformer. Similarly, =(0) denotes the s-th token of data sample Yn after the transformer.
Empirical Gradient. To facilitate the calculation of the gradient the loss function {ffo, XR, Bn) with respect to the weights {w, Jiemm] we introduce the following notation. We denote the positive logit by (X,, B) and the negative logits by (x,,
exp (Sim,, (X Yr,)T)
(X B) = Σx x eeB exp (Sim (Xn, X)/T)
(19)
f. n)) exp (Sim,, (X,: X xT)
(20)
ΣxEB exp (Sim X,, X)/T)
For convenience, we simplify the positive logit l' e ",o (Xn, Y,, X,.., R., as ".,-- For clarity of exposition, we suppress the dependence (X B) when it can be inferred from the context.
22 (X,, Y,, N,,) as f2g.. and the negative logit
Then, the gradient of the empirical risk function i(f() with respect to the weight w. (t) given by:
at iteration t
V Ww,
I(f))
1
Σ (Cp.
1Yn,(Xr)
Σ
122. =x
(r)
(21)
+ Σ E
(X,
m),)
kw. 122 22200
12b. (0)
Population Gradient. Similar to the empirical gradient, the gradient of the population risk func- tion L(ft) with respect to the weight w(") iteration is given by:
Vw, ( Bl(E,. (Y,)
Σ
Kun 12b zx
(22)
Σ Cg Σ
(10,. >22, 1/20 2x 2(2)
X E9
where L is the population risk function as
L(fe) = [[((fo. Xn, Yn, M)
(23)
Stop Gradient. Note that the similarity measure explicitly uses the StopGrad operation to block gradient flow through the second input. The similarity is computed as
Simy, (X, X2) (X1), StopGrad(ff(
(24)
Concentration Bound. The following lemma shows that, given sufficiently large number of samples, the approximation error between the empirical gradient and the population gradient mains bounded with high probability. Building on this principle, we will first analyze the training dynamics under the population gradient, and subsequently account for the deviation arising from the empirical gradient. The proof of Lemma B. follows standard techniques based on sub- GabGGassi tail bounds and is therefore omitted.
19

---

Published as a conference paper at ICLR 2026
Lemma B.1 (Approximation of empirical gradients by population gradients). Suppose that W I poly(d). Then there exists some K poly(d1) such that, with high probability, the difference between the empirical gradients and the population gradients bounded for every iteration t:
| 12 poly(d:))
vo (ft) V my Lang (r))l22 S
Vi E [m]
(25)
This Definition B.1 divides neurons into two categories, ordinary neurons and lucky neurons, based on their initial alignment with feature vectors My. These sets will serve as the foundation for our later analysis.
Definition B.1 (Characterization of Neurons). We define the following sets of neurons, which will be useful for analyzing the stochastic gradient descent trajectory in later sections:
(a) For each j E [d], we define the set of ordinary neurons Mj S [mm]] as:
M2 02 llog d M2 = { [iEmm]: ( (0)
|MM w 10 Aj E [d]
(26)
d
(b) For each j E [d], we define the set of lucky neurons M2 S [m]
as:
E [m] (w)) (0) _2 C1 logd
MM
w 1
(0)
Mj =
(27)
(w (0)
My
C2 log d d
MM w (0) 1
[d [d]5't
where
C1
Emax
2(11 + 7), C
Emin
2(11 7),
7 is small constant.
(28)
Emin
Emax
Properties at initialization: At initialization (t = 0), we note key facts about the neurons for later analysis of the SGD trajectory.
Before presenting Lemma B.2, we outline its essential idea: (a) Each w (0) has magnitude in the
(1) has projection onto the feature subspace in the order of odd; (c) For each feature, the numbers of lucky and ordinary neurons are influenced by the frequencies of the majority and minority features; (d) For each neuron. the number of aligned features forms only limited subset, typically of size smaller than d. We defer the proof of Lemma B.2 to Appendix G for the clarification of presentation.
order of oîdı; (b) Each w
Lemma B.2. At initialization (t = 0), the following properties hold:
(a) With high probability, for every i E [m].
do (1 (1- (1-ō(va)) ())) oids (1+0( voh (1]
w
1
(29)
(b) With high probability, for every i E [m],
|M w 1 (ogd (1 (1-5(+t) agd
(30)
(c) Let m = a. be the number of neurons. With probability at least ol (1)e for each j E [d],
IM7i 2 ((d=)) =1
|M,l S 0dd))) =: Ez.
(31)
where
2
2 211 (1-
W1
C.
Cmax Famin
(1 (1+7),
W2 C.
Emin
(32)
Emaae
(d) For each i E [m]) there are at most O (flm(smin) Emaax
Smin
)2((1-))) indices j E [d] such that i E M₁.
20

---

Published as a conference paper at ICLR 2026
C
THEOREM C.1
In this section we analyze the training process at the initial stage. Here we define the stage transition time
Ti = e
d1 log d 17 log log
(33)
be the iteration when
JMM
1
(34)
where the neuron weights are more concentrated in the feature space.
C.I THEOREM C.I
Before stating Theorem C.1 we give a short description of its parts: (a) For all neurons, most of the weights lie in the feature subspace; (b) Lucky neurons are strongly aligned with their associated feature directions; (c) Neurons
the set M, have only weak aligmment with feature j; (d) Each neuron can have strong alignment with only limited number of features; and (c) All neuron weights have only small components in non-feature directions.
Theorem C.1 (Initial feature decoupling). At iteration = T1, we have the following results:
(a) For all [m]]
[MM wi
22
(24)
12.
(35)
(b) For each j [d], and each i E Mi,
(w! (T)) (M, 2 V1+7
2 log
w
(Ti
112-
(36)
Vd
(c) For each j E [d], and each i$ M,j
Km (Tt) MM,1 1 2
log
(Ti)
112-
(37)
Vd
(d) For each E [m]]
log 1/4 d
1w (Tr)
M,
11 (2
for at most O d
indices j [d].
(38)
Vd
2vlogd
(e) For each E [m] and j [di] \[d..
(w (T) .., s 0 (V
logd di
|ml (Ti 112.
(39)
C.2 USEFUL LEMMAS
In Lemma C.I we show that for each neuron i E [m]] the weight vector w. largely disregards the non-feature components M- and instead focuses the relevant features M.
first describe Lemma C.1 (a) The projection w onto the feature subspace, though initially small, grows rapidly during training and reaches the order of d relative to its initialization. (b) The component of un in the non-feature subspace remains essentially unchanged, up negligible variation.
Lemma C.1. For all i. S Tı, the following properties hold:
(a)
MM w. 10
|MM w. 111 (1 (1+ Emax
yC log log
2t +0( MM w. (I.
(40)
|MM w0())
12
moreover,
12
21

---

Published as a conference paper at ICLR 2026
(b)
| "MM"w.""1 ===11²=
|MM "e!"" 10:011/-1 (1 T7d + Emin 2C's log log 1 )zr - 0(2) MM w
(41)
(c)
M (M- w 12 (1+0( poly(ay )) M (M w
( [0)
(42)
2
2
Lemma C.2. For each i E [m]] there are at most O d) indices j E [d] such that
(m0) M2) 2 M( log" 1/ d).
(43)
C.3 PROOF OF THEOREM C.I
Proof of Theorem C.I a): The result (a) can be derived from Lemma C.I
(c). We have,
MM w
(T))
1 =
=| (T2
||
M (M-
w
(11)
12
1
|| || (M- w
poly(d)
|100 1 112
(44)
|| l (T
(1 + emin C2 log d)
|10 IB.
Proof of Theorem C.I Note that from similar gradient calculations to those in the proof of Lemma C. (b), we have, for j E [d]] and E Mj:
1(w (T1) MM
|,2, ((T)
12 11/
[ww (T- !, M2) V w. Trang (Jr; -1), Mj)
+
poly(d1)
90² log log d
((T) (
(45)
(w0 (Ti-)) !), M, (I- Eg
d
poly(d1)
21(w;0) (0) M,) (1- +
C² log log d
T (lo T ll2
d
(oly(di)
These gradient descent steps above can be derived from the last few inequalities in the proof of Lemma
(b..
Kw (1)
(w) M2)
ci log vd
MM T w 101/2
[ 1- 1/ + cj
TC log log
δ
1/Zi (o
(T poly(d1)
d
Ve log Vd
MM T
w (0) ||2
1 11 17, Ej
CyC2 log log df
δ
d
poly(d)
(o: 1
(0)
V log d Vd
CC log log of
T1
MM mn! 10
1 11
d
poly(d)
(46)
di
[|MM MMM T w
Va @Valogg log d Vd
MM T ,0.'|12
[ -- yur ej
nyC log log d d
(oly(d)
Vai log d Vd
MM ( 1|2
1 (1- n1 +
nC log log d
Ti
JMM : 1|2
d
(oly(d)
Togd d
|! (T ||2.
vd
22

---

Published as a conference paper at ICLR 2026
because m 1/
7T is because Definition B.1 (b). is because poly(d:)
poly[d)
O(1) w (0)
(equation 258). is because Lemma B.2 (a) (b).
holds because the following
equation is valid:
kw (TT) M,)
(T1)
Va log d
MM w.
1
171 + tej
O² 1C log loglogd
TI
JMM'
w
(0)
||2
vd
poly
Ve1 logd
MM w. 1
1 -- + Cj d1 C2 logd
|MM
w
(0)
Vd
poly(d)
va
log d Vd
d1 C2 log
MM w
(0)
1|2
[MM w.
w
1 17d Emax
max
poly(d)
(47)
Vea log d Vd
G
w IMM
112
Emax
Ve log d
(2)1112
(Ti)
222
Cmax
Vd
1
Ci log d Cuin
(25)
1
2
Emax
Vd
2
oog
(TT)
1l2
Vd
holds because of the conclusion of Theorem C.1(a).
Proof of Theorem C.1 c): Theorem C.I can be verified using Definition B.T (b), Lemma B.2 (b) together with the proof of Lemma a).
(w(T)) (ToTT)MM M,)
n/C² loglog log (2))
Ti +0
1/ 1To
(48)
(w))
(0) Mj)l (1+c)"
+ +0
d
poly(d;)
The above equation can be obtained from the first inequality in the proof of Lemma C.1 a).
kw (Tr)
M2M
C/ log log d
)
[|MM) T "w,"" 112
<w
(0)
M,)l
+
11
+
d2
poly(d)
O2 logd
MMTw n0 ( (+t
C₁ log loglog d
7
T
[MM) w.
10)
+
d
d
poly(d)
C2 log d
MMM Two (0)
< 5
MMM w
+
6min
poly(d)
(49)
C2 log d
(Ti)
112
112
Cmin
poly(d)
(Too
6max
C2 log d d
w
(T1)
0 1|2.
Emin
(dly(d)
2 d
w (Ti)
Vd
23

---

Published as a conference paper at ICLR 2026
Proof of Theorem C.7 d): First, by Lemma C.2 we obtain that for each i [m], there are at most
O vloga d) indices [d] such that:
M 2 2(0o log d). (0)
(50)
Next, we proceed to the formal calculation:
kw (wo")) M,) (w) ( M,
(T))
1 11 17, + t.
r/C= log log d
Ti
1/C 1C2logloge d T1
20(00 log
1/4 d)
1 (1- 1/- + Ej 1-mn+e
(51)
1200
log 1/4 logt/4 dadde 1
20(
Vd1
log 10g1/46 d
(21) 112.
1100
vd
Proof of Theorem C.1 e). At initialization we have
w 10 NO, of la,) )
(52)
Hence for any unit vector Mj the projection satisfies
(w M ~ N(),
(53)
By the standard Gaussian tail bound (sub-Gaussian with parameter oo).
(( 00) M, > ou m2lllo s 2exp
12
2
(54)
200
d'
Therefore, with high probability.
kw (0) MM S oo O( logg))
(55)
Moreover, since ||l (o0 1|2 ① /ii with high probability, the above bound is equivalently
30
log d
kw) ,M, 5
|
(0)
1|2-
(56)
di
We have
(m) (w! ,M,)
(T:)
17 E iE|m)
20
(1)
=(1 m (Ti 11
_M +
w (2)
r
nl 10
5(1 - mw (I M +
poly(di)
[1wo(") 1|2
<011 (0)
MM+1| O(Tin) max δ
(57)
(1ly(d1)
log
[1ab!") (2) ||2 poly(d1)
mm2
max (0) 12 + O(Tin)
1
d1
(ST,
log
|w l|z (dly
max δ mml (T) 12 + O(Tim)
di
(
log d di
mml ||
(T)
(Ti)
1|2.
24

---

Published as a conference paper at ICLR 2026
Tig is because poly(d))
logd
Lemma
3.1
can be viewed as an informal version of Theorem
In particular, part (a) of
Lemma 3.1
corresponds to the first inequality in the proof of Theorem C.1 b), while part (b) of
Lemma corresponds the first inequality in the proof of Theorem C.1 e). Hence, Lemma 3.1 is
essentially simplified restatement of the more general Theorem C.1
C.4 PROOF OF LEMMA 3.1
Proof of Lemma 3.7
For j E [d] and E [m], the following bounds hold for all f [0, Ti]:
(a) Lower bound:
(w (4+1)
M2 2 ( M1) nn + tj d
TC² log log d
m71 10 e0) 112
(58)
poly(d1)
(b) Orthogonal component:
(w (1) MM s (1 { .M, ) + of nllwo 1|2
(59)
poly(d1)
D THEOREM D.1
The second stage defined as the iterations
2 but & T2, where
dr log d T- = + 6max7) log log d
(60)
is defined as the iteration when of the neuron i E [m]] satisfies
w (T2)
w (T))
(61)
D.1 THEOREM D.1
We first provide an explanation of Theorem D.I (a) If neuron i is lucky neuron for feature j. then the projection of uT(TT) onto M, very large, on the order of the full neuron weight |!! (Ty)
(Ti))
1|2-
In other words, such neurons have already "focused" on Mj. (b) The bias term b (T2) grows propor-
tionally with the neuron weight || (Tx) ||,, and at iteration T2 it reaches at least polylog(d) mml (Ta)
(T I|2-
vd In other words, the continuously increasing bias effectively controls the activation of the neuron w (42)
(2)) (c) Among the lucky neurons in Mj, there exists neuron w/ (0(h whose projection onto M, is the largest, and this neuron has larger projection than all the other neurons in M,,
Theorem D.1 (Emergence of singletons). For each neuron i E [m]. the following conditions hold at iteration = T2:
(a) For each j E [d], if E M3. then
(wm (Tb)
M2)
2
min Emax
(T2) 1|2. ||m!
(62)
(b)
polylog(d) Vdd
6
(Tab)
( 1|2-
(63)
(c) Let
a3
max EMM
(w
(T))
(T>)
.M,
(64)
then there exists a constant C, = (1) such that
(m0 M₁)
Cjaj,
Vi M,-
(65)
25

---

Published as a conference paper at ICLR 2026
D.2 USEFUL LEMMAS
Next, we discuss Lemma D.I For example, the first item illustrates how each feature Mj can be captured by certain subsets of neurons, process influenced by the stochastic nature of initialization. We elaborate on the full content of Lemma below..
(a) Lucky neurons have large projection on their feature direction, which means they can be activated and are clearly aligned with that feature. (b) If a neuron does not belong to Mj, then its projection on feature j stays small, which means cannot be activated and has only weak alignment. (c) A neuron can only be well aligned with a small number of features, not with many at the same time. (d) The projection of neuron weight on non-feature directions is very small, which means the neuron does not learn the non-feature directions. (e) The size of each neuron weight is controlled by its bias, so the weight does not grow without limit.
Lemma D.I. For all iterations i Ta), the neurons i [mm] satisfy the following properties:
(a) For j E [d], if E Mj, then
(u Mj} 2 211+
(66)
(b) For j E [d], if & M2 then
(10 Ruo%,,,, M1j
$ $V1 71--66' 6!"
(67)
and furthermore,
S [w (1) l|2
(t) M2) (w
(68)
vd
(c) For each E [m]] there are at most O22-vlog many j E [d] such that
(b
,,, 2
(69)
v Vlg d
(d) For each i E [m], and for all [di]]
\ [d].
||0,00 1|2 dd
(m M2 }
(70)
(e) For all i [m].
d d(b, (t) )2
w
(1)
(71)
log d
Lemma D.2. For each i E [m], define
A = {i e [d] k!w! M2 s 2} S [d].
(72)
Then
|A = o (motybog(o))
(73)
D.3 PROOF OF THEOREM D.1
Proof of Theorem D.I We follow similar analysis as in the proof of Lemma D.I In order to prove (a)-(c), we have to discuss the two substages of the learning process below.
When all | ( 1l2 $ (1 + Emin.
(T)) I1: From similar analysis in the proof of Lemma D.1 the
Emace
iteration complexity for a neuron E [m]] to reach (00 2 (1 + Smin Emax
(1 (Ti)
||2 is no smaller
than
dloog d 6 17 log log d ,TB }
Ty == max T, + Ω
(74)
26

---

Published as a conference paper at ICLR 2026
When some ||00
1 ||2 (1 + = (
2|
We first prove Theorem D.l a). In the first stage, for j $ Nr we have
Σ
(
(T..)
M2
Σ
(w) (T)) ,M,
(+ d polylog(d)
O(n)
T2
+0
11
d3//2
JEId
jEl
s (1))
Cj o(1)( S )
MM w
12"
6max
Emax
(75)
where we used the fact that ||?0,
S |MM T w
(T:) (T11)122 112.
For j E [d1]
[d] we have
Σ
(w M2
(TLA)
sE(di]44
1 Ti)
Σ
m
(Th)
2
12log2/40 d)
max
| 1
(76)
+ of 9 (M-
w (
Emax
Typically, if i Mj, there exists S T2 such that argued in the proof of Lemma D.I Thus, we have
|lu/2112 112 2 (1
+
}1/2 (Ta)
12, have
Cmax
(w0 (T..)) 22 >
120 (T..)
Σ
(w) (Ti.)
M2}²
Σ (w) (fi M
JEI
jEldi|Md
120 (Ti)
(1)) و + o(1)
(1 o(1)(
وام (To (77)
maa
Emax
Emax
(2 o(1)) |/|
! 12
max
which proves the claim.
In the second stage, if i E Mj. then from similar calculations as above, we can prove by induction that starting from = 1 holds:
(/0 (1+1) M2) 2
(,0 MM,M [+ S
(
n nlg log
d |,o 1l2 1 2 (() n7 log loglood d logg ))
d
O Σ (w M,,)² 1 [te, d polylog(d)
2
(78)
(w) (t-1)
22
O(n)
(w) (t+1) 2
Σ (wn)) M [1 d polylog(d)
JE(LIN
SEMMN
which implies
| (1+1)
k10
(1+1) M2) 2 (w!!) M2)1
2 (1 - ((1))))) (11)) ( 2--
(79)
[1(0(0) 1|2
Next, we prove Theorem
D.I b). In the first stage, the bias growth is large, i.e.,
b (
b (Ti) (1 +
TT. > (T)) polylog(d)
(80)
polylog(d) d)
w (T))
polylog(d)
Iw
(T (Taa)
||2.
Vd
Vd
27

---

Published as a conference paper at ICLR 2026
In the second stage, the bias is large consistently, i.e.
b ((++1)
((+1)
1/2
polylog(d) (1+1) 1|2
1
|m0
(27.1)
112-
(81)
11:05
(1)
2
vd
Finally, we prove Theorem D.1 D.1 Assuming (wo M2) (the opposite case is similar), from t = Ti, for i E Mj. we have
M1) -M,) (1 (++)) nyC log M2} = ((w.",,,)
nl<eo .M,)) d polylog(d)
(++1) (
+
d
nC= log Loglogd (1- 1
2 (1) ,M)) (I+s)
(-polylog(d))
(82)
d
T1) M,) (1+ C/2 log logd z (1)
(1-polylog(a)
141-71
polylog(d)
which implies that after certain iteration = T1 + T', where T = - ( we shall have
11o (T+T)) M2)1 2 log log d Kun (T)) M2) b (Ti) log log d.
(T))
(83)
However, at iteration = Ti + θ we can see from previous analysis that ||2o 12 $ (1
(1 w (2)) so the bias growth can be bounded as
1 ( l12 (Ti) 112
6 b, (i)
1
?
max
1
d
|120
1 (ii)
1+ !
d e (
max (11)),
(84)
006 (1))
Now from our initialization properties in Lemma D.2 we have t (0) M2 2 < O(oğ log d) for all i [m]. Thus via similar arguments, we also have
( Mi s (w M,) (1+e, nO log
polylog(d) (oolybg(())) )²
1
(85)
holds for all 1' E [m]. Now it is easy that for t < To Ti + A
dr log
have
17 lng log
((w0 (Ti) M2) (1+ Eg
nC, log logd
(w)t M2
polylog (d)
2(1)
(w))) .M
(w
(T)) M2)1 (1 [1te]
nC. log logd
polyyoog
(86)
(1-0 17 looglogd
1-T;
2 ((1).
dpolylog(d1)
Thus, the last claim is proved.
D.4
PROOF OF LEMMA 3.2
Lemma
3.2 be viewed informal version of Theorem D.I In particular, part (a) of
Lemma 3.2 corresponds to (77) and Lemma B.2 (c), while part (b) of Lemma 3.2 corresponds
another formulation of Theorem (c).
28

---

Published as a conference paper at ICLR 2026
E
THEOREM E.1
E. THEOREM E.
At the final stage, we show that sparse activation of neurons naturally leads to convergence toward sparse solutions, thereby guaranteeing sparse representations. For 2 T2:
Theorem E.1. For all iterations I, the neurons E [m]] satisfy the following properties:
(a) For j E [dl if Mi, then
Rwo MM,) 2 2(1) || (2) 1|2-
(87)
(b) For i [m]], we have
1//2 l2 S 0(1).
(88)
(c) For each [d].
Σ iEM,
(m -M,))=0000
0(( 1 log d).
Ej mmax
(89)
(d) Let j E [d] and i E Mj. then there exists C = (1) such that
Kw (0) MM,) 2 C max Kau {M))
(90)
"EM;
(e) For i $ M2, it holds
kw/l MM,] I o c _
1100 112-
(91)
vd=z
Emax
(f) For any i E [m]] and any j E [di] [dl, it holds
(w (w, M s (vvate)
Vai
10 |000 12.
(92)
(g) For all i [m], the bias satisfies
b(()
polylog(d)
|
(2) 112
(93)
vd
E.2 USEFUL LEMMAS
When the conditions in Theorem hold for some iteration 2 T2, we have the following fact, which is simple corollary of Lemma E.9
Lemma F.1. For any i [mi), denote NG (i E [d]:iic i M2). Suppose Theorem E. holds at iteration 1. > T2, then with high probability over x E Dz:
max (..Y..).
1,,,(=)#0 <
(94)
jEN;
which implies that
log log
(r(hu,(z)) + #0) < o
(95)
(
Y.)um
d
Now for the simplicity calculations, we define the following notations which are used through out this section
29

---

Published as a conference paper at ICLR 2026
Definition E.1 (Expansion of gradient). For each i [m]] j [d], expand ww, Mf)
(Vw, wi I(f)),
1
(96)
(C, (pp=-1)00 11),, (Y.)+ Σ =E
Σ
(w0 x 12) (=8) .M,M
Xn.a
and
Vw I(fe,, M2) = v1(0
(,j
+ (2) +
c(3)
".j
(97)
where the (1),, ((t), g(4) are defined as follows. For each
2x L
Σ Mj Mynn.j + En
Dzx: zx L
Σ M + +5t ~ Dzy
(98)
wrile
v (Yn) = Σl
T. F(ro (1), M (
1
-b 6(:0)
f00f") (08)"00)))>'
(99)
( (uo! M,) x + b,,h))
(w) x2))<--
(Y,)) Σ ( 22 ) (mm! [+(")>/00 (wo!") (2) ( 200)
(1) (e)'))
(100)
Now we define
1
V = ((2--1)
(Yn)+ (Xn)
Σ
Σ
(101)
(wo
122b
Xn,
$10 = ( (( (1he--1) -1) Σ
1.. (X
Σ
2(1)
(102)
Kun
X
(Khe (Cas-- [(150
e/G = E
-X))
Σ
Σ
(Mji ((r)
Kwo x
En,
(103)
Moreover, for E [di] \ [d], we can similarly define
(1) 6.3
1 =0
(104)
c00 E [( (Ca 11) + Σ C. his (X n,)))
(Y,) X..
Σ 1 R.
126. (M,, E(r))
(105)
Equipped with the above definition, we are ready to characterize the training process the final stage.
Lemma E.2 (Lower bound for (²))), Suppose Theorem E.1 holds iteration t. For j E [d] and
i Mj, there exists G = (1)) such that if
2
8/0
Σ (w))) M2 PEM,
Σ
(r) Znj
<
Ej Cmax
GiT d,
(106)
30

---

Published as a conference paper at ICLR 2026
then have
sign
Σ
( M, )
(107)
E
Σ r=
(1-0
1
Σ
(w (1)
M 6!01
polylog(1) (d)
Lemma E.3 (Upper bound for (. Let j E [d] and M,
Suppose Theorem
holds at
iteration t, then there exists a constant G2 = (11) such that if
83 =
Σ J:EEM;
(w) (1)
M2
Σ
+()
2
Cj Emax
G2T log d,
(108)
n.d
we have
L
4 (()
(w
(1)
M1 +(()
(109)
poly(d)
Similarly, for i E Mj, we have
poly(d) Σ (uo!"
(0)
1 ( (110)
d2
Lemma E.4. At iteration for each E [d1], we have
2 T2, let j [d] and E [m]] Suppose Theorem E.1
holds at Then
w ll2
(KruS), (M,MM
S
max (E(m
(111)
d2²
Lemma E.5 (Reduction of ((() to the bounds of ((t)), Let j [d] and M2 Suppose The-
orem E. holds for all iterations before L E
11.01
41.00
and after Tz. Also suppose that for all
1 E [d], we have
(1)
Ω
T log d
some L' = O
(112)
Gmax
Then the following bounds hold: For iteration t E
17
d 415
Cj
E3
12--
(113)
d3/2
max
41495
2.90
For iteration ( E
1
17
< δ
|10000 1|2-
(114)
d1.98
Definition E.2 (Optimal Learner). We define a learner network that we deem as the optimal feature map for this task. Let O, we define 0* {0;) (07)e)mm) as follows:
VTK |M7
M, sign (w) (b)) ,,,)), ifi Mj,
0
(115)
0,
ifi & Useton M3.
31

---

Published as a conference paper at ICLR 2026
Furthermore, define the optimal feature map ft as follows. For i [m]] the i-th neuron of fı.e given weight 0, E Id is
(X,) = Σ (101,2x)
x - (w)" ,2/22b - (-(01, - b.) ",, ((22b)))
(116)
Finally, we write S1,0 as the concatenation
st,o(-) =
f.)
(117)
Lemma E.6 (Optimality). Let {0, Jie(ma and fe,o be defined as in Definition
E.1
When Theo-
rem E. define the pseudo loss function
elli. (X.)
T St) = E - - log
(118)
Exes ei (xx).1((X1177
Then by choosing k = (E2), and assuming
Σ
Kuw MMM 2 8
(119)
Eg
KEM;
obtain the following loss guarantee:
U(fe, fi) S 0
lng
(120)
Lemma E.7 (Pre-activation size 1). Let 2
(r) M
+
((r)
Dzx
wi E Rd Define
6) (Erthrete 2. + Z(r)
xx
Then the following results hold:
(a) Naive Chebyshev bound: For any X > 0,
((w,, (01) + Hw,, M))= > A Alw:lliloog z0(2).
(121)
SON Eir)
The same tail bound applies to (w,, xxx (w,, - and (wi, 3
(b) High probability bound for sparse signal:
Pr (un ||wa max M2 log' d) S -S(llog d)
(122)
jE(d]
(c) High probability bound for dense signal: Let 7 = (ur; (lt)) Then
Pr
[m [/w.11/1og" log" d S E 2Oloo d)
(123)
Lemma E.8 (Pre-activation size II). Suppose the following conditions hold:
(w)) (t)
M2)² 2 (2))))
for at most O(1) indices j E [d],
(124)
(b )2
(0 M23 2 Ω
for at most Ole -RVlogg d) indices j E [d].
(125)
Vlogd
d(b(0)
11 500
(126)
log
Then, for any 0.0001,
Pr ((w
(1) ,, (())) 2 26!1) S 52llog1/4 d)
(127)
and
Pr ( (n!
(x) )
2 Ab Se Allog 1/4 d).
(128)
32

---

Published as a conference paper at ICLR 2026
Lemma E.9 aPre-activation size III., Ler i E [m]. Suppose there exists set Ng S [d] with [MG ((1) such that
(w/) (0) M2) so
(b
Vi Ni
(129)
polylog(d)
and
w (00 o
d(b("))2
(130)
polylog(d)
Then, for any À E [0.01,0.0 0.99).
Pr
Σ (uo)
2 + (wss ((
2 2b
S Se-allog d)
(131)
JEN;
Lemma E.10 (Gradient for sparse features). denote events
Suppose D.1
holds at iteration t 2 O. For j E [d].
A = S/ 2 b 1.3
C2
A2 = ( 2 6.0 - -0/0
(132)
A3 : { s + 23 2 2 0 C2 - D, b,t)}
A4 :
SV 11 2 a/
10 b(() }}
and quantities L1, L2, L3, L1 as
E 2 (IA, +1A))l
L
L2 : :P((A),
El(w
(133)
E1 (W +a)l
La
L4 = Pr(As.
(1)
Then we have the following results:
(0) (a) (all features) For all E [m]. if of 2 0, we have (when o < 2 the opposite inequality holds)
h4
Y,
xw,, 22 12b 2(rr)
Kw,
4
L
E
nj
(())
nn
extex 20:+l/w0
Ex #x ;=x
2
(134)
==1
+
(
0
B.
(()) +
(1))
|=(0)
O(In T2).
(b) (lucky features) If [
(1)
(10) we have
E
h4
(Y))
22)))
12bb
T. Lij
E
n.j
n.j
1w (0)
Ex
1wbb+ (1) 2x +x
,=,xx=
(135)
(VIEI
L
2(0)+ + (()
+
+
E
O + L4).
2
33

---

Published as a conference paper at ICLR 2026
(6)
<bb,,, then the opposite inequality holds with (0 b/ replaced by (o b, b(
If 0..
Lemma E.11 (Gradient from dense signals). Let i E [m] and j [d]. Suppose D.1 holds for the current iteration t. Then
I 1l2
E ha (n)
Σ Rull 2222000 (E), Mj)
Pr((s,((s)) + #0)
(136)
da
For dense features Mj j E [d1] [d]. we have a similar result:
| 122 dydi
E h4 Σ 1 R00
1 2 (E1) M1
Pr(hi,((r)) #1))))))) #0).
(137)
E.3 PROOF OF THEOREM E.
Proof of Theorem E.1 First we need to prove all the Theorem E. hold for t = T2. Indeed, (I), (4), (5), (6), (7) is valid at T2 from Lemma E.9 and Theorem D.I (2) and (3) holds at T2 obviously.
Now suppose it hold for some t 2 T2, we will prove that it still hold for t + 1. We first deal with the case where j E [d] and of Mj. where it holds that
(w: (w!++)), (M) (w _2 nA) + mle
Kw 12b. (=X M,)]
(138)
ME
Σ X
h (X, n.a)
(w (w.
(2x
(r) M2)
+
poly(d1)
In this case, to calculate the expectation, we need to use Lemma pute the probability of events A1 Ai by using Lemma E.7
E.10
Lemma
E,4
First we com-
Lemma
E.8
Lemma
E.9 and our
Theorem to obtain
Pr(A2)
(139)
poly (d) Olog
which implies
L1, La
(140)
poly (logo
Furthermore, from Fact E.
we also have
E
(r)
(
12 (mm Pall
S E
poly (d)
(141)
[m0,
Now we further take into considerations Lemma E.1I Lemma
E.4
We can obtain
I ll2
(e (++1) .MM < (wm M1 ())) + δ
" poly(d1)
(142)
d²
Indeed, since we have chosen learning rate rj = poly(a) and À E [artor, a is easy to prove (5) as follows:
For & Mj, ((w (0) M2 < o
This easy since by using Lemma E.10
vd=g
Lemma R.4 we can prove the following inequality by contradiction
1 d2
(w0 M3 (kw (t-1)
M,][[++5
nA +
w
l|2
L2
(143)
0
E,
Émax
ViE
34

---

Published as a conference paper at ICLR 2026
Now we begin to prove (6). For all i [m]] we have max k20! M S O
vd
iteration & T2. Now, by expanding the gradient updates of (w0 _M²), we can see that
(4)
Kw ( MM - (w MM 77)) + 1
|
+
E
(144)
S f Mj I(1 TTA) +
1 (4) 112-
TV did
where the last inequality are obtained as follows: From Lemma E.4 we have
1
(+)
12
(Kou (9 Mf31)
max ('(mm)
d²T
|100 (1)
12
0
(since max (1u 10 (e[m]]
M M
(145)
d²r
ddz
Vd
<ō
| (1)
ll2--
After (5) and (6) are proven, it is easy to observe (1) is true at Below we shall prove (2), (3) and (4), after which (7) can be also trivially proven.
Indeed, (2) is a corollary of (3) and (4), since if 8/8 S O(T log" d) and (4) holds, we simply have
1 2 w
Σ
w (t)
2 +
Σ
(w (4)
22
+
Σ
(w (9) .M, )
JEN.
JEN.JE
JEld
(
(*) BB
220 (2) 13
yw) (t) M2 + O(d) O
(1
Emax
+ O(dd) o
di
1
(146)
jEN;
(t) M2)² +
(t)
Ej )2
1
|120!
(40
113
=10
mmax
JEN.
which implies (2).
1 10 S Σ (w) 2M,2
(
20 12
E
Cmax
jENr
Σ (w! M2
(147)
JEN.
((1)( polylod(d)
d
< 01))
Thus we only need to prove (3) and (4). Indeed, for (3), letting i E Mj. we proceed as follows: we first write the updates of (w!² (t)
M2,) as
(w (++1) MM,) (w!) (9) M1) n) + v ! + +
3
(148)
(w (0) M1)(1 n7 + (1)
5
|/ 112-
+
d2
where the last inequality comes again from Lemma E.4 Q(( S ) log" d), by Lemma E.3 we have
Now suppose for some we have 8/6 z
Cmax
(w (wo(t11
((+1) M2) = =w M2)
1 1 poly(d) A
+
10
ll
d2
(149)
<(w(()) (1) MM))
1 1 poly(d)
2
35

---

Published as a conference paper at ICLR 2026
which means that (wo (o M,) < (w, (9) M2). This in fact gives 83 (/(+11 proven.
8/0/ so that (3) is which is the specifie iteration when 8 2 G2 log d, where G2 is defined in Lemma E.3 The induction of (4) follows from
dllogd d Now for (4), we need to induct as follows: for t S T == vllg log
similar proof in Theorem
D.1 After T we discuss as follows
d' 2'-48 [T,
When &
from above calculations, for each i' E MM. we have
kw (++))
M,
((0 (1) .M,1(11 - 7) + 1 o
v Iva
[
112
(150)
1 ((+1)
M2)1
K 1 M2 - nA) m // + O
| 112
On one hand, for those i' My such that (un) M,) | S b
o
( I1
(0 (12))
(0)1)))
s
we can
safely get (wn (t)1)) M2}
>> (mo ((+1)
M2)
On the other hand, if
(w0
(0)
6y - then
we have
T
1,j
(1)
V
(t)
0 de (w M₁)
00
0
Ey
(!0)
(151)
(w
(1)
M2}
(w M2}
P3
td
Thus by letting = (wo? ,M,)"
then
{w) (1-1) _Mj)
(m (t) M2) (1 + nT/ m1) + 0
3x
||00 (0)
112
Iva
(152)
(10 (1-1) MM²
(mn Mj) (1 + - m) + 0
Eg
|1200 (2)
||
w/d
Since at iteration t E Ty., d2.40
T6 - s (0(( (5)
it is easy to obtain that
Thus we have
(np! (o((+11), M2)
(wo (1+1) Mj)
(w!!) .MM
(1+ 17( (1) (t) A)(1 )
(w))) M2)
(1 + n 16 (1) - A)(1+ 32))
_ ) n (T/0 A)(1 + 52))
1 ++ n (0)0
A)(1
(1oS)) M2)
(153)
(im (0) M2)
1 (1- (5
)) =2
)
Vdl
(m M2,)
(ww M,)
2 (1-
wva)
(2 M,)
(T)) M,,)
II
E 1-0
))
(w
2 2(1..
vvva)
(T))
(w
M,)
where in the last inequality we have used our Theorem F.1 at TV
36

---

Published as a conference paper at ICLR 2026
41 d1.00 largely similar to the above. The only difference
The proof for iterations E here is that we rely slightly different comparison here: Indeed, we have
w ((+1) M2)
(w (0) M2}
(1 + } 77
(4) 1
1|
(154)
mm (1+1) M₁)
(w M2
(1 + 17 - -
11
/ 2/ 0 km
Here we can use similar techniques as above to require
Now we also have
(w (++1) M,)
(wo M2)
77 A)(1 - )
(w (w((+11) M2)
(20 M2)
7 3 A) ))
(wo MM
(155)
tvd
(1o M,)
(10 (1) (1..0//m) /n) M2)
(1-
E3 t' d'
-2(1).
(d1.40/m) (w
M),
d'=d2.49/ /v)
Now (4) are proven. (7) is immediate result of our update scheme.
E.4 PROOF OF THEOREM 3.1
The first part proves the convergence the loss function. The second part is a further extension of Theorem E.
Proof of Theorem 3.1 We start with the proof of convergence [[I in Theorem 3.1
Denote w(t) (w) (t)
w since our update is
((+1) (1()) - V Tuutnng (Ji) poly(di)>
(156)
have
77 V u Lung (ft), (1 - -0"
2 |1 V m ang (ft) + (1 00 -1 -0' poly(d.)
(157)
2n² (oly(d) + 1 1lw((() - 10 1 [|w((tt)) - -- + poly(d1)'
where the inequality comes from
V n Long (foll
Σ VV Lang
(158)
Each term O(1), and since m poly(d), the overall complexity is poly
Now we will use the tools from online learning to obtain loss guarantee: define pseudo objective for parameter
T Taaus, (0) T((
Σ
110,12
iclm] elfralXa))aaaa
(159)
- log
Σ [E(m]
100.12
Which is convex function over 0 since it is linear in (for fixed fú, we consider i(f 9. fi) to be convex with respect to 0, because fr,o (z) is linear, and softmax + log is a convex composition; the regularization term is convex).
37

---

Published as a conference paper at ICLR 2026
Moreover, have
Trang (w = Tang (fe);
(160)
and
Vo, Laug, (z=,²)= V Laug( (fi).
(161)
Thus we have
(V Lun (ft), w() - 8r) w(()) )((()))))())))))) 8*)
=n Vo VoLuug, (wn(t))
①
Laug, (wi()) - L'aus, (8*)
)
- -
Laug, (w(())) -B -T log
Σ Exs el (X.))))))
Σ aE|m]
10; 11
(162)
Iraug, -
0(1100113)
loga
im(m)
Laug (Ji) (Ji)- -0( log
① is because the surrogate objective function Lang, is convex function with respect to 0, so we can use first-order convex lower bound: f(0) f(0') I (V)()), 0 8*) is because Eiem)
X|10; 1 Eje{a EieM; 1110 Σje(d) ΣiEM; A TMiTT Σje(d) M M
Are"
Now choosing i ((2)) S 1. (so that Eiemm 1100 a) and by telescoping summation, we have
7++--1
Ta+T-1 -1
(Tang (ft) -0 logg
S T
17 ( Jang (ft) 10 - 6r)
T
((IIu(Ts) 0*11²)
Tn
(11 + 110 2 2TT (((s)))T0)))))
Tn
(163)
(1) + |10^
Tn
mm| (Ta) 12
T'
O
m3 T'
Since Tn 2 m= this proves the claim.
For 12 in Theorem
3.1 we have
38

---

Published as a conference paper at ICLR 2026
w.
Σ
(w))
M,
Σ
M2)
+
Σ M
JEN
JENi
jE|d/Nd
E
(w)
M,MM,
Σ
o
Gj
M,
Σ
|lof 1|2 vdz Eq
M
Emax
vdzz
JENu,jEld
JENd
sE(d]]M
Σ JEN
M2
+
E
d
M3
+
E
M
JEN.JEE
jEldi\id
(164)
From Lemma c., we know that for each j E [d], there is at least one neuron that can fully learn the feature M,, and at most Eg neurons can learn the feature Mj. Combining this with Theorem E.Ic)::
Σ (w M2) 1 iEM,
((
log d
(165)
Emax
we can conclude that the range of (w) M2) is Emax B2 Fmax
7]. and hence the range of aij is
L Ee¹ Fanns
E/ r]. Furthermore, from Theorem E.1 (e) and (f), can obtain that df, S ol cmax va
and Bi, S ol ) respectively.
Next, we compute the upper bound of [NG]. As a first step, we calculate the expectation of [W]]
m
B||MG|]
Σ
WG
M2 d. Odd)))
m
m
m
1
O(d
0 +
(166)
dcm
O (1-( <min
)(((->))
Fix neuron i, we have: μ4 = E[[ By Bernstein's 's inequality,
Pr[[ 2 ++ < exp
t 2 0.
(167)
2(u; + 1/3)
We set = 3(VpLL L.) and plug this into the inequality above. Then we obtain
Pr Wil 2 14 + 3-+3(vmLL 1)] I ct
(168)
Hence, for any constant >> O, taking L = clog yields
Pr NG S μ + 3(vp3 /cc d clog 2 1 do
(169)
Next, we apply the union bound. For the event
A = { Wi S us +((
(170)
the union bound gives
X7
m
Pr
A
Σ Pr(Af) 2 me -4
(171)
Taking L = clog(md)) we obtain
Pr Vi E [m], IMI S i + + (V/uclog(md) + clog(md)) 2 1 - (mm))
(172)
39

---

Published as a conference paper at ICLR 2026
We know μ4 >> log(md), have
IMG μi 1 = VV
log(md) N/
Mi
+ o(1)
(173)
o d=
Eimi
(2--))
with probability least 1 (md)
Finally, for each dictionary atom M, there are at least R(d=))) neurons i [m]] such that M; = {j).
From Lemma B.2 (c), we recall that MMj] 2 ((d")). Morcover, a neuron belongs to Mj, then cannot belong to Myr.
For in Theorem our proof is complete.
F THEOREM F.1
From Lemma B.2 c). we know that for each j E [d]. there is at least one neuron that can fully learn the minority feature M,.. When we prune out the lucky neurons that learn these minority features during the forward pass, the network will force the lucky neurons to further strengthen their feature learning ability the minority features during the backward pass.
After magnitude pruning, neurons encoding a specific minority feature are removed. Pruning these lucky neurons reduces si (X,, n,; Y,)) during the forward pass. The decrease in similarity reduces the positive logit l which in turn increases the gradient of the loss function, thereby encour- p.o maak
aging these lucky neurons to further enhance their learning ability on the minority features.
Fix one specific minority feature Mj and let M2.. [m] denote the subset of neurons primarily aligned with it, with [M n. For pruning rate [1/m, n/m], the number of pruned neurons am n. Let P C Mj. be the pruned set with [P] cm.
F.1 THEOREM F.1
Theorem F.1 (Feature Dynamics After Pruning). Starting from the pruning stage T4 with pruming ratio a, the following statements hold.
(a) When i E Mí.. we have
(w M,,) 2 (1- -y- TJe
C. logloog
(o (panllaaaa)
(2 aame, log log logd)) )) (w My.).
(174)
(b) When i & M2. and j + jr, we have
(itt,,,M) s (1-1A+
Os log logd 0 (panslaata)
t my,
(175)
+ ej
log log
polslog(a)) (o(porlogta)) + of
amer, log log oi
mm
M,).
d
dEg
(c) For each neuron E P and E [Ta, Ts), contrastive learning learns the following decomposition:
w = O'iji* Mj. +
Σ Mi +
Σ
Br
(176)
jeN:
SE(diN
where
O8. E
TT
oie 5 o (1+ 2 Td
| ! |Bub 5 (
Vol
|10 !000"122
(177)
40

---

Published as a conference paper at ICLR 2026
F.2
USEFUL LEMMAS
Lemma E.1 (Expected values of neuron activations after Ta). From T4 onward, the following results
hold:
(a) For positive pair,
7²
log log d
R
ha x
2 52 om
m 6j-
(178)
d
iEP
(b) For negative pair,
B
Σ ((())) h. (x(o)))))
= 0.
(179)
EP
(c) For negative pair,
E (N,(( (V Wm IX))) M3.) = 0.
(180)
Lemma F22 (Hffect of Pruning on Positive Logit Weight). At the pruning stage, for the data follow- ing distribution Dr, the post-pruning positive logit
satisfies
p. m.ak
P. (.
+
+
am
Ej*
log log
d
(181)
Lemma F.3 (Positive gradient). Let hig())) denote the i-th neuron iteration t Ti (so that
(4) 0). Then the following hold: (a) For each j [d].
Elbi,e(Yn)) (Vw, (X.),,,,
1 7, (w,
M[ En
(182)
(b) For each j E [di] [d].
E B[
(183)
F.3 PROOF OF THEOREM
Overview of the proof: first, the data be divided into parts: the samples that contain M, and those that do not. The former follow distribution D1, while the latter follow distribution D2. Next, let examine &
ptt(t) The values of l' differ depending the distribution: for samples from
mack
p,0 mark
D1, we have E' () mack
= -1-000 +
Ω
E
Ej¹
log loglogd), whereas for samples from D2, E' p. mack
- O(²). Since the latter do not contain My., pruning does not affect them.
Proof Theorem E.I. For any neuron it E P have
(wo M,,))
|!m ll2
-ww(") ,,,,) - 17 ( (Vw, Laug (f.), poly(d)
=(1 - (w My+)
(184)
+n Ex.. (1 E (1)()) mack
(Xn, B)) hur e (Yr) (V w. his (Xn), M,
+ |10 |
1/ (Xne B) his o (Xn,s) (V hi (Xn), M;
poly (d)
E
At stage T4, pruning is applied. We regard & p _
and l' (0)
fixed, and by combining
maok
Lemma FTc) with the law of total probability, obtain
41

---

Published as a conference paper at ICLR 2026
(w) (1+1) .M.+)
=(1 (20 M)-))
(2)
+n Ex.. vn (1 E"
ho (X.). M,.) Ex.. v. ( (V 1.
p
mark
-(1 -n))(e) (t)
.M.+)
+ Ex Y (1 - E'
p,0 ) Ex., Y,,,i
(Yn) (V 2o his (Xn), (x)) P..D.
maak
tn Y (1
l'
E Ex..NaDD
(Yn) (V w h4 (Xn),, M,.) Y,
.X...Y
p
mark
(185)
Combining Lemma F33 a) with 181 in Lemma
F22
we obtain
(++1) .M.+)
=(1 - mw (1) M,.)
+1/ Bx.. Y Di
(1 f'
p.o mack ) Ex..
I (wn (1)
M2 T [ n.j Emm ]3]
P DanD
tn Ex D2 (1 l' EXx.X
mola manch )
(w)
M2
E 2 n. 2n.g") ]]
P PX..XanDi
-(1- nd)
(w:
(1) M).)
log loggog dl )) (mn M2) Cz-
log log d
+1
θ
+ Ω
Cj*
1
0. 1- (w (t) M₁,)
+n
=(1 n)) u (
M,-) ) +
(
+ 5
am
E
E"*
log logel dl ) (1o Mj.) Gg-
d
log log d
(186)
Hence, the post-pruning one-step update along Mj-
(wof (t+1)) M,,) 2 1 1- n1 nesr
C2 log log (e
polsloz(a) ) + 52 ame,; log log ol )) (w MM,))
d
dEg
(187)
Similarly to 186) for any neuron & P, we have:
(w) (c+1) M,)
=(1 (w) (t) MM))
+ 17 Ex, Yn DD [1 (1-6 8'
Ex.. Dn
I (w M2) E[ 1 knnد) ))
+ 17 Ex Yn D2
(1
Ex. N.
1 ww
M,) ) _ kn
PonDYaDD
=(1 7)) (wo .M,)
+ 17
-
1
+22
am
E E5-
log lgglog! dl ))
,MM EE
log log log log
d Ej. d
d
1
log log
+ 0
1. (M,)
d
(188)
42

---

Published as a conference paper at ICLR 2026
Hence, the post-pruning one-step update along M, is (w) (++1) .2) = (1 (1-7) 17 + ny
C2 log log d ( (0 polvioog(a)
d
(189)
+ Cgr d
log log dl
(0( (0(notioott)) + O
aomey log log of d=g
(w) M2).
The above constitutes the proof of Theorem F.1 regarding pruning.
F.4 PROOF OF THEOREM 3.2
Theorem 3.2 (a) and (b) can be derived as simplifications of Theorem F.1 (a) and (b). Theorem 3.2
(c) coincides with Theorem F.1 (c). By taking the elapsed time T
((Fmss /5p 1)d
and simplifying
yor. C. log log
(a) and (b), then substituting into the conclusion of Theorem 3.1, the proof follows.
F.5 PROOF OF LEMMA
Proof of Lemmal FT The alignment with the target minority feature Mj- is (w, M,-). and we have (w,, M,: second part of the proof of Theorem 3.1 For the positive pair (X,, In). the latent variables Zn.j* and =. are correlated through the augmentation process. For a negative sample Xn,ss its latent variable zn..." is independent of those of the positive pair (Zm.g" + #n... )r so we have:
2 ) at T4 (This is the conclusion of Theorem 3.1 which can be found in the
= IL #n,s"""
(190)
For the anchor Xn and its positive Y,, we have
L
(((x)) = Σ
(r)
=
1 [(W,, WT,
M
((r)
n
+
2(r)
St
(191)
r=1
r=1
r=1
L
L
h₂())) = Σ
Wh, 24,28 22
(a)
=
wii
M
(+)
+
2+(+)
(192)
8=1
==1
L
2n
x+(8)
E =
E(r)
-+ (+)
(193)
n
Sn ,
r=1
We can write the outputs as:
1
(x,),) = 7 (wi, M2, + En))
hz(Yn) = = (ws,, M + E.)
(194)
For a negative sample Xn,s: En.8 = E/e 2n,0:
[(g) En, = Eye Sn.s, the output is:
1
ha(Xn,,) = 7, Wwi M 2n,a Em.s)
(195)
We first establish a lower bound for (X,,,(X,,))t) (h. Expanding and using zero-mean and independence of latent variables and noises, we have
1
E[A (Y.))
(w M2,) (w; M
L
1
Wi M,) 2 E 2n +
12
j=1
(196)
1
2
(wi My.)
) F$ Em $
I
log log
2 Ω (
d
43

---

Published as a conference paper at ICLR 2026
Therefore
(X,) 2 96( 7²
log log d
=
ejr
(197)
d
Next, we compute the expectation of h, (X he (Xn,, , (X) hol (x).))
wi M + En)) wi: Man, + En+)
(198)
By the assumption, the latent variables of X,, are independent of those of the negative X,,.x and all noises are mean-zero and independent. Therefore,
(w; M 22) (w:, M 2n, O, ((ww,,s)) Sn.)
(199)
Therefore, we conclude that
=0 Let P be the pruned set with |P] = am. Summing the per-neuron bounds over E P, we obtain
(200)
h X (n))s((Y.) 2 aaam 72
E
=g Pp+))
(201)
REP
h x (1 n, (x.,)) = 0.
(202)
iEP
This completes the proof of Lemma Finally, we compute the expectation of hi,e, (Xn,)), (V (x,),,, M,.), and we have
F.T (a)(b).
hi x,s) (V w. hi (n,), My+)
(T, wis Mên,s + Em,)) T, (Man + Ens M2.) )) Wi Mžn, + En (Enj" (En, My-))
(203)
E Ke,, M 2n.))) 2
2
0.
This completes the proof of Lemma C.
E hie X n, (V (V,; (X.), M,))
(204)
F.6 PROOF OF LEMMA F2
Proof of Lemma F.2
We link the logit to the pruning ratio and plug into the gradient growth. Recall the softmax weights and partial derivatives
up/T
ua/T
6
C 1 Eo
(205)
e + E caalt
euy/t + E5 eua/t"
Up Sim (Xns Yn),
Simy (X Xn.o).
(206)
06,
0e, Ju.
(1 (2);
C
(207)
Oup 7
44

---

Published as a conference paper at ICLR 2026
Pruning the size changes the similarities by
Aup Σ (x,()))
Au Σ (x.)))))ss
(208)
SEP
iEP
Next, calculate the first order change of E. we know:
u (Up: 41 Us),
Au = (Aup. Aug, Aus).
(209)
Using multivariate Taylor expansion up second order with remainder:
(u + Au) (u) = (u) Au + 2 Au Hp(u)Au + o(llAul)?), Au Au - 0.
(210)
By first order Taylor expansion, have
de Sup
8 de
A
uup +
Au, + o (llul)))
Ous
(211)
s
1 6e Σ 8=1
(1 - (g) Dup
C Aug.
We note that at T4, by the convergence of the loss function, we obtain E 1-) and both fy and E take fixed values. Then, by taking expectations over DeS and using the relation Σ. la Eg. we obtain:
B[AGT] - O
(222 E
Σ EEP
h. h,,,,,,,
E Σ h4 (X,,))) iEP
(212)
Also, by Lemma
given that
y²
log log d
ha (x. (..
- E Σ
2 2(am 6
6
(213)
NEP
SEP
Hence,
E[Ag] = - am
log log d
Ej*
0.
(214)
d
Hence,
Bfe² pt)(t) mask
1= Eie p.000 + 1] pl
EiL -Eee,,0oo 1opco-10 52( am Σソ
log log d
Ej-
(215)
= =1 0 - 5
log log d
Ev-
=
Now, converting to the form of 1 - E::
E El
= B(C²
(216)
P.O ())(() mark
Substituting the previous expression gives
E El
(0 mankk
1
-
2
a.
=3 Ej-
log log d
d
(217)
45

---

Published as a conference paper at ICLR 2026
F.7
PROOF OF LEMMA H3 A)::
Proof of Lemma F.3 a):
1e (Y, (V w. (X,,),, M,)
=E
he
(.10,.1)))/ 2x M,}
-E
Wi, (+))
Σ (2x Mi)
E(ME!) + +EH +EP.M)))
W M (0))
E
L
w (2) Σ (M (aMP)..M)))
E
E
Kei,, M 2 ((, M2))]
L
(218)
1 L
E
(w;, M M Σ
My &y Mj)
j'E(d]
1
E
(wi, M M
Σ
(My, M2 2n.) 2m0g
T22
S'EM
E.
(wi, M M Em.il
T,2
E
(wi, M ) fn ann
L
Σ
(wi, M 2n)
L2
j'E(d]
In the final step, we have
Σ
(w.,
MB
g {mms
(ws, M2)B &
(219)
17 J'Eld
This completes the proof.
46

---

Published as a conference paper at ICLR 2026
F.8
PROOF OF LEMMA F.3 B):
Proof of Lemma
E33 b):
E
14 Y (V we 14 (,,M))) M]
R
h
=x M1))
w - 1/2
-E
(Y,))
Σ
20
1
1
=x2 -2.M))))
=E
wi z/
200
Σ
(w -2/00
-E
x My
E
Σ (wi, M 2 (00))
Σ
M ME + Eir) M²²)
1,2
W M (+))
Σ (M M,
(220)
I.
E
(uD M
(+((s)
(M Σ
2((r) Mf}
L
8=1
r=1
E
[w,,, M )
(M M,,))
2
E
(un; M
2n.3', 1
I
j'l(d)
1 E
(w, M
Σ f'E(d)
My' M2 2mm
7,2
1 E
(wi M
Σ
(My, M2 )
L
j'Eld]
[(DD) M o]
1,2
0.
G PROOF OF LEMMAS IN APPENDIX B
G.1 PROOF OF LEMMA (.2
Proof of Lemma B.2 a): At initialization, the neuron weight w0! is a high dimensional Gaussian
vector:
w 00) o(0,o ).
(221)
with (o!") E Rdi and each coordinate w (k) -N(0, i.i.d.
47

---

Published as a conference paper at ICLR 2026
di
(0)
|| = Σ
2
ww
(0)
(k) 2
(222)
We know that u (k) ~ ~N(0, oZ), so:
of w (0)
x
(d1).
(223)
According to the concentration inequality of the chi-square distribution:
If X ~ (d1), then for any 0 < < 1, we have:
X
d1E²
Pr
2
S
exx
(224)
di
Therefore, have:
(0) 2
w
d1e²
Pr
1
2
s 2
(225)
ofd1
4
Choose suitable E to derive the precision range and we choose: E =Ō At this time, the probability of deviation is:
(
(+))
Vdg
Pr
w
(0)
2 oo
S 50 di)
21 1
(226)
poly(d)
That is:
uo (0)
on (1-ō 1
(dn)) (1-0(1)))
ofd 1 11+ ( 1
(227)
Vdi
This holds with high probability (1 poly(aj). ).
G.2 PROOF OF LEMMA (B):
Proof of Lemma B.2 b): Let:
Z,
do w (0)
Ia).
Ia.)
(228)
Then we have:
MM
w. (0)
o
MMM T
2
(229)
We regard T T as a rank-d projection matrix, projecting Z, E Rd onto the column space of M
we use the following property:
MM a fixed rank-d projection matrix, and Z, ~ No00 It. ), then:
|MM 2 z (MMJ) MM T Z, Z MM Z, | Z,
(230)
M Z, ~ I2)
(231)
Therefore, we can conclude:
|M 2. 2112 x2d)
= E MIIM 2/122] = d.
(232)
48

---

Published as a conference paper at ICLR 2026
And it satisfies the following Chi-square concentration inequality:
IP ( z.1l2 -d S <d)d 2 11 - exx (-ceEd)
(233)
Choose ((1/vd), and the result holds with high probability. We substitute back w
(0)
(0) MM T wi | 1 M.MM 2.112 E ond (1 1 (1-Ō (0) (1a ) ond 1+
(1+0 Ö
//2
) (234)
Vdl
G.3 PROOF LEMMA B.2 C):
Proof of Lemma B.2 c): Recall if g is standard Gaussian, then for every > 0,
12 + e
1/2
Pr go,1)
l t]
1 1/
(235)
2
2
Therefore, for every i E [m]] and j E [d].
Ci log d d
P1 = Pr
w M2)²
MM T
w of
Pr
w
.M,
2 Vei logd
(236)
2 Ω
do1/2
1
Lmax (2.(1+%) d' "min
and
P2 = Pr
1 22
C2 lolog d
M MM'
w
(0)
d
w .2
Pr
2
C2 log
(237)
00
Vlog d
doa/2
Vlog d
Smin emax
(7.(1-))
We define the following events in definition B.1
A; Lucky neuron satisfies conditions 1(i.e,, the response is large enough and in the correct direction)
B;: for all j' # j. lucky neuron i satisfies condition 2 (i.e., small responses in other directions)
We now compute the probability of the intersection event A, 0 B::
Pr[A.]
P1 22
2 d
(mmax Crann
.(1+5))
Pr[B:] = (1 - -2p2)) d1 8 -(d-1)p = -dd 1)d- = -atta =1,
(238)
Pr[A,n = P1 (1- p2) 1 =
df (mmax) )² (1+7))
2
Vlogd
49

---

Published as a conference paper at ICLR 2026
(1) We now have dm neurons. Therefore, the expected number is:
-(147)
E [MMi m Pr(A, BB] dm 1 d
Smiin
(239)
Cm
Emax
((1+7))
Ω
Chernoff bound (Lower-tail form): For any (0,1), we have:
Pr
X; < (1 (21-6)1 5)10 S se 31
(240)
Let 1, 2'
we obtain:
Pr Σ X; < 2/14 _ e//8
Pr [M < (d Se- -2(d11)
(241)
Pr [[M]] > 0 (d 2 1--6 1
(2) We now have m = dm neurons. Therefore, the expected number is:
c .O M|l = m P2 = C
(Volg d
(umia (²-(1-))
VVog
(242)
dm -
(.min "min
(111--))
o
emaxx
Vlog
d
Chernoff bound (upper tail) tells that for any 0 < 8 < 1, we have:
Pr
Σ X; (1 + 8)μ
3 ())
Pr
vlogd VVoga )
1
0 Vigg dua
(243)
1
MyI < Pr
Vloga )
2 1
d4
1
Pr[|M ~0MMMM~000 <0
2 1
14
G.4 PROOF OF LEMMA B.2 D):
Proof of Lemma B.2 d): We know: IM, < O(d²). There are d indices E [d]. Therefore, the total number of pairs (i.j) such that E My is at most:
dd =((tt)) tw).
(244)
On the other hand, the total number of neurons is m cm, So for any fixed we define:
Nr := (j E [d] : My).
(245)
Then,
d
Σ
N =
M,l < tus)
(246)
1=1
/=1
Therefore,
m
N₂ < ( tun = O -O(d'((imi)) (d1=(((mm)²:010)))
)(
E
N. = m
(247)
50

---

Published as a conference paper at ICLR 2026
Then:
Pr (10, (0) M2} 2 al log d)) s 2exp
200 (-207)
82
2 2-n(Voog
(248)
Fix E [m]] and consider d different j. Each has probability 2 2-22lyloRd) to exceed the threshold. Therefore, the expectation
[ E [d]
(w0 M,) 2 2 (0o 1/4 d)) d)} = (2 .d..
vlogo.....
(249)
H PROOF OF LEMMAS IN APPENDIX C
This section can be found in the Supplementary Material.
I PROOF OF LEMMAS IN APPENDIX
This section can be found in the Supplementary Material.
J PROOF OF LEMMAS IN APPENDIX
This section can be found in the Supplementary Material.
K PROOF OF ADDITIONAL LEMMAS
This section can be found in the Supplementary Material.
51


Reviews:
[{"review_id": "HNJQHy5Kbs", "reviewer": "Reviewer_8inz", "summary": "The paper deals with contrastive learning, a self-supervised framework in which input data are paired into positive or negative pairs based on their similarity in semantic meaning. In particular, data are generated according to the Sparse Coding Model: they are linear combinations of feature vectors plus noise, feature $j$ has an activation probability controlled by a parameter $\\epsilon_j$, and features are in general imbalanced (majority features have $\\epsilon_j = \\epsilon_{\\rm max}$, minority ones $\\epsilon_j = \\epsilon_{\\rm min}$). Positive pairs are formed with inputs sharing the same token-aggregate features. The model is a Transformer-MLP built with a single head attention layer followed by an MLP. Training is performed by minimizing the InfoNCE loss, suitable for contrastive learning. A pruning mask filtering temporarily small magnitude neurons is applied when computing the gradients at each training step, and released before updating the neurons. \n\nThe paper provides formal results on the dynamics of training of this model. In particular, Lemmas 3.1, 3.2 and Theorem 3.1 show the existence of 3 temporal regimes during training with no pruning: in the first, neuron weights grow in feature directions but are suppressed in non-feature directions, with the growth rate in a feature direction depending on its frequency; in the second, few (lucky) neurons align significantly with single features, while ordinary neurons align in composite directions; in the last (at convergence), the training error is small and neurons are strongly aligned with a subset of features, weakly aligned with the remaining features, and small in the non-feature directions (still, only a limited number of neurons specialize in learning a single feature). Theorem 3.2 shows that pruning amplifies the learning of minority features.", "questions": "- Can the authors make the main text more self-contained? For example:\n    - I could find the hypothesis on the batch size $K$ (such that empirical gradients approximate population ones) only in Lemma B.1;\n    - I am assuming $C_m$, $C_z$ to be positive constants, is that the case?\n    - The scaling of $T_1$, $T_2$ is not given in the main text;\n    - $\\Xi_2$ is not defined in the main text, etc.\n- In appendix A.2, the authors provide numerical experiments with synthetic data. Can they report in the main text numerical evidence for the 3 training regimes they are able to identify theoretically, for better illustration and check?\n- How are negative pairs chosen for the case of real data of section 4?\n\nMinor:\n\nSometimes * is used instead of $\\star$ to denote lucky neurons.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Both the architecture and the training protocol are relevant. Theoretical predictions for the training dynamics of Transformer-based encoders are of utmost interest. Narratives on the assumptions behind the data model and on the formal results are provided. Numerical results on real data support the theoretical claims on the advantage of pruning.", "weaknesses": "The formal results are hard to read, as the main text is not really self-contained (see Questions below), despite the commendable effort of Table 1. Numerical illustrations in the vanilla setting, even with synthetic data, could help explaining the practical relevance of the bounds provided (for example, by tracking the inner products of lucky/non-lucky neurons with features during training in the 3 regimes, and comparing with theoretical bounds). Considering that reviewing proofs in Appendix is out of question due to time limitations, numerical checks should help strengthening the claims provided by the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PVWBOpF5JF", "reviewer": "Reviewer_xbUv", "summary": "The paper analyzes why contrastive learning deteriorates under imbalanced data. Using a sparse latent-feature model and a Transformer-MLP encoder with InfoNCE, it shows that rare (minority) features grow more slowly, fewer neurons specialize in them, and neurons are forced to represent multiple features, effectively increasing the capacity needed to cover all features. To mitigate this, the authors propose a magnitude-based, forward-masked but backward-unmasked pruning scheme: small-magnitude parameters are masked out only in the forward pass, but all parameters are still updated. This selectively amplifies gradients along minority-feature directions and restores more balanced representations. Experiments on CIFAR-LT and ImageNet-LT confirm better linear-probe accuracy and smaller head–tail gaps.", "questions": "- The paper sometimes reasons at the neuron level but prunes at the parameter level; please clarify the exact relationship between “parameter-level masking” and the claimed neuron-level effects.\n- Which parts of the analysis rely essentially on (i) identity/self attention, (ii) sparse independent features, or (iii) single-layer MLP, and which parts you believe could be relaxed or extended? A short discussion of “what breaks / what survives” under more realistic assumptions would help.\n- Can you provide empirical evidence of “lucky” vs. mixed-feature neurons (e.g., alignment plots)?\n- In the experiments, how does the method behave when the imbalance ratio = 1 (i.e., no explicit long tail)? Since the proposed pruning is fairly general, do you observe gains in “implicitly imbalanced” settings (e.g., feature-frequency skew induced by augmentations or views) even without a constructed LT distribution?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Provides a rare, neuron-level theoretical analysis of contrastive learning under data imbalance, clearly explaining how minority features are under-learned.\n- Connects the analysis to a simple, practical fix (magnitude-based forward-masked, backward-unmasked pruning), making the work actionable.\n- Writing and structure are generally clear, making a dense theoretical contribution reasonably accessible.", "weaknesses": "- Experiments mainly compare “with vs. without pruning” and lack baselines from other long-tailed methods.\n- Sensitivity to pruning ratio/schedule is not deeply analyzed.\n- Paper could more explicitly discuss limitations and when the proposed analysis may not apply.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m2nErhKTrj", "reviewer": "Reviewer_yCvM", "summary": "This paper provides a robust theoretical analysis of the training dynamics of Transformer-MLP models in learning feature representations through contrastive learning in imbalanced data scenarios. Specifically, the analysis of neuron weight evolution reveals how a minority of features undermines overall model performance. Building upon this, they revisit amplitude-based pruning methods, theoretically demonstrating that pruning yields more robust and balanced representations.", "questions": "As described in “Weaknesses”.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "This paper explores how imbalanced data degrades representation quality in contrastive learning from a novel perspective of neural weight evolution. Through extensive theoretical analysis, the authors demonstrate that a minority of features weakens representational power while increasing the demand for complex architectures. Building upon this, they further prove that pruning techniques enhance gradient updates along these dominant features, thereby mitigating performance degradation caused by imbalance. This forward-looking approach holds significant promise for advancing the field.", "weaknesses": "1：This paper uses numerous notations, some of which lack clear definitions upon their first appearance. Additionally, maintaining consistent notation throughout the text would improve readability.\n\n2：Could the authors please clarify the meaning of \"feature frequency\"? Specifically, how are the majority and minority features identified within the unsupervised learning framework?\n\n3: The study is confined to the Transformer-MLP model. Could the authors discuss the generalizability of their approach to other architectures or tasks, such as long-tailed visual recognition?\n\n4: While the authors provide substantial theoretical proofs, the effectiveness of the pruning strategy for imbalanced scenarios remains unverified by strong experimental evidence. Additionally, could its potential as a plug-and-play module for existing methods be discussed?\n\n5:The authors are advised to provide precise citations to the appendix for their key conclusions, and to thoroughly review the manuscript to correct various notational errors and inconsistencies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9hqhUrVX72", "reviewer": "Reviewer_Zaui", "summary": "This paper develops a theoretical framework to analyze how contrastive learning behaves under imbalanced data distributions, with a specific focus on Transformer-based encoders. It identifies how feature frequency imbalance (i.e. majority vs. minority features) affects neuron-level learning dynamics and the quality of representation learning.\n\nThe paper shows that the training progresses through three distinct stages—initial feature growth, specialization of “lucky” neurons, and final convergence. It also shows minority features are learned more weakly, leading to fewer neurons specializing in them and resulting in degraded representations. It further shows introducing magnitude-based pruning during training mitigates these effects by amplifying updates for neurons aligned with minority features, thus improving representation balance.\n\nTheoretical results are corroborated by empirical experiments on various dataset (CIFAR10-LT, CIFAR100-LT, and ImageNet-LT), demonstrating consistent performance gains and reduced accuracy gaps between majority class and minority classes.\nIn addition, the paper provides theoretical proofs for convergence, establishes feature alignment properties, and rigorously derives the effects of pruning on the learning dynamics.", "questions": "(1) In Transformers, neurons often encode superposed features that are only linearly separable after training. The analysis assumes near-orthogonal features and pure specialization at convergence. Can the authors comment on whether their theory predicts or forbids feature superposition (neurons simultaneously encoding multiple correlated features)?\n\n(2) Could the authors maybe clarify how the temperature parameter ττ influences imbalance sensitivity in their theory?\nSpecifically, does a smaller ττ (sharper similarity weighting) amplify majority-feature dominance by concentrating gradients, while a larger ττ mitigates it by smoothing updates?\n\n(3) If the attention mechanism were trainable, pruning neurons downstream of attention would alter gradient flow through $W_K, W_Q, W_V$​. Could the authors speculate on whether such coupling might reinforce or dampen the minority-feature amplification effect?\n\n(4) Could the authors comment on how neuron specialization (Theorem 3.2) is connected to the generalization performance?\nIn particular, does specialization provably enhance linear separability or any other preferred properties of the learned representations?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Originality** The paper seems quite novel. Given my knowledge I am not aware of many existing works exploring the statistical generalization theory of contrastive learning and data imbalance. While prior works have addressed imbalance heuristically or empirically, this paper seems the first formal theoretical treatment of the phenomenon. \n\n**Quality**\nThe math of the paper seems rigorous, supported by clearly stated lemmas and theorems. Empirical experimental results seem to align with statistical theory. The synthetic data experiments (Appendix A.2) further validate theoretical predictions in controlled settings.\n\n**Clarity**\nThe clarity of the paper is good. Although this paper is theoretically heavy and thus not easy to read, its clear presentation of results has made the task easier. The three training stages are clearly delineated, and key insights are summarized upfront (Section 3.1). The figures and tables are easy to interpret, and a table of notation is provided. The proof sketch is helpful too.\n\n**Significance** \nThis work has strong implications for self-supervised learning on real-world, imbalanced datasets. It offers theoretical understanding in the study of contrastive learning, showing how imbalance alters neuron specialization and model complexity requirements.", "weaknesses": "There is no major technical flaw detected. However, I do have the following concern.\n\nAssumptions might be a bit strong.\n\n(1) The entire analysis seems to focus on the sparse coding model with orthogonal features and independent Gaussian noise. I understand the technical challenge in analyzing a more general setting. But the current setting feels a bit simple, because in reality features are usually correlated. This orthogonality assumption prevents the theory from describing how contrastive learning handles overlapping or dependent semantic features. \n\n(2) The paper’s assumption that self-attention remains identity-fixed fundamentally limits the scope of its theory. In other words, $W_K, W_Q, W_V$​ are not learnable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)