Test like you Train in Implicit Deep Learning

The paper investigates implicit deep learning, particularly in meta-learning and Deep Equilibrium Networks (DEQs). A common belief is that increasing the number of iterative solutions (inner iterations) during inference improves performance. This study challenges that notion, providing a theoretical analysis that highlights overparametrization as a crucial factor. The findings reveal that overparametrized networks, like DEQs, don't benefit from extra iterations at inference, while meta-learning does.

This paper focuses on the inner iterations overfitting problem of overparameterized models in implicit deep learning, and provides theoretical results in a simplified affine setting to show that increasing the number of iterations at test time cannot improve performance for overparametrized models. Two typical implicit deep learning methods, DEQ and iMAML, are considered in the paper. Experiments on diverse tasks verifies the theorem on both DEQ and iMAML.

This paper studies when increasing the number of fixed-point iterations in implicit deep learning during test time improves performance. In implicit deep learning, the network must implicitly minimize the objective $\ell(z(\theta))$. Specifically, the network parameters $\theta$ control some intermediary output $z(\theta)$ which tends to be the solution of some inner rootfinding problem $f(z, \theta, D_{train}) = 0$. The solution $z(\theta)$ is often identified through some $N$ fixed-point updates during training, then during test time, the solution is updated using test data using more fixed-point updates. This helps in certain applications like meta-learning and not others like DEQ.

The paper attempts to prove that this difference (helping in certain cases and not others) can be explained by the degree $\theta$ is overparameterized for the inner problem. They prove lower bounds for how much the training loss can change when changing the number of fixed-point iterations for an affine inner problem.

We thank the reviewers for engaging in the review process and hope that the rebuttal period will be fruitful. The reviewers praised our work for being “well written and easy to follow.” (YcSM) and found our “empirical results [...] support our hypotheses” (UkB2). One of the main points raised by the reviewers was the discrepancy between train and test. We give here a common answer: C1: Discrepancy train/test: Indeed, the discrepancy between train and test data is a major point that we highlighted (see page 6.). However, in the case of DEQs, we also provided a generalization result as stated on page 6 “We provide a generalization result in Corollary D.1 (in Appendix D) for DEQs stacked sample-wise with some additional smoothness assumptions.”. This result shows that if we have access to enough data and the functions $K_\text{in}(x_i), B(x_i), \ldots$ are sufficiently smooth, then theorem 1 holds on average for test data. This is equivalent to the kinds of results one usually gets when applying the Probably Approximately Correct (PAC) framework, which is generally used to discuss the generalization of machine learning approaches (our analysis is inspired by such literature) for example Mohri et al. 2018. The takeaway is that we proved that given sufficient data (which incidentally is also a necessary assumption for a model to generalize) our analysis holds in the test case. This is further validated by our experiments. We further highlight this aspect in our revised version by saying in the introduction “We also characterize the implications of this result for non-affine f in Theorem I.10, and for the generalization setting in Corollary D.1.” and also mention it in the empirical part. To be more specific on the evolution of overparametrization in the case of increasingly large training datasets, one can see that it’s possible to remain overparametrized even in this case, i.e. have $\mathcal{B}\_N(\mathcal{D}\_\text{train}) = 0$ (defined in eq 132 in Appendix). This typically happens in the trivial case where all the functions $U, B, K_\text{in}, \ldots$ are constant. More generally, the bound does not diverge even when increasing the number of data points if the data support is compact. We added a short paragraph covering this in the generalization appendix.

For this last day of the discussion period we would like to summarize our responses. Mainly we have answered the point raised by 2 reviewers on the discrepancy between train and test, highlighting our current results (both theoretical and empirical) on this point. We also proceeded to some clarifications regarding our results, and modified a Figure (fig 4) to better illustrate our point. We think that all the points raised by the reviewers have been addressed satisfactorily and we would be grateful to the reviewers should they consider raising their score to accept if they think similarly. Thanks again for your commitment to the review.

This paper studies the role of the number of inner iterations in implicit deep learning during test time on the performance. The paper attributes this role to the degree of overparameterized for the inner problem. They prove lower bounds for how much the training loss can change when changing the number of fixed-point iterations for a simple affine inner problem. The theory is validated on DEQs and meta-learning problems.

The reviewers found the paper clear to read, technically correct, and having good empirical support for their theory. However, multiple reviewers pointed out a major issue which is that the main theoretical result considers the training loss rather than the test loss, which is inconsistent with the main message of the paper. The authors pointed out a generalization result in Appendix D as a resolution of this issue, but that is based on a PAC-style argument and it isn't clear to the reviewers and the area chair how this applies to the setting considered, especially with overparameterization. The authors are encouraged to prove an explicit theorem on test data in future revisions.

Reject