Trade-Offs of Diagonal Fisher Information Matrix Estimators

Thank you for recognizing our contributions and their soundness.

Re: Significance vs [37]

Please see our global response to all reviewers.

Re: Better clarity on which estimator is preferred

Through our analysis, We have identified certain scenarios where one estimator is superior to the other.

For example, ReLU networks can only apply $\\hat{\\mathcal{I}}_1$. As another example, $\\hat{\\mathcal{I}}_2$ has zero variance in the last layer and is always preferable than $\\hat{\\mathcal{I}}_1$. Similarly, in the second last layer, $\\hat{\\mathcal{I}}_2$ has a simple closed form and preferable for neurons in their linear regions (see Remark 4.5).

In general, Theorem 4.1 suggests that the variance of both estimators depend on three factors: (1) the number of samples; (2) the derivatives of the neural network output wrt its parameters; (3) moments of the sufficient statistics $\\mathbf{t}(\\mathbf{y})$ (see Table 1). One has to incorporate these factors based on their specific neural network and settings, to decide which estimator to use. We agree with the reviewer to recall these points in the conclusion section.

Re: Numerical examples are fairly toy

Our contributions are mainly theoretical and do not depend on the empirical results. In the paper and appendix, we use MNIST MLP with different activations (sigmoid and log-sigmoid) and different number of layers (4 and 5) to showcase the bounds of the variance and how the variance evolves over training time in different layers of the neural network. These numerical examples mainly serves the purpose of providing intuition. We do agree with the reviewer that more extensive empirical studies are meaningful as future work.

Re: Numeric example for regression similar to Fig (2)

We thank the reviewer for this suggestion. As remarked by Proposition 5.1, for regression, all bounds in Theorem 4.1 becomes equalities. Therefore it is less meaningful to examine the case of regression, where all interested quantities $\\mathcal{I}$, $\\mathcal{V}_1$, and $\\mathcal{V}_2$ reduce to evaluation of the Jacobian/Hessian of the parameter-output mapping. The limited time for the rebuttal does not allow for redesigning the pipeline of FIM estimation on a new dataset and learning task. We therefore hesitate to overstate our empirical discoveries or to promise experimental extensions without clearly understanding their implications.

We would also like to iterate that aim of this paper's contribution is primarily theoretical. All though additional numerical experiments would be nice, a thorough independent empirical study is out of this paper's scope and is arguably would make an independent piece of work (examining various settings, optimizers, and architectures etc).

Thank you for recognizing our contributions and praising our writing. We respond to your remarks and suggestions as below.

Re: The notation $\\hat{\\mathcal{I}}_{1}(\\theta_i)$

We propose to explicitly define it as $\\hat{\\mathcal{I}}_{1}(\\theta_i) := (\\hat{\\mathcal{I}}_1(\\theta))\_{i i}$ at the first appearance of this notation in the beginning of Section 3, and add a sentence to remind the reader that $\\hat{\\mathcal{I}}_1(\\theta_i)$ is an abuse of notation and actually depends on the whole $\\mathbf{\\theta}$ vector.

Re: The notation $\\hat{\\mathcal{I}}_j(\\theta_i \\mid \\mathbf{x})$ with $\\mathbf{x}=\\mathbf{x}_1=\\cdots=\\mathbf{x}_N$

We will rewrite the sentence to clarify that $y_1,\\cdots,y_N$ are i.i.d. from $p(y \\mid \\mathbf{x}; \\mathbf{\\theta})$ as the reviewer has suggested.

Re: Explicit proof of Lemma B.1 to make the paper self-contained

We agree to include such a proof that is consistent with our current notation.

Re: Line 168, Eqs. (4) and (6), the scale of $\\mathcal{I}(\\theta_i \\mid \\mathbf{x})$ and $\\mathcal{V}_2(\\theta_i \\mid \\mathbf{x})$

Yes, this is a typo and the corrections suggested by the reviewers are correct. Eqs. (4) and (5) are the bounds which depend on the network Jacobian.

Re: Line 246, sample complexity should not be $\\mathcal{O}(\\frac{1}{N_xN_y})$

We thank the reviewer for pointing out this mistake. By the total variance decomposition in Eq. (13), the variance wrt $q(\\mathbf{x})$ scales inversely to $N_x$, the number of $\\mathbf{x}$-samples, and the variance wrt $p(y \\mid \\mathbf{x})$ scales inversely to $N_y$. Therefore the total variance is in the order of $\\mathcal{O}\\left(\\frac{1}{N_x}+\\frac{1}{N_y}\\right)$ as the reviewer suggested. We will make sure to correct this sentence and related places in this paragraph.

