Binary Hypothesis Testing for Softmax Models and Leverage Score Models

We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. Below, we want to respond to the weaknesses and questions.

Concern 1: The primary issue lies in the disconnect between the stated motivation—understanding large language models (LLMs) through the softmax attention mechanism—and the actual focus of the work, which is on sample complexity for binary hypothesis testing of softmax distributions.

Answer: We thank you very much for your insightful comments. In Sections 3 and 4, we study the distinguishability of models (softmax and leverage score based) through the lens of binary hypothesis testing, establishing tight sample complexity bounds. These results directly address the challenge of determining how much information (or how many queries) is needed to tell apart closely related models, a theoretical formulation aligned with understanding model "abilities" via limited parameter access. Moreover, our framework sets a path toward identifying distinguishable components of large models. For instance, showing that certain parameters contribute more significantly to distinguishability (via Hellinger distance bounds) offers insight into what might constitute an “ability region” within a model, aligning with the introductory motivation.

Thank you very much for your help checking our typos, and we will fix them in the revised version of our paper.

We thank you again for your insightful comments.

We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. Below, we want to respond to the weaknesses and questions.

Concern 1: However, in the latter sections of the paper, the contents are not cycling back to the theme.

Answer: We thank you very much for your insightful comments. In Sections 3 and 4, we study the distinguishability of models (softmax and leverage score based) through the lens of binary hypothesis testing, establishing tight sample complexity bounds. These results directly address the challenge of determining how much information (or how many queries) is needed to tell apart closely related models, a theoretical formulation aligned with understanding model "abilities" via limited parameter access. Moreover, our framework sets a path toward identifying distinguishable components of large models. For instance, showing that certain parameters contribute more significantly to distinguishability (via Hellinger distance bounds) offers insight into what might constitute an “ability region” within a model — aligning with the introductory motivation.

Concern 2: In addition, the claim in the first page "As we delve deeper .... self-attention" needs reference to support.

Answer: We completely agree with your comment. In the revised version of our paper, we will include the following citations:

$\bullet$ [1] introduced the Transformer model and the use of softmax in computing attention weights.

$\bullet$ [2] discussed variants and performance implications of softmax in attention computation.

$\bullet$ [3] and [4] directly analyze optimization on the softmax regression problem and show the impact of the softmax unit on self-attention and in-context learning.

Concern 3: Authors study a softmax models and leverage score models which can be formulated as a matrix. However, it is not clear how the conclusions made on this single-layer model can be generalized to the LLM and transformer as a whole system.

Answer: We thank you for your insightful comments and kindly refer you to our Response to Weakness 1 to Reviewer 8u1R.

Concern 4: The discussion on leverage scores appears forced and disconnected from the original motivation, aside from some vague remarks about their usefulness. While highlighting similarities may be valuable, the authors need to clarify how this perspective relates to the central question at hand.

Answer: ​​Our primary motivation is to understand how model components, especially those arising in LLMs, can be distinguished via limited parameter sampling. While the softmax mechanism arises directly in self-attention, the leverage score model serves as a broader, more general abstraction of distributional behavior driven by matrix-parameterized functions, similar in form to softmax. Both softmax and leverage score models define distributions over outputs conditioned on structured inputs, and both are parameterized by matrices, making them amenable to a unified theoretical treatment through hypothesis testing. Rather than being a disconnected addition, the leverage score model allows us to extend our analysis framework and highlight that the difficulty of distinguishing close models under sample constraints is not unique to softmax, but also arises in other distributional settings relevant to algorithms and data analysis.

We thank you again for your insightful comments.

[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. "Attention is all you need." NeurIPS’17.

[2] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins et al. "Rethinking attention with performers." ICLR’21.

[3] Zhihang Li, Zhizhou Sha, Zhao Song, Mingda Wan. "Attention scheme inspired softmax regression." ICLR’25 workshop.

[4] Yeqi Gao, Zhao Song, and Junze Yin. "An iterative algorithm for rescaled hyperbolic functions regression." AISTATS’25.

We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. Below, we want to respond to the weaknesses and questions.

Response to Weakness 1: We thank the reviewer for raising this important point. Our work intentionally focuses on single-layer softmax and leverage score models as a foundational step in the theory of binary hypothesis testing over these widely-used classes of distributions. The novelty of our contribution lies in formalizing and analyzing the binary hypothesis testing problem in these structured settings, which, to the best of our knowledge, has not been previously addressed in the literature.

Regarding motivation, we highlight that both softmax and leverage score distributions arise naturally in machine learning and numerical linear algebra, but their sample complexity under hypothesis testing had remained unexplored. We clarify this explicitly in the introduction (Section 1) and further elaborate on the motivation from LLMs and self-attention mechanisms, where softmax distributions are central.

As for the extension to multiple layers of softmax, we agree that this is an important direction and appreciate the opportunity to clarify. In deep neural networks, softmax layers often appear at the output or in attention blocks. While our current theoretical framework addresses a single softmax layer parameterized by matrix $A$, multi-layer architectures can be viewed as compositions of functions, where only the final layer produces a distribution over outputs (which can be seen in Figure 1 of the famous “attention is all you need” paper https://arxiv.org/pdf/1706.03762). Therefore, the final softmax layer can still be abstracted and analyzed using our current hypothesis testing framework.

In the revised version of our paper, we plan to discuss this direction more explicitly and include formal remarks about the potential for generalization to multi-layer or multi-head softmax settings.

Response to Weakness 2: Thank you very much for your insightful comments. Leverage scores are a central concept in numerical linear algebra, machine learning, and graph algorithms. As detailed in Section 1 and Section 4 of our paper, they appear in numerous algorithmic applications including matrix approximation (e.g., CUR decomposition), randomized linear algebra, graph sparsification, maximum flow and matching, and random spanning tree generation.

In all these applications, a leverage score distribution defines a probability distribution over data points, rows, or graph elements, used for randomized sampling or importance weighting. Given their probabilistic nature and sensitivity to the underlying data matrix, it is natural to ask whether two leverage score models correspond to the same or different underlying structures, particularly when the models are accessed as black boxes, which motivates our formulation of the binary hypothesis testing problem for leverage score distributions.

Response to Weakness 3: The core novelty of our proofs lies in adapting classical binary hypothesis testing, typically studied in the context of generic distributions, to the structured, parameterized families of distributions induced by softmax and leverage score models. Unlike arbitrary distributions, these models produce distributions that are nonlinear functions of the input and matrix parameters, which poses unique analytical challenges.

For the softmax model, one key difficulty is that different parameter matrices $A$ and $B$ can induce indistinguishable distributions due to invariance under certain transformations (e.g., row shifts). To handle this, we introduce structural constraints (e.g.,$\\|A - B\\|\_{2 \to \infty}$) and prove that the Hellinger distance between softmax outputs under constrained inputs governs the sample complexity. This leads to a tight upper and lower bound framework via careful analysis of the sensitivity of softmax distributions to perturbations in the parameter matrix.

For the leverage score model, the challenge is even greater due to the nonlinear matrix expressions involved, including matrix inversion and normalization, and the fact that the input sss is a vector that rescales rows of the matrix. We overcome this by establishing operator norm bounds and using perturbation theory to relate changes in the parameter matrix to changes in the output distribution. Our proof carefully propagates these changes through the matrix expressions and yields tight dependence on both the model difference $\\|A - B\\|$ and input constraints.

Response to the question: Thank you very much for carefully checking this. We will fix this in the revised version of our paper.

We thank you again for your insightful comments.