Re: Relationship with reference [37]

Please see our global response. We shifted the closed-form expression of the variances to the appendix mainly due to space limitation. Should more space be permitted, we agree to improve our self-contained narrative by moving them back into the main paper.

Thank you for recognizing our strength and potential usefulness in application areas including continual learning.

Re: Significance vs [37]

We kindly refer the reviewer to our global response.

Re: Assumption of $\\theta$ being the true value

We clarify that our result holds for any $\\theta$ in the parameter space of neural networks, aka the neuromanifold. Our formal statements and their proofs do not depend on setting $\\theta$ to the true value or its MLE, nor to be "nearby" to the MLE in any sense. Similar to how the FIM $\\mathcal{I}(\\theta)$ is a PSD tensor when $\\theta$ varies on the neuromanifold, our central subjects $\\hat{\\mathcal{I}}_1(\\theta)$, $\\hat{\\mathcal{I}}_2(\\theta)$, $\\mathcal{V}_1(\\theta)$, and $\\mathcal{V}_2(\\theta)$ are all tensors defined for any $\\theta$.

Re: What are the effects of model miss-specification

Model miss-specification can be examined from the perspective of empirical data deviating from the joint model distribution $p(\\mathbf{x}, \\mathbf{y}; \\mathbf{\\theta}) = q(\\mathbf{x})p(\\mathbf{y} \\mid \\mathbf{x}; \\mathbf{\\theta})$ in Line 24. Our analysis does provide convenient tools to study this phenomenon.

When the observed samples $\\mathbf{x}$ deviates from the "true" $q(\\mathbf{x})$, the error in estimating the FIM is described by the first term in Eq. (13): a high variance of the conditional FIM $\\mathcal{I}(\\theta_i \\mid \\mathbf{x})$ leads to a large error in estimating the FIM.

On the other hand, when the observed $(\\mathbf{x},\\hat{\\mathbf{y}})$'s sampling deviates from the predictive model $p(\\mathbf{y} \\mid \\mathbf{x}; \\mathbf{\\theta})$, data also plays a larger role. In this case, the empirical Fisher becomes a biased estimation of the FIM, which is described by Lemma 6.1 and its surrounding text. Our FIM estimators do not depends on the observed $\\hat{\\mathbf{y}}$ in the dataset and is not affected by such a bias. Note the difference in sampling of the labels $\\mathbf{y}$ / $\\hat{\\mathbf{y}}$ in the data FIM versus the FIM.

When the model is well specified, the data FIM and FIM are equivalent (see Line 328).

We thank the reviewer for praising the technical developments of the paper.

Re: Usefulness of FIM estimation, e.g. in NN optimization

As stated in the submitted paper, the FIM can be used (Line 117) "to study the singular structure of the neuromanifold [2,40], the curvature of the loss [8], to quantify model sensitivity [28], and to evaluate the quality of the local optimum [15,16]." Other applications include continual learning as pointed out by Reviewer wDHQ.

We believe studying the estimation quality of the FIM in the general setting -- without examining the specific application and implications of these use cases -- is meaningful on its own. First, the bias and variance always exists for tackling the computational cost of estimating the (diagonal) FIM, but has not been thoroughly analyzed. Second, such a study can be potentially applied to different scenarios. For example, in NN optimization, the bounds of the diagonal FIM leads to bounds of a learning step, which is determined by the product of the inverse of the diagonal FIM and the gradient vector. In our toy example in Figure 1, we demonstrated how optimization can be interfered with the variance of the FIM. To apply our results into specific optimizers would be another independent work.

Re: FIM for analyzing NN structure

In the paper, we mentioned that one application of the FIM is to study the "intrinsic structure of the neuromanifold". Here, "structure" does not refers to the NN structure but the Riemannian geometry structure in the space of neural networks (aka the neuromanifold), where the FIM plays the role of a local metric tensor. Our main results are universal in the sense that they do not rely on specific NN structure. On the other hand, the FIM is indeed affected by the NN structure, which in the simplest case includes the depth and width of neural networks [15,16]. This is beyond the scope of the current paper.

Re: Organization of technical contents

A main reason for the technical density is space limitation. We tried hard to layout the main results while providing enough examples with intuitions. To address the reviewer's concern, we propose to not further increase the density and use extra space, if provided, to remark on the connections between the equations -- to improve the paper's self-contained narrative (also suggested by Reviewer qMTe) -- and remark on takeaways in the conclusion (suggested by Reviewer QMbw).