Original: Path Independent Equilibrium Models Can Better Exploit Test-Time Computation
Revised: Path Independent Equilibrium Models Can Better Exploit Test-Time Computation

Original: Abstract
Revised: Abstract

Original: Designing networks capable of attaining better performance with an increased inference budget is important to facilitate generalization to harder problem instances.
Revised: Designing networks capable of attaining better performance with an increased inference budget is important to facilitate generalization to harder problem instances.

Original: Recent efforts have shown promising results in this direction by making use of depth-wise recurrent networks.
Revised: Recent efforts have shown promising results in this direction by making use of depth-wise recurrent networks.

Original: We show that a broad class of architectures named equilibrium models display strong upwards generalization, and find that stronger performance on harder examples (which require more iterations of inference to get correct) strongly correlates with the path independence of the system—its tendency to converge to the same steady-state behaviour regardless of initialization, given enough computation.
Revised: We show that a broad class of architectures named equilibrium models display strong upwards generalization, and find that stronger performance on harder examples (which require more iterations of inference to get correct) strongly correlates with the path independence of the system—its tendency to converge to the same steady-state behaviour regardless of initialization, given enough computation.

Original: Experimental interventions made to promote path independence result in improved generalization on harder problem instances, while those that penalize it degrade this ability.
Revised: Experimental interventions made to promote path independence result in improved generalization on harder problem instances, while those that penalize it degrade this ability.

Original: Path independence analyses are also useful on a per-example basis: for equilibrium models that have good in-distribution performance, path independence on out-of-distribution samples strongly correlates with accuracy.
Revised: Path independence analyses are also useful on a per-example basis: for equilibrium models that have good in-distribution performance, path independence on out-of-distribution samples strongly correlates with accuracy.

Original: Our results help explain why equilibrium models are capable of strong upwards generalization and motivates future work that harnesses path independence as a general modelling principle to facilitate scalable test-time usage.
Revised: Our results help explain why equilibrium models are capable of strong upwards generalization and motivates future work that harnesses path independence as a general modelling principle to facilitate scalable test-time usage.

Original:  Introduction
Revised:  Introduction

Original: One of the main challenges limiting the practical applicability of modern deep learning systems is the ability to generalize outside the training distribution [Koh et al., 2021].
Revised: One of the main challenges limiting the practical applicability of modern deep learning systems is the ability to generalize outside the training distribution [Koh et al., 2021].

Original: One particularly important type of out-of-distribution (OOD) generalization is upwards generalization , or the ability to generalize to more difficult problem instances than those encountered at training time [Selsam et al., 2018, Bansal et al., 2022, Schwarzschild et al., 2021b, Nye et al., 2021].
Revised: One particularly important type of out-of-distribution (OOD) generalization is upwards generalization , or the ability to generalize to more difficult problem instances than those encountered at training time [Selsam et al., 2018, Bansal et al., 2022, Schwarzschild et al., 2021b, Nye et al., 2021].

Original: Often, good performance on more difficult instances will require a larger amount of test-time computation, so a natural question arises: how can we design neural net architectures which can reliably exploit additional test-time computation to achieve better accuracy?
Revised: Often, good performance on more difficult instances will require a larger amount of test-time computation, so a natural question arises: how can we design neural net architectures which can reliably exploit additional test-time computation to achieve better accuracy?

Original: Equilibrium models, a broad class of architectures whose outputs are the fixed points of learned dynamical systems, are particularly suited to meet this challenge.
Revised: Equilibrium models, a broad class of architectures whose outputs are the fixed points of learned dynamical systems, are particularly suited to meet this challenge.

Original: Closely related to weight-tied recurrent models – networks that apply the same fixed neural network module repeatedly to hiddenlayer activations – equilibrium models are capable of adapting their compute budget based on the input they are given.
Revised: Closely related to weight-tied recurrent models – networks that apply the same fixed neural network module repeatedly to hiddenlayer activations – equilibrium models are capable of adapting their compute budget based on the input they are given.

Original: Under what conditions, if any, can this input-dependent ability to scale-up testtime compute actually lead to upwards generalization?
Revised: Under what conditions, if any, can this input-dependent ability to scale-up testtime compute actually lead to upwards generalization?

Original: We argue that a key determiner of whether a learned model can exploit additional test-time computation is whether the dynamical system corresponding to the model is path independent ; that is, whether the learned model’s hidden layer activations converge to the same asymptotic behaviour (i.e. fixed point or limit cycle), regardless of the initialization of the system.
Revised: We argue that a key determiner of whether a learned model can exploit additional test-time computation is whether the dynamical system corresponding to the model is path independent ; that is, whether the learned model’s hidden layer activations converge to the same asymptotic behaviour (i.e. fixed point or limit cycle), regardless of the initialization of the system.

Original: For example, a simple integrator x t +1 = x t + 1 is clearly not path independent, as its final state depends on the initial state x 0 and the number of iterations run; conversely, the system x t +1 =
Revised: For example, a simple integrator x t +1 = x t + 1 is clearly not path independent, as its final state depends on the initial state x 0 and the number of iterations run; conversely, the system x t +1 =

Original: ( x t + 1) / 2 is path independent, as it will converge to the solution x T = 1 as T → ∞
Revised: ( x t + 1) / 2 is path independent, as it will converge to the solution x T = 1 as T → ∞

Original: regardless of the initial condition of x 0 .
Revised: regardless of the initial condition of x 0 .

Original: Path independence is closely related to the concept of global stability control theory (
Revised: Path independence is closely related to the concept of global stability control theory (

Original: Section
Revised: Section

Original: Intuitively, path independent systems can more easily take advantage of additional test-time iterations than path dependent ones.
Revised: Intuitively, path independent systems can more easily take advantage of additional test-time iterations than path dependent ones.

Original: For instance, gradient descent applied to a convex objective is path independent, and correspondingly when confronted with a more ill-conditioned problem instance, one can compensate by increasing the number of iterations.
Revised: For instance, gradient descent applied to a convex objective is path independent, and correspondingly when confronted with a more ill-conditioned problem instance, one can compensate by increasing the number of iterations.

Original: Conversely, a weather simulation is path dependent, and extending the simulation won’t yield more accurate predictions of a given day’s weather.
Revised: Conversely, a weather simulation is path dependent, and extending the simulation won’t yield more accurate predictions of a given day’s weather.

Original: Based on this intuition, we hypothesize that path independence of a learned model is a key determiner of whether it can take advantage of an increased test-time iteration budget when generalizing to harder problem instances:
Revised: Based on this intuition, we hypothesize that path independence of a learned model is a key determiner of whether it can take advantage of an increased test-time iteration budget when generalizing to harder problem instances:

Original: Path Independence Hypothesis: Models which successfully fit the training distribution with a path independent function are better able to exploit more test-time iterations to achieve higher accuracy, compared to those which fit the training distribution with a path-dependent function.
Revised: Path Independence Hypothesis: Models which successfully fit the training distribution with a path independent function are better able to exploit more test-time iterations to achieve higher accuracy, compared to those which fit the training distribution with a path-dependent function.

Original: We first introduce a metric for measuring the path independence of a learned model, the Asymptotic
Revised: We first introduce a metric for measuring the path independence of a learned model, the Asymptotic

Original: Alignment (AA) score.
Revised: Alignment (AA) score.

Original: On a wide range of tasks including sequence prediction, visual reasoning, image classification, continuous optimization and graph regression, we show that a model’s AA score is strongly correlated with its performance when far more iterations are used at test time than at training time.
Revised: On a wide range of tasks including sequence prediction, visual reasoning, image classification, continuous optimization and graph regression, we show that a model’s AA score is strongly correlated with its performance when far more iterations are used at test time than at training time.

Original: In general, we find that path independent models increase their performance monotonically with the number of test-time iterations, whereas path dependent models degrade when the number of test time iterations exceeds those at training time.
Revised: In general, we find that path independent models increase their performance monotonically with the number of test-time iterations, whereas path dependent models degrade when the number of test time iterations exceeds those at training time.

Original: We find that input injection and implicit updates improve both the path independence and the accuracy.
Revised: We find that input injection and implicit updates improve both the path independence and the accuracy.

Original: Furthermore, we perform an experimental manipulation by introducing regularizers which directly promote or punish path independence.
Revised: Furthermore, we perform an experimental manipulation by introducing regularizers which directly promote or punish path independence.

Original: We find that these manipulations, while increasing or decreasing the AA score, also have the corresponding effect on accuracy.
Revised: We find that these manipulations, while increasing or decreasing the AA score, also have the corresponding effect on accuracy.

Original:  Background
Revised:  Background

Original: Equilibrium Models Equilibrium models treat computing internal representations as a fixed-point finding problem.
Revised: Equilibrium Models Equilibrium models treat computing internal representations as a fixed-point finding problem.

Original: [McClelland and Rumelhart, 1989, Liao et al., 2018, Bai et al., 2019]
Revised: [McClelland and Rumelhart, 1989, Liao et al., 2018, Bai et al., 2019]

Original: Concretely, letting x ∈ R n x and f w : R n x × n z (cid:55)→ R n z stand for an input and the equilibrium model function (or
Revised: Concretely, letting x ∈ R n x and f w : R n x × n z (cid:55)→ R n z stand for an input and the equilibrium model function (or

Original: “cell") parametrized by w ∈ R n w respectively, equilibrium models aim to solve for the fixed point z ∗ ∈ R n z that satisfies z ∗ =
Revised: “cell") parametrized by w ∈ R n w respectively, equilibrium models aim to solve for the fixed point z ∗ ∈ R n z that satisfies z ∗ =

Original: f w ( x , z ∗ ) .
Revised: f w ( x , z ∗ ) .

Original: The cell f w might represent anything from a fully connected layer to a transformer block [Vaswani et al., 2017].
Revised: The cell f w might represent anything from a fully connected layer to a transformer block [Vaswani et al., 2017].

Original: We emphasize that f w directly depends on the input x ; following existing literature, we refer to this as input injection .
Revised: We emphasize that f w directly depends on the input x ; following existing literature, we refer to this as input injection .

Original: The outcome of the fixed point finding process might depend on the initial fixed point guess.
Revised: The outcome of the fixed point finding process might depend on the initial fixed point guess.

Original: To make this dependency explicit, we define the function FIX : R n x × n z (cid:55)→ R n z that maps an input x and an initial guess for the fixed point z 0 to an output that satisfies the fixed point equation FIX f w ( x , z 0 ) :
Revised: To make this dependency explicit, we define the function FIX : R n x × n z (cid:55)→ R n z that maps an input x and an initial guess for the fixed point z 0 to an output that satisfies the fixed point equation FIX f w ( x , z 0 ) :

Original: = f w ( x , FIX f w ( x , z 0 )) .
Revised: = f w ( x , FIX f w ( x , z 0 )) .

Original: The behaviour of FIX depends on the solver that’s used to find fixed points.
Revised: The behaviour of FIX depends on the solver that’s used to find fixed points.

Original: The most straightforward approach to solve for fixed-points is the fixed point iteration method, which recursively applies the function f w on the internal representations z (i.e. z t +1 =
Revised: The most straightforward approach to solve for fixed-points is the fixed point iteration method, which recursively applies the function f w on the internal representations z (i.e. z t +1 =

Original: f w ( x , z t ) ).
Revised: f w ( x , z t ) ).

Original: If certain conditions are satisfied (such as the fixed iterations being globally contractive – more general conditions are discussed below), this procedure converges 3 to a fixed point: z ∗ = f w ( x , z ∗ ) .
Revised: If certain conditions are satisfied (such as the fixed iterations being globally contractive – more general conditions are discussed below), this procedure converges 3 to a fixed point: z ∗ = f w ( x , z ∗ ) .

Original: As solving for fixed points exactly is expensive, fixed point iterations are often terminated after a fixed number of steps or when the norm of the difference between subsequent iterates falls below a preselected threshold.
Revised: As solving for fixed points exactly is expensive, fixed point iterations are often terminated after a fixed number of steps or when the norm of the difference between subsequent iterates falls below a preselected threshold.

Original: The model weights can be updated using gradients computed via backpropagating through the full forward computational graph.
Revised: The model weights can be updated using gradients computed via backpropagating through the full forward computational graph.

Original: If one commits to using fixed point iterations as the root solver, then the output z ∗ of equilibrium models can be interpreted as the infinite-depth limit of an input-injected, weight-tied model f ∞ w ( x , z ∗ ) =
Revised: If one commits to using fixed point iterations as the root solver, then the output z ∗ of equilibrium models can be interpreted as the infinite-depth limit of an input-injected, weight-tied model f ∞ w ( x , z ∗ ) =

Original: lim n →∞ f ( n ) w ( x , z 0 ) where the notation f ( n ) stands for n repeated applications of f on its own output, and z 0 stands for the fixed point initialization.
Revised: lim n →∞ f ( n ) w ( x , z 0 ) where the notation f ( n ) stands for n repeated applications of f on its own output, and z 0 stands for the fixed point initialization.

Original: Implicit Training of Equilibrium Models
Revised: Implicit Training of Equilibrium Models

Original: Different training algorithms for equilibrium models can be derived by considering their implicit nature.
Revised: Different training algorithms for equilibrium models can be derived by considering their implicit nature.

Original: Bai et al.
Revised: Bai et al.

Original: [2019] solve for fixed points explicitly using black-box root finders, such as Broyden’s method [Broyden, 1965] or Anderson acceleration [Anderson, 1965].
Revised: [2019] solve for fixed points explicitly using black-box root finders, such as Broyden’s method [Broyden, 1965] or Anderson acceleration [Anderson, 1965].

Original: In order to avoid explicitly differentiating through the root-finding procedure, they utilize implicit differentiation to compute gradients.
Revised: In order to avoid explicitly differentiating through the root-finding procedure, they utilize implicit differentiation to compute gradients.

Original: 4 Concretely, letting g w ( x , z ) =
Revised: 4 Concretely, letting g w ( x , z ) =

Original: f w ( x , z ) − z for a fixed point z ∗ , the Jacobian of z ∗ with respect to the equilibrium model weights can be given by:
Revised: f w ( x , z ) − z for a fixed point z ∗ , the Jacobian of z ∗ with respect to the equilibrium model weights can be given by:

Original: [Equation]
Revised: [Equation]

Original: Inverting a Jacobian matrix can become computationally expensive.
Revised: Inverting a Jacobian matrix can become computationally expensive.

Original: Recent works [Geng et al., 2021a,
Revised: Recent works [Geng et al., 2021a,

Original: Fung et al., 2021] have shown that the inverse-Jacobian term in Eq.
Revised: Fung et al., 2021] have shown that the inverse-Jacobian term in Eq.

Original: (1) can be replaced with an identity matrix i.e. Jacobian-free or an approximate inverse-Jacobian
Revised: (1) can be replaced with an identity matrix i.e. Jacobian-free or an approximate inverse-Jacobian

Original: [Geng et al., 2021b] without affecting the final performance.
Revised: [Geng et al., 2021b] without affecting the final performance.

Original: This approximation makes the backward pass inexpensive and lightweight.
Revised: This approximation makes the backward pass inexpensive and lightweight.

Original: Equilibrium models have been shown to achieve state-of-the-art performance on various tasks including language modelling [Bai et al., 2019], image recognition, semantic segmentation [Bai et al., 2020], object detection [Wang et al., 2020], and graph modeling [Gu et al., 2020, Liu et al., 2021a].
Revised: Equilibrium models have been shown to achieve state-of-the-art performance on various tasks including language modelling [Bai et al., 2019], image recognition, semantic segmentation [Bai et al., 2020], object detection [Wang et al., 2020], and graph modeling [Gu et al., 2020, Liu et al., 2021a].

Original: Equilibrium Models vs. Depthwise Recurrent Models Both equilibrium models and input-injected depthwise recurrent (i.e. weight-tied, fixed-depth) networks leverage weight-tying i.e., they apply the same transformation at each layer, f
Revised: Equilibrium Models vs. Depthwise Recurrent Models Both equilibrium models and input-injected depthwise recurrent (i.e. weight-tied, fixed-depth) networks leverage weight-tying i.e., they apply the same transformation at each layer, f

Original: [ i ] ω = f ω ∀ i .
Revised: [ i ] ω = f ω ∀ i .

Original: The two models differ in the ultimate aim of the forward pass: while depthwise recurrent models compute a (weight-tied) fixed depth computation (which may or may not approach a fixed point), the stated goal of equilibrium models is explicitly to find a fixed point.
Revised: The two models differ in the ultimate aim of the forward pass: while depthwise recurrent models compute a (weight-tied) fixed depth computation (which may or may not approach a fixed point), the stated goal of equilibrium models is explicitly to find a fixed point.

Original: Weight-tied fixed depth networks by definition require backpropagation through an explicit stack of layers.
Revised: Weight-tied fixed depth networks by definition require backpropagation through an explicit stack of layers.

Original: Equilibrium models, however, directly solve for fixed points using (potentially black-box) solvers during the forward pass and may be trained using implicit differentiation.
Revised: Equilibrium models, however, directly solve for fixed points using (potentially black-box) solvers during the forward pass and may be trained using implicit differentiation.

Original: Convergence As alluded above, in order to guarantee convergence to a unique fixed point, it suffices for the cell of the equilibrium model to be contractive over its input domain 5 (i.e. the singular values of its Jacobian all lie below 1 ).
Revised: Convergence As alluded above, in order to guarantee convergence to a unique fixed point, it suffices for the cell of the equilibrium model to be contractive over its input domain 5 (i.e. the singular values of its Jacobian all lie below 1 ).

Original: Previous work has leveraged Lipschitz constrained cells to ensure contractivity [Revay et al., 2020].
Revised: Previous work has leveraged Lipschitz constrained cells to ensure contractivity [Revay et al., 2020].

Original: Other approaches for ensuring global convergence exist: the monotone equilibrium model architecture guarantees global convergence by utilizing an equilibrium model parametrization that bears similarities to solutions to a particular form of monotone operator splitting problem [Winston and Kolter, 2020].
Revised: Other approaches for ensuring global convergence exist: the monotone equilibrium model architecture guarantees global convergence by utilizing an equilibrium model parametrization that bears similarities to solutions to a particular form of monotone operator splitting problem [Winston and Kolter, 2020].

Original: Unrestricted equilibrium models aren’t constrained enough to guarantee convergence: they can easily express globally divergent vector fields that prohibit the existence of fixed points.
Revised: Unrestricted equilibrium models aren’t constrained enough to guarantee convergence: they can easily express globally divergent vector fields that prohibit the existence of fixed points.

Original: It is, therefore, interesting that they can (and often) learn path independent solutions.
Revised: It is, therefore, interesting that they can (and often) learn path independent solutions.

Original: Also note that the “infinite-depth weight tied network" interpretation of equilibrium models is less general than the implicit formulation presented above, as the latter admits unstable fixed points as well.
Revised: Also note that the “infinite-depth weight tied network" interpretation of equilibrium models is less general than the implicit formulation presented above, as the latter admits unstable fixed points as well.

Original: Terminology and Abbreviations We use the term “equilibrium models” to refer to the general class of networks that explicitly solve for a fixed point in the forward pass.
Revised: Terminology and Abbreviations We use the term “equilibrium models” to refer to the general class of networks that explicitly solve for a fixed point in the forward pass.

Original: We use the term ‘solver’ to refer to the use of black-box root finders like Anderson acceleration to find fixed points of an implicitly trained equilibrium model.
Revised: We use the term ‘solver’ to refer to the use of black-box root finders like Anderson acceleration to find fixed points of an implicitly trained equilibrium model.

Original: These networks can use implicit gradients computed via implicit function theorem (IFT), Jacobian-free backward pass, or with an approximation of inverse-Jacobian.
Revised: These networks can use implicit gradients computed via implicit function theorem (IFT), Jacobian-free backward pass, or with an approximation of inverse-Jacobian.

Original: The term ‘unroll’ refers to equilibrium models that use regular fixed-point iterations to compute the equilibrium point.
Revised: The term ‘unroll’ refers to equilibrium models that use regular fixed-point iterations to compute the equilibrium point.

Original: We use the abbreviation ‘bp’ to refer to backpropagation gradients, and ‘inj’ to refer to input injection.
Revised: We use the abbreviation ‘bp’ to refer to backpropagation gradients, and ‘inj’ to refer to input injection.

Original: The term ‘progressive net’ refers to the deep thinking networks trained with progressive training as proposed by Bansal et al.
Revised: The term ‘progressive net’ refers to the deep thinking networks trained with progressive training as proposed by Bansal et al.

Original: We use ‘PI’ and ‘non-PI’ to refer to path independent and path dependent networks, respectively.
Revised: We use ‘PI’ and ‘non-PI’ to refer to path independent and path dependent networks, respectively.

Original:  Upwards Generalization with Equilibrium Models
Revised:  Upwards Generalization with Equilibrium Models

Original: In this section, we establish that equilibrium models are capable of strong upwards generalization.
Revised: In this section, we establish that equilibrium models are capable of strong upwards generalization.

Original: To study the effects of test time computation, it is useful to consider tasks with an explicit difficulty parameter, so that the learned models can be tested on more difficult instances which require a large number of iterations to solve correctly.
Revised: To study the effects of test time computation, it is useful to consider tasks with an explicit difficulty parameter, so that the learned models can be tested on more difficult instances which require a large number of iterations to solve correctly.

Original: We focus on multiple algorithmic generalization tasks: prefix sum and mazes by Schwarzschild et al.
Revised: We focus on multiple algorithmic generalization tasks: prefix sum and mazes by Schwarzschild et al.

Original: [2021a,b], blurry MNIST , matrix inversion and edge copy by Du et al.
Revised: [2021a,b], blurry MNIST , matrix inversion and edge copy by Du et al.

Original: Taken together, these tasks cover a wide range of problems from different domains, namely sequence prediction, visual reasoning, image classification, continuous optimization and graph regression.
Revised: Taken together, these tasks cover a wide range of problems from different domains, namely sequence prediction, visual reasoning, image classification, continuous optimization and graph regression.

Original: To maintain clarity and focus, we run our detailed analysis on the prefix sum and mazes tasks, and provide complementary results for the remaining tasks in the Supplementary Material (SM).
Revised: To maintain clarity and focus, we run our detailed analysis on the prefix sum and mazes tasks, and provide complementary results for the remaining tasks in the Supplementary Material (SM).

Original: Tasks Prefix-sum is a sequence-to-sequence task whereby the network is given a sequence of 0bits, and is trained to output, for each bit, the parity of all of the bits received since the beginning of the sequence until the current bit.
Revised: Tasks Prefix-sum is a sequence-to-sequence task whereby the network is given a sequence of 0bits, and is trained to output, for each bit, the parity of all of the bits received since the beginning of the sequence until the current bit.

Original: We train on 10,000 unique 32-bit binary strings, and report results on binary strings of other lengths.
Revised: We train on 10,000 unique 32-bit binary strings, and report results on binary strings of other lengths.

Original: The mazes task is also an image-to-image task, where the input is a three-channel RGB image.
Revised: The mazes task is also an image-to-image task, where the input is a three-channel RGB image.

Original: The ‘start’ and ‘finish’ positions are marked by a red and a green square respectively; walls are marked in black.
Revised: The ‘start’ and ‘finish’ positions are marked by a red and a green square respectively; walls are marked in black.

Original: The output is the optimal path in the maze that connects these two points without passing through the walls.
Revised: The output is the optimal path in the maze that connects these two points without passing through the walls.

Original: We train on 50,000 small mazes of size 9 × 9 , and report upward generalization results on larger mazes.
Revised: We train on 50,000 small mazes of size 9 × 9 , and report upward generalization results on larger mazes.

Original: Instances of each of these problems, as well as additional image classification and continuous optimization results can be found in the supplementary material.
Revised: Instances of each of these problems, as well as additional image classification and continuous optimization results can be found in the supplementary material.

Original: Blurry MNIST [Liang et al.] is a robustness-to-corruption task: one has to learn to do
Revised: Blurry MNIST [Liang et al.] is a robustness-to-corruption task: one has to learn to do

Original: MNIST classification from lightly blurred images and generalize zero-shot to highly blurred ones.
Revised: MNIST classification from lightly blurred images and generalize zero-shot to highly blurred ones.

Original: In the matrix inversion task [Du et al., 2022], the goal is to learn to invert 10 × 10 matrices in a way that generalizes to matrices that have worse condition number than those observed during training.
Revised: In the matrix inversion task [Du et al., 2022], the goal is to learn to invert 10 × 10 matrices in a way that generalizes to matrices that have worse condition number than those observed during training.

Original: Edge copy [Du et al., 2022] is a simple graph regression tasks that requires learning to output the input edge features, in a way that generalizes to larger graph sizes.
Revised: Edge copy [Du et al., 2022] is a simple graph regression tasks that requires learning to output the input edge features, in a way that generalizes to larger graph sizes.

Original: Note that the training and test data in the latter two tasks are generated with noise added on-the-fly, as done by Du et al.
Revised: Note that the training and test data in the latter two tasks are generated with noise added on-the-fly, as done by Du et al.

Original: Strong Upward Generalization Fig.
Revised: Strong Upward Generalization Fig.

Original: 2a shows that equilibrium models demonstrate very strong upward generalization performance compared to non weight-tied fixed-depth models.
Revised: 2a shows that equilibrium models demonstrate very strong upward generalization performance compared to non weight-tied fixed-depth models.

Original: Moreover,
Revised: Moreover,

Original: Fig.
Revised: Fig.

Original: 2b shows that increasing inference depth consistently improves performance—especially on harder problem instances.
Revised: 2b shows that increasing inference depth consistently improves performance—especially on harder problem instances.

Original:  Path Independence
Revised:  Path Independence

Original: Having intuitively motivated the idea of path independence in Sec. 1, we now define it formally: we say that the computation performed by a recurrent operator computing function f w on an input x is path independent if it converges to the same limiting behavior regardless of the current state z t .
Revised: Having intuitively motivated the idea of path independence in Sec. 1, we now define it formally: we say that the computation performed by a recurrent operator computing function f w on an input x is path independent if it converges to the same limiting behavior regardless of the current state z t .

Original: As a special case, if the computation is convergent, this property is equivalent to the existence of a unique fixed point z ∗ such that f ∞ w ( x , z 0 ) =
Revised: As a special case, if the computation is convergent, this property is equivalent to the existence of a unique fixed point z ∗ such that f ∞ w ( x , z 0 ) =

Original: z ∗ for any z 0 .
Revised: z ∗ for any z 0 .

Original: However, our definition allows for other behaviors such as limit cycles (see Sec. 7).
Revised: However, our definition allows for other behaviors such as limit cycles (see Sec. 7).

Original: [Figure]
Revised: [Figure]

Original: Some architectures guarantee the path independence property (see Sec. 2).
Revised: Some architectures guarantee the path independence property (see Sec. 2).

Original: However, most common
Revised: However, most common

Original: DEQ architectures—and the ones we use throughout this paper—have the expressive power to learn multiple fixed points per input.
Revised: DEQ architectures—and the ones we use throughout this paper—have the expressive power to learn multiple fixed points per input.

Original: Since it is unclear whether architectures enforcing the contraction property lose expressiveness [Bai et al., 2021], we focus our investigation on unrestricted architectures.
Revised: Since it is unclear whether architectures enforcing the contraction property lose expressiveness [Bai et al., 2021], we focus our investigation on unrestricted architectures.

Original: PI networks represent a different model for computation than standard feed-forward networks: instead of learning an entire computational graph to map inputs to outputs, they only have to learn where to stop .
Revised: PI networks represent a different model for computation than standard feed-forward networks: instead of learning an entire computational graph to map inputs to outputs, they only have to learn where to stop .

Original: We dedicate the rest of the paper on exploring the Path Independence Hypothesis —the idea that models which successfully fit the training distribution with a path independent function are better able to exploit more test-time iterations to achieve higher accuracy, compared to those which fit the training distribution with a path-dependent function.
Revised: We dedicate the rest of the paper on exploring the Path Independence Hypothesis —the idea that models which successfully fit the training distribution with a path independent function are better able to exploit more test-time iterations to achieve higher accuracy, compared to those which fit the training distribution with a path-dependent function.

Original: Before establishing a connection between path independence and out-of-distribution generalization, we first describe two architectural components that are necessary conditions for achieving path independence.
Revised: Before establishing a connection between path independence and out-of-distribution generalization, we first describe two architectural components that are necessary conditions for achieving path independence.

Original: Afterward, we describe a metric to quantify how path-independent a trained network is.
Revised: Afterward, we describe a metric to quantify how path-independent a trained network is.

Original: 4.1 Architectural Components Necessary for Path Independence
Revised: 4.1 Architectural Components Necessary for Path Independence

Original: Past work has observed that weight tying and input injection are both crucial for upwards generalization [Bansal et al., 2022].
Revised: Past work has observed that weight tying and input injection are both crucial for upwards generalization [Bansal et al., 2022].

Original: We observe that both architectural components are also necessary for a learned model to be PI.
Revised: We observe that both architectural components are also necessary for a learned model to be PI.

Original: 6 Without weight tying, the network is constrained to have a fixed forward depth, so it is meaningless to talk about the limiting behavior in large depth.
Revised: 6 Without weight tying, the network is constrained to have a fixed forward depth, so it is meaningless to talk about the limiting behavior in large depth.

Original: Input injection ensures that the equilibrium point depends on the input despite having an “infinite depth".
Revised: Input injection ensures that the equilibrium point depends on the input despite having an “infinite depth".

Original: Without input injection, a PI network would necessarily forget the input; hence, any model which successfully fits the training distribution must be path dependent.
Revised: Without input injection, a PI network would necessarily forget the input; hence, any model which successfully fits the training distribution must be path dependent.

Original: Interestingly, both architectural motifs are also key components of deep equilibrium models [Bai et al., 2019]; in that work, the motivation was to enable efficient gradient estimation via the implicit function theorem (IFT) — a concept closely related to path independence, since the premise of the
Revised: Interestingly, both architectural motifs are also key components of deep equilibrium models [Bai et al., 2019]; in that work, the motivation was to enable efficient gradient estimation via the implicit function theorem (IFT) — a concept closely related to path independence, since the premise of the

Original: IFT gradient estimator is that only the final hidden state matters, not the path taken to get there.
Revised: IFT gradient estimator is that only the final hidden state matters, not the path taken to get there.

Original: It is striking that two separate lines of work would converge on the same architectural motifs, one motivated by generalization and the other by a variant of path independence.
Revised: It is striking that two separate lines of work would converge on the same architectural motifs, one motivated by generalization and the other by a variant of path independence.

Original: Reproducing the results of Bansal et al.
Revised: Reproducing the results of Bansal et al.

Original: [2022], in Fig.
Revised: [2022], in Fig.

Original: 2a we show upward generalization performance using both equilibrium models and progressive nets [Bansal et al., 2022] – and the lack thereof using non-input-injected networks.
Revised: 2a we show upward generalization performance using both equilibrium models and progressive nets [Bansal et al., 2022] – and the lack thereof using non-input-injected networks.

Original: For the remainder of this paper, we focus on architectures with both input injection and weight tying.
Revised: For the remainder of this paper, we focus on architectures with both input injection and weight tying.

Original: Input: A batched input (cid:20) xx(cid:21) , an operator f w
Revised: Input: A batched input (cid:20) xx(cid:21) , an operator f w

Original: Initialize: (cid:20) zz(cid:21) =
Revised: Initialize: (cid:20) zz(cid:21) =

Original: Define: h ( y 1 , y 2 ) = y∥ y 1 ∥ 2 · y∥ y 2 ∥
Revised: Define: h ( y 1 , y 2 ) = y∥ y 1 ∥ 2 · y∥ y 2 ∥

Original: Compute (cid:20) z ′z ′(cid:21) =
Revised: Compute (cid:20) z ′z ′(cid:21) =

Original: FIX f w (cid:18)(cid:20) xx(cid:21) , (cid:20) zz(cid:21)(cid:19)
Revised: FIX f w (cid:18)(cid:20) xx(cid:21) , (cid:20) zz(cid:21)(cid:19)

Original: # Interchange and reinitialize iterates
Revised: # Interchange and reinitialize iterates

Original: Compute (cid:20) z ′′z ′′(cid:21) =
Revised: Compute (cid:20) z ′′z ′′(cid:21) =

Original: FIX f w (cid:18)(cid:20) xx(cid:21) , (cid:20) z ′z ′(cid:21)(cid:19) return average( h ( z ′′ 1 , z ′ 1 ) , h ( z ′′ 2 , z ′ 2 ))
Revised: FIX f w (cid:18)(cid:20) xx(cid:21) , (cid:20) z ′z ′(cid:21)(cid:19) return average( h ( z ′′ 1 , z ′ 1 ) , h ( z ′′ 2 , z ′ 2 ))

Original: 4.2 Quantifying Path Independence
Revised: 4.2 Quantifying Path Independence

Original: We propose a simple metric to quantify path independence based on the directional alignment of the fixed points computed with the same input, but different initializations.
Revised: We propose a simple metric to quantify path independence based on the directional alignment of the fixed points computed with the same input, but different initializations.

Original: We name this metric the Asymptotic Alignment (AA) score .
Revised: We name this metric the Asymptotic Alignment (AA) score .

Original: Pseudocode to compute the metric is given in Alg.
Revised: Pseudocode to compute the metric is given in Alg.

Original: The
Revised: The

Original: AA score is the average cosine similarity between the fixed points obtained with the training time initialization (often simply the zero vector) and the fixed points obtained when one initializes the solver using the fixed points computed on different inputs .
Revised: AA score is the average cosine similarity between the fixed points obtained with the training time initialization (often simply the zero vector) and the fixed points obtained when one initializes the solver using the fixed points computed on different inputs .

Original: Higher AA scores (with 1 being the highest value) imply higher degrees of path independence.
Revised: Higher AA scores (with 1 being the highest value) imply higher degrees of path independence.

Original: In Sec. 5, we show a strong correlation between path independence and accuracy using the AA score.
Revised: In Sec. 5, we show a strong correlation between path independence and accuracy using the AA score.

Original: The AA score is cheap to compute, is a reliable indicator of path independence (see below), and is unitless, meaning that networks obtained from different training runs can be compared on equal footing.
Revised: The AA score is cheap to compute, is a reliable indicator of path independence (see below), and is unitless, meaning that networks obtained from different training runs can be compared on equal footing.

Original: See the supplementary material for other metrics we’ve considered for quantifying path independence and why we found AA score to be preferable.
Revised: See the supplementary material for other metrics we’ve considered for quantifying path independence and why we found AA score to be preferable.

Original: Stress-testing the AA score To stress-test the extent to which the AA score really measures path independence, we search for adversarial initializations that are optimized to result in distinct fixed points, hence low AA values.
Revised: Stress-testing the AA score To stress-test the extent to which the AA score really measures path independence, we search for adversarial initializations that are optimized to result in distinct fixed points, hence low AA values.

Original: (Unlike adversarial examples, this attack is not constrained to an ε -ball.)
Revised: (Unlike adversarial examples, this attack is not constrained to an ε -ball.)

Original: We use the L-BFGS
Revised: We use the L-BFGS

Original: [Liu and Nocedal, 1989] optimizer, and repeat the search multiple times starting from different fixed point initializations.
Revised: [Liu and Nocedal, 1989] optimizer, and repeat the search multiple times starting from different fixed point initializations.

Original: We include pseudocode in the supplementary material.
Revised: We include pseudocode in the supplementary material.

Original: Results of the adversarial stress test can be seen in Tab. 1.
Revised: Results of the adversarial stress test can be seen in Tab. 1.

Original: The results corroborate that the AA score is indeed a reliable measure of path independence; while it isn’t possible to find adversarial initializations for high AA score networks (indicating high path independence), low AA score networks can easily be adversarially initialized to be steered away from the original fixed point estimate.
Revised: The results corroborate that the AA score is indeed a reliable measure of path independence; while it isn’t possible to find adversarial initializations for high AA score networks (indicating high path independence), low AA score networks can easily be adversarially initialized to be steered away from the original fixed point estimate.

Original:  Path Independence Correlates with Upward Generalization
Revised:  Path Independence Correlates with Upward Generalization

Original: Is path independence (as measured by the AA score) a strong predictor of upwards generalization?
Revised: Is path independence (as measured by the AA score) a strong predictor of upwards generalization?

Original: We took the trained networks from Sec. 3, computed their average AA scores on in- and out-ofdistribution splits and inspected whether the AA scores are correlated with upward generalization.
Revised: We took the trained networks from Sec. 3, computed their average AA scores on in- and out-ofdistribution splits and inspected whether the AA scores are correlated with upward generalization.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: Accuracies and AA scores are capped at 0.999 for compatibility with the probit transform.
Revised: Accuracies and AA scores are capped at 0.999 for compatibility with the probit transform.

Original: On prefix sum experiments, we varied 1) network depth, 2) whether or not weight norm (wnorm)
Revised: On prefix sum experiments, we varied 1) network depth, 2) whether or not weight norm (wnorm)

Original: [Salimans and Kingma, 2016] was used or not, 7 3) learning rate (one of [0.01, 0.001, 0.0001]), 4) forward solver (fixed point iterations or Anderson acceleration [Anderson, 1965], and 5) the gradient estimator (backprop or implicit gradients).
Revised: [Salimans and Kingma, 2016] was used or not, 7 3) learning rate (one of [0.01, 0.001, 0.0001]), 4) forward solver (fixed point iterations or Anderson acceleration [Anderson, 1965], and 5) the gradient estimator (backprop or implicit gradients).

Original: On the maze experiments, we varied 1) network depth, 2) use of weight norm, 3) forward solver (fixed point iterations or Broyden solver [Broyden, 1965]), and 5) the gradient estimator (backprop or implicit gradients).
Revised: On the maze experiments, we varied 1) network depth, 2) use of weight norm, 3) forward solver (fixed point iterations or Broyden solver [Broyden, 1965]), and 5) the gradient estimator (backprop or implicit gradients).

Original: Fig.
Revised: Fig.

Original: 4 displays our findings.
Revised: 4 displays our findings.

Original: We evaluated performance on a mixture of in- and OOD validation data; results on individual data splits can be found in the supplementary material.
Revised: We evaluated performance on a mixture of in- and OOD validation data; results on individual data splits can be found in the supplementary material.

Original: The results show a strong correlation between AA score and accuracy when the inference depth is large enough.
Revised: The results show a strong correlation between AA score and accuracy when the inference depth is large enough.

Original: This shows that PI networks allow for scaling test-time compute to improve test-time accuracy (see also
Revised: This shows that PI networks allow for scaling test-time compute to improve test-time accuracy (see also

Original: Fig. 2b).
Revised: Fig. 2b).

Original: The in-distribution validation performance of non-PI networks degrades with deeper inference depths.
Revised: The in-distribution validation performance of non-PI networks degrades with deeper inference depths.

Original: Unsurprisingly, these networks generalize poorly on harder problem instances that require deeper inference depths (i.e. problem instances that provably require at least a given number of layers to handle).
Revised: Unsurprisingly, these networks generalize poorly on harder problem instances that require deeper inference depths (i.e. problem instances that provably require at least a given number of layers to handle).

Original: Further results on the BlurryMNIST, matrix inversion and edge copy tasks can be found in Supplementary Material H, I and J.
Revised: Further results on the BlurryMNIST, matrix inversion and edge copy tasks can be found in Supplementary Material H, I and J.

Original:  Experimental Manipulations of Path Independence
Revised:  Experimental Manipulations of Path Independence

Original: The previous section demonstrates a strong correlation between path independence and the ability to exploit additional test-time iterations.
Revised: The previous section demonstrates a strong correlation between path independence and the ability to exploit additional test-time iterations.

Original: Unfortunately, we can’t make a causal claim based on these studies: the observed effect could have been due to an unobserved confounder.
Revised: Unfortunately, we can’t make a causal claim based on these studies: the observed effect could have been due to an unobserved confounder.

Original: In this section, we intervene directly on path independence by imposing regularizers which directly encourage or penalize path independence.
Revised: In this section, we intervene directly on path independence by imposing regularizers which directly encourage or penalize path independence.

Original: We find that interventions designed to promote path independence also improve generalization, while interventions designed to reduce path independence also hurt generalization.
Revised: We find that interventions designed to promote path independence also improve generalization, while interventions designed to reduce path independence also hurt generalization.

Original: 6.1 Promoting Path Independence via Randomized Forward Passes
Revised: 6.1 Promoting Path Independence via Randomized Forward Passes

Original: A straightforward way to encourage path independence is simply to initialize the hidden states with random noise during training.
Revised: A straightforward way to encourage path independence is simply to initialize the hidden states with random noise during training.

Original: To this end, we experimented with initializing the hidden states with zeros on half of the examples in the batch, and with standard Gaussian noise on the rest of the examples.
Revised: To this end, we experimented with initializing the hidden states with zeros on half of the examples in the batch, and with standard Gaussian noise on the rest of the examples.

Original: The reason to include the zero-initializations at training time is that we initialize from zeros at test time - not including this initialization during traning time causes a distribution shift.
Revised: The reason to include the zero-initializations at training time is that we initialize from zeros at test time - not including this initialization during traning time causes a distribution shift.

Original: Another way to promote path independence is simply running the forward solver with randomized compute budgets/depths during training time.
Revised: Another way to promote path independence is simply running the forward solver with randomized compute budgets/depths during training time.

Original: While a path independent solution can be expected to be robust against this intervention, a path dependent one will fail.
Revised: While a path independent solution can be expected to be robust against this intervention, a path dependent one will fail.

Original: We took the training configurations of the 12 prefix-sum networks described in Sec. 5 that use fixed point iterations in their forward pass, and backpropagation gradient in their backward pass, and retrained them separately with the aforementioned mixed initialization and randomized depth strategies without modifying any other experimental conditions.
Revised: We took the training configurations of the 12 prefix-sum networks described in Sec. 5 that use fixed point iterations in their forward pass, and backpropagation gradient in their backward pass, and retrained them separately with the aforementioned mixed initialization and randomized depth strategies without modifying any other experimental conditions.

Original: As can be seen in Fig.
Revised: As can be seen in Fig.

Original: 3, the interventions lead to strong test-time path independent neural networks, while also reliably improving in- and out-of-distribution validation accuracy.
Revised: 3, the interventions lead to strong test-time path independent neural networks, while also reliably improving in- and out-of-distribution validation accuracy.

Original: We especially emphasize that shallow networks trained with mixed initialization actually remain far from having high AA scores using the trainingtime forward pass conditions due to lack of convergence.
Revised: We especially emphasize that shallow networks trained with mixed initialization actually remain far from having high AA scores using the trainingtime forward pass conditions due to lack of convergence.

Original: However, since the mixed initialization strategy results in path independent networks, scaling up test-time compute budget leads to high AA scores, and therefore high upwards generalization.
Revised: However, since the mixed initialization strategy results in path independent networks, scaling up test-time compute budget leads to high AA scores, and therefore high upwards generalization.

Original: 6.2 Penalizing Path Independence via the Fixed Point Alignment Penalty
Revised: 6.2 Penalizing Path Independence via the Fixed Point Alignment Penalty

Original: Does an intervention that results in less path independence also result in poorer upwards generalization?
Revised: Does an intervention that results in less path independence also result in poorer upwards generalization?

Original: Like in the mixed initialization experiment, we retrained the 12 unroll + backpropagation networks with an additional auxiliary loss term that penalizes the dot product between the fixed points computed from the same input, but different initializations sampled from standard Gaussian noise.
Revised: Like in the mixed initialization experiment, we retrained the 12 unroll + backpropagation networks with an additional auxiliary loss term that penalizes the dot product between the fixed points computed from the same input, but different initializations sampled from standard Gaussian noise.

Original: Fig.
Revised: Fig.

Original: 3 shows that this intervention succeeded in pushing the AA scores down, while also keeping the accuracy on the same trend line.
Revised: 3 shows that this intervention succeeded in pushing the AA scores down, while also keeping the accuracy on the same trend line.

Original:  Disambiguating Convergence and Path Independence
Revised:  Disambiguating Convergence and Path Independence

Original: Is convergence necessary for path independence?
Revised: Is convergence necessary for path independence?

Original: We answer this statement in the negative, and show that neither training-time convergence nor test-time convergence is required for path independence.
Revised: We answer this statement in the negative, and show that neither training-time convergence nor test-time convergence is required for path independence.

Original: Instead convergence to the same limiting behavior regardless of initialization is important.
Revised: Instead convergence to the same limiting behavior regardless of initialization is important.

Original: Training Time Convergence We consider two implicitly trained equilibrium models trained on the mazes task—one trained with implicit gradients computed via implicit function theorem (IFT), and the other trained with an approximation of the (inverse) Jacobian, called phantom gradients [Geng et al., 2021b].
Revised: Training Time Convergence We consider two implicitly trained equilibrium models trained on the mazes task—one trained with implicit gradients computed via implicit function theorem (IFT), and the other trained with an approximation of the (inverse) Jacobian, called phantom gradients [Geng et al., 2021b].

Original: We report the values of residuals ( i.e., ∥ f w ( x , z ) − z ∥ 2 ), AA scores and accuracies observed for in- and out-of-distribution data for mazes in Tab. 2.
Revised: We report the values of residuals ( i.e., ∥ f w ( x , z ) − z ∥ 2 ), AA scores and accuracies observed for in- and out-of-distribution data for mazes in Tab. 2.

Original: We observe that DEQs trained with phantom gradients have higher values of in-distribution residuals but are path independent, as indicated by their high AA scores, and show strong upward generalization as indicated by their good accuracy.
Revised: We observe that DEQs trained with phantom gradients have higher values of in-distribution residuals but are path independent, as indicated by their high AA scores, and show strong upward generalization as indicated by their good accuracy.

Original: The mixed-initialization intervention described in Sec. 6.1 also leads to a separation between trainingtime convergence and path independence.
Revised: The mixed-initialization intervention described in Sec. 6.1 also leads to a separation between trainingtime convergence and path independence.

Original: We found that it is possible to train very shallow ( i.e.,layer) unrolled networks that, while being very far from converging during training and attaining poor in-distribution generalization, are able to converge and achieve perfect performance when run for many more iterations during test time.
Revised: We found that it is possible to train very shallow ( i.e.,layer) unrolled networks that, while being very far from converging during training and attaining poor in-distribution generalization, are able to converge and achieve perfect performance when run for many more iterations during test time.

Original: Details are provided in the supplementary material.
Revised: Details are provided in the supplementary material.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: Test Time Convergence From Tab. 2, one might conclude that test time convergence is important for path independence.
Revised: Test Time Convergence From Tab. 2, one might conclude that test time convergence is important for path independence.

Original: However, we show that this connection is not necessary, and convergence to the same fixed point is not a required condition for path independence.
Revised: However, we show that this connection is not necessary, and convergence to the same fixed point is not a required condition for path independence.

Original: We study test time convergence properties of an unrolled weight-tied input-injected network trained with backpropagation under different solvers.
Revised: We study test time convergence properties of an unrolled weight-tied input-injected network trained with backpropagation under different solvers.

Original: This network is highly path independent using either the Broyden solver or fixed point iterations, as indicated by its high AA scores (0.99) on both in- and out-of-distribution data.
Revised: This network is highly path independent using either the Broyden solver or fixed point iterations, as indicated by its high AA scores (0.99) on both in- and out-of-distribution data.

Original: We visualize the values of test-time residuals with fixed point iterations and Broyden’s method in Fig.
Revised: We visualize the values of test-time residuals with fixed point iterations and Broyden’s method in Fig.

Original: 5a.
Revised: 5a.

Original: Both these solvers converge to different limit cycles but still show good upward generalization.
Revised: Both these solvers converge to different limit cycles but still show good upward generalization.

Original:  Path Independence on a Per-Example Level
Revised:  Path Independence on a Per-Example Level

Original: The connection between path independence and prediction correctness also largely holds on a perinstance basis.
Revised: The connection between path independence and prediction correctness also largely holds on a perinstance basis.

Original: Using the prefix-sum networks trained with the mixed-initialization strategy (the most performant group of networks in our intervention experiments), we plotted the distribution of per-instance fixed point alignment scores, colored by whether the prediction on that instance was correct or not in Figure 5b.
Revised: Using the prefix-sum networks trained with the mixed-initialization strategy (the most performant group of networks in our intervention experiments), we plotted the distribution of per-instance fixed point alignment scores, colored by whether the prediction on that instance was correct or not in Figure 5b.

Original: This suggests that path independence can be used as a valuable sanity-check to determine whether a prediction is correct or not without the need for any label data, both in- and out-of-distribution.
Revised: This suggests that path independence can be used as a valuable sanity-check to determine whether a prediction is correct or not without the need for any label data, both in- and out-of-distribution.

Original: We provide a more in-depth per-instance analysis in the supplementary material.
Revised: We provide a more in-depth per-instance analysis in the supplementary material.

Original:  Related Work
Revised:  Related Work

Original: There is a long line of research on neural networks that can adapt their computational budget based on the complexity of the task they are learning to solve—akin to the intrinsic mechanism in humans to reason and solve problems.
Revised: There is a long line of research on neural networks that can adapt their computational budget based on the complexity of the task they are learning to solve—akin to the intrinsic mechanism in humans to reason and solve problems.

Original: Schmidhuber [2012] introduced self-delimiting neural networks which are a type of recurrent neural networks (RNNs) that adapt their compute based on the output of a special "halt" neuron.
Revised: Schmidhuber [2012] introduced self-delimiting neural networks which are a type of recurrent neural networks (RNNs) that adapt their compute based on the output of a special "halt" neuron.

Original: Adaptive computation time (ACT)
Revised: Adaptive computation time (ACT)

Original: [Graves, 2016a] also uses the output of a sigmoidal halting unit to determine the termination condition of an RNN, but it avoids long "thinking" time by explicitly penalizing it.
Revised: [Graves, 2016a] also uses the output of a sigmoidal halting unit to determine the termination condition of an RNN, but it avoids long "thinking" time by explicitly penalizing it.

Original: Subsequent works have successfully applied variants of ACT in image classification and object detection [Figurnov et al., 2017], visual reasoning [Eyzaguirre and
Revised: Subsequent works have successfully applied variants of ACT in image classification and object detection [Figurnov et al., 2017], visual reasoning [Eyzaguirre and

Original: Soto, 2020], Transformers [Vaswani et al., 2017] for language modelling [Dehghani et al., 2019,
Revised: Soto, 2020], Transformers [Vaswani et al., 2017] for language modelling [Dehghani et al., 2019,

Original: Elbayad et al., 2020a, Liu et al., 2021b], and recognizing textual entailment [Neumann et al., 2016].
Revised: Elbayad et al., 2020a, Liu et al., 2021b], and recognizing textual entailment [Neumann et al., 2016].

Original: PonderNet [Banino et al., 2021] reformulates the halting policy of ACT as a probabilistic model, and adds a regularization term in the loss objective to encourage exploration.
Revised: PonderNet [Banino et al., 2021] reformulates the halting policy of ACT as a probabilistic model, and adds a regularization term in the loss objective to encourage exploration.

Original: With these additions,
Revised: With these additions,

Original: PonderNet can extrapolate to more difficult examples on the parity task, first proposed by Graves [2016b]: in a vector with entries of 0, -1, and 1, output 1 for odd number of ones, and 0 otherwise.
Revised: PonderNet can extrapolate to more difficult examples on the parity task, first proposed by Graves [2016b]: in a vector with entries of 0, -1, and 1, output 1 for odd number of ones, and 0 otherwise.

Original: In this work, we do not optimize or penalize the network for the number of computational steps.
Revised: In this work, we do not optimize or penalize the network for the number of computational steps.

Original: Our main goal is to understand the underlying mechanism that results in scalable generalization of equilibrium models on harder problem instances.
Revised: Our main goal is to understand the underlying mechanism that results in scalable generalization of equilibrium models on harder problem instances.

Original: Our current work is closely related to previous work by Schwarzschild et al.
Revised: Our current work is closely related to previous work by Schwarzschild et al.

Original: [2021b] and [Bansal et al., 2022] that propose architectural choices and training mechanisms that enables weight tied networks to generalize on harder problem instances.
Revised: [2021b] and [Bansal et al., 2022] that propose architectural choices and training mechanisms that enable weight tied networks to generalize on harder problem instances.

Original: We relate these papers’ contributions to ours in Section 4.1.
Revised: We relate these papers’ contributions to ours in Section 4.1.

Original: Another family of models with the property of adaptive inference compute budget is early exit networks [Teerapittayanon et al., 2016, Laskaridis et al., 2021].
Revised: Another family of models with the property of adaptive inference compute budget is early exit networks [Teerapittayanon et al., 2016, Laskaridis et al., 2021].

Original: These networks have multiple additional “exit” prediction heads along their depth.
Revised: These networks have multiple additional “exit” prediction heads along their depth.

Original: At inference time, the result that satisfies an exit policy is selected as the prediction output.
Revised: At inference time, the result that satisfies an exit policy is selected as the prediction output.

Original: This approach of designing adaptive networks has been adapted both in natural language processing [Schwartz et al., 2020, Soldaini and Moschitti, 2020,
Revised: This approach of designing adaptive networks has been adapted both in natural language processing [Schwartz et al., 2020, Soldaini and Moschitti, 2020,

Original: Elbayad et al., 2020b, Zhou et al., 2020, Liu et al., 2020] and vision [Li et al., 2017, Wang et al., 2018,
Revised: Elbayad et al., 2020b, Zhou et al., 2020, Liu et al., 2020] and vision [Li et al., 2017, Wang et al., 2018,

Original: Xing et al., 2020, Kouris et al., 2021].
Revised: Xing et al., 2020, Kouris et al., 2021].

Original: Most of these architectures have complex sub-modules that are trained in multiple stages, and require complex exit policies.
Revised: Most of these architectures have complex sub-modules that are trained in multiple stages, and require complex exit policies.

Original: In contrast, equilibrium models have a simple architecture, and can use root solvers to efficiently solve for the fixed point at inference.
Revised: In contrast, equilibrium models have a simple architecture, and can use root solvers to efficiently solve for the fixed point at inference.

Original: More complex transformer-based language models like GPT-3 also struggle to generalize well on simple algorithmic tasks like addition [Brown et al., 2020].
Revised: More complex transformer-based language models like GPT-3 also struggle to generalize well on simple algorithmic tasks like addition [Brown et al., 2020].

Original: Recent work by Nye et al.
Revised: Recent work by Nye et al.

Original: [2021] shows that transformers can be trained to perform well on algorithmic tasks and generalize on OOD data by emitting the intermediate steps of an algorithm to a buffer called “scratchpad”.
Revised: [2021] shows that transformers can be trained to perform well on algorithmic tasks and generalize on OOD data by emitting the intermediate steps of an algorithm to a buffer called “scratchpad”.

Original: Using a scratchpad enables the model to revisit its errors and correct them.
Revised: Using a scratchpad enables the model to revisit its errors and correct them.

Original: Path independence is closely related to the concept of global stability and global convergence in control theory and optimization.
Revised: Path independence is closely related to the concept of global stability and global convergence in control theory and optimization.

Original: This concept is somewhat overloaded, as it sometimes requires convergence to a single point [Slotine et al., 1991], and sometimes implies the system is convergent everywhere, even if to different points [Wang et al., 2003, Sriperumbudur and Lanckriet, 2009].
Revised: This concept is somewhat overloaded, as it sometimes requires convergence to a single point [Slotine et al., 1991], and sometimes implies the system is convergent everywhere, even if to different points [Wang et al., 2003, Sriperumbudur and Lanckriet, 2009].

Original: We thus choose the term path independence to refer specifically to the fact that the system will converge to the same limiting behavior (whatever that might be) regardless of the initial state of the system.
Revised: We thus choose the term path independence to refer specifically to the fact that the system will converge to the same limiting behavior (whatever that might be) regardless of the initial state of the system.

Original:  Conclusion
Revised:  Conclusion

Original: Being able to attain better levels of performance using a larger inference-time compute budget is a feat that eludes most standard deep learning architectures.
Revised: Being able to attain better levels of performance using a larger inference-time compute budget is a feat that eludes most standard deep learning architectures.

Original: This is especially relevant for tasks that require upwards generalization , i.e., the ability to generalize from easy problem instances to hard ones.
Revised: This is especially relevant for tasks that require upwards generalization , i.e., the ability to generalize from easy problem instances to hard ones.

Original: We show that equilibrium models are capable of displaying upwards generalization by exploiting scalable test-time compute.
Revised: We show that equilibrium models are capable of displaying upwards generalization by exploiting scalable test-time compute.

Original: We link this to a phenomenon we call path independence : the tendency of an equilibrium network to converge to the same limiting behavior given an input, regardless of the initial conditions.
Revised: We link this to a phenomenon we call path independence : the tendency of an equilibrium network to converge to the same limiting behavior given an input, regardless of the initial conditions.

Original: We investigate this phenomenon through careful experiments and verify that path independent networks indeed generalize well on harder problem instances by exploiting more test time compute.
Revised: We investigate this phenomenon through careful experiments and verify that path independent networks indeed generalize well on harder problem instances by exploiting more test time compute.

Original: Moreover, interventions on training conditions that promote path independence also improve upwards generalization, while those that penalize it hurt this capability.
Revised: Moreover, interventions on training conditions that promote path independence also improve upwards generalization, while those that penalize it hurt this capability.

Original: Our findings suggest that path independent equilibrium models are a promising direction towards building general purpose learning systems whose test-time performance improves with more compute.
Revised: Our findings suggest that path independent equilibrium models are a promising direction towards building general purpose learning systems whose test-time performance improves with more compute.

Original:  Acknowledgements
Revised:  Acknowledgements

Original: AP is supported by a grant from the Bosch Center for Artificial Intelligence.JT acknowledges support from the Center on Long-Term Risk Fund. CA is supported by NSERC Canada Graduate School - Doctorate scholarship.
Revised: AP is supported by a grant from the Bosch Center for Artificial Intelligence.

Original: A Alternative Approaches for Quantifying Path Independence

Original: We review alternative methods for quantifying how path-independent a given equilibrium model is, and why the Fixed Point Alignment score is the most suitable one amongst them for our analyses.

Original: We want a path independence metric to satisfy three criteria: 1) Dimensionless: Metrics computed from different training runs should be directly comparable with each other.

Original: To ensure this, one has to make sure that one uses dimensionless metrics (i.e. doesn’t have units).

Original: 2) local and global path independence: The metric should test for path independence not just in a local neighborhood of fixed points, but over a sufficiently large portion of the initialization domain.

Original: 3) efficiency: It should be computationally tractable.

Original: With these in mind, we review three alternative ways of quantifying path independence.

Original: The summary of the analysis is described in Table :

Original: • Jacobian Norm: The Frobenius norm of the Jacobian of the output of an equilibrium model with respect to its fixed point initialization (i.e. || ∂ FIX( x , z 0 ) ∂ z|| 2 ) gives a robust measure of how locally sensitive the network is to initializations.

Original: If this quantity is very small (and approaching 0 with more root finding steps), one can conclude that the given equilibrium model is path-independent on the given input.

Original: The main issues with this approach are: 1) It’s extremely computationally intensive to compute (since both the input and the output of FIX are high-dimensional, neither forward nor backward differentiation can estimate this Jacobian efficiently) 2)

Original: It has units, meaning it isn’t comparable across different training runs, 3) It only measures local path independence.

Original: • Agreement with Adversarial Initializations: This method quantifies to what extent it is possible to directly optimize initializations such that they result in different fixed points.

Original: If the network is path independent, there shouldn’t exist an adversary that can find adversarial initializations.

Original: The exact procedure for how adversarial fixed points can be found in Section B.4.

Original: The downsides with this approach is that 1) it’s expensive - one has to solve an optimization problem for every problem instance, potentially using different optimizers and initializations, and 2) it only checks for local path independence.

Original: • Agreement with Swapped Initializations (Fixed Point Alignment): The main idea behind this approach is to initialize the forward pass of an equilibrium model with fixed points obtained from other problem instances, and checking if the network still displays the same asymptotic behaviour (i.e. finds the same fixed point).

Original: If one uses cosine similarity to measure consistency, then one recovers Fixed Point Alignment score proposed in Section 4.2.

Original: Another option is to use the average Euclidean distance to check for consistency.

Original: This is identical to computing the trace of the covariance matrix of the distance between the computed fixed points.

Original: This approach – especially the one that utilizes cosine similarity to quantify consistency – is appealing, as 1) because it’s unitless, it can be compared between training runs on the same task, 2) it checks for non-local convergence by sampling from a distribution of relevant fixed points 3)

Original: it’s very efficient, requiring only two forward passes.

Original: The version that utilizes Euclidean distance does have units, hence cannot be used for crossmodel comparisons 9 )

Original: To further make sure that this metric is correct, and does subsume with local path independence, we ran stress tests described in Section 4.2 to make sure this metric produces results consistent with the agreement with adversarial initializations method described above.

Original: Alternatives to Cosine Similarity in Measuring Directional Alignment While cosine similarity is a conceptually clean way of quantifying directional alignment, one can consider alternative kernels as well.

Original: Ideally, our main results should not depend on the specific implementation details of path independence metrics (as long as they satisfy the criteria we set out above).

Original: To check how sensitive our results are on the precise functional form of the AA score, we replaced it with three other kernels and re-assessed whether the results pointed to the same high-level takeaways.

Original: Concretely, we tried the Gaussian kernel ϕ g ( r ) =

Original: exp( − ( ϵr ) 2 ) , Laplacian kernel ϕ l ( r ) =

Original: exp( −| ϵr | ) and inverse multiquadratic kernel ϕ

Original: i ( r ) = 1 √ 1+( ϵr ) 2 where the input to the kernel function is the Euclidean distance of the normalized fixed points r = ||

Original: z 1 ||

Original: z

Original: || − z

Original: 2 || z

Original: 2 ||

Original: || .

Original: We used ϵ = 5000 , which gave a reasonable dynamic range in the resulting values.

Original: The results can be seen in Figure 6.

Original: The takeaways from the plots remain identical: path independence is strongly correlated with generalization, regardless of the specific details of how path independence is quantified (as long as it satisfies the criteria we set in Appendix A).

Original: Importance of using dimensionless metrics:

Original: Being unitless is necessary for a metric to be comparable across training runs, though obviously not sufficient.

Original: Consider two equilibrium models

Original: M1 and M2, where the fixed points computed by M2 have the same direction as those computed by M1, but have twice the Euclidean norm.

Original: The behaviour of this M2 is qualitatively the same as that of M1, but any metric that depends on the Euclidean metric (or Jacobian L2 norm, or any non-unitless metric) would report this network to be less path independent.

Original: Note that this is not a purely theoretical consideration: simply adding L2 regularization, or penalizing the magnitude of fixed points to encourage convergence will directly impact the scale of the fixed points (things that practitioners often use), thereby rendering non-unitless metrics unreliable.

Original: Path Independence Metrics

Original: Dimensionless

Original: Local &

Original: [Table]

Original: Coverage

Original: Efficient

Original: [Table]

Original: Norm ✗ ✗ ✗

Original: Using Adversarial Initializations ✓✓ ✓ ✗

Original: [Table]

Original: ✓✓ ✓ ✓✓

Original: B Experimental Details

Original: This section contains the additional information needed to reproduce our experimental findings.

Original: This section complements our code release.

Original: B.1 Details for Experiments on Prefix Sum

Original: Unless otherwise specified, all of the prefix-sum networks were trained under the following conditions:

Original: Training details: We used the Adam [Kingma and Ba, 2014] optimizer and trained forgradient steps with a batch size of 150 .

Original: We varied the initial learning rate on different runs ( 0 . 001 , 0 . 0001 and 0 . 00001 ).

Original: We applied a stepwise learning rate decay both with a factor of 0 .

Original: 5 halfway and three-quarters through the learning.

Original: We also used gradient clipping at L2 norm equals 1.

Original: Model selection: Since we didn’t observe overfitting in our experiments, we used the last model saved model checkpoint during evaluation.

Original: Training data: We used the length-32 training split of the prefix sum dataset described in

Original: Schwarzschild et al.

Original: [2021a].

Original: Architecture: We used a 1D-convolutional depth-wise recurrent ResNet architecture roughly analogous to that used by Bansal et al.

Original: [2022] in their prefix sum experiments.

Original: The main differences include: (1) We used a single convolutional readout layer, as well as single convolutional layer to right before the depth-wise recurrent section of the network.

Original: (2) Unless otherwise specified, we initialized the fixed point at the 0 vector, unlike Bansal et al.

Original: [2022] who initialized the fixed point vector directly with the input.

Original: (3) While Bansal et al.

Original: [2022] uses concatenation to do input injection, we directly added the input (or a projection of it) to the hidden activations of the residual blocks, the same way the Kolter et al.

Original: [2020] does in their convolutional architecture.

Original: (4) We ablated using Weight Normalization [Salimans and Kingma, 2016], as this is something that has been reported to stabilize training.

Original: Forward solvers: We used either fixed point iterations (of depth 6 or 32 ) or Anderson acceleration [Kolter et al., 2020] (with an iteration count of 6 or 32 and a memory of 3 ) depending on the experimental purpose of the experiment.

Original: Backwards pass: We used either the standard backpropagation algorithm, the IFT gradients to train the networks or truncated backpropagation, depending on the purpose of the experiment.

Original: We used Anderson Acceleration to estimate the IFT gradients with an iteration budget of 10 and memory of 3 . We found that rescaling the Jacobian term that appears in Equation 1 with a constant less thanstabilized training with the IFT gradients. We used a rescaling factor of 0 . 8 .
Revised: JT acknowledges support from the Center on Long-Term Risk Fund.

Original: While this is relatively non-standard practice, we confirmed that using this stabilization strategy doesn’t yield qualitatively different results than those reported in the paper.

Original: When we used truncated backpropagation, we dropped the gradients halfwah through the network (i.e. if the forward pass consisted of 10 iterations of the recurrent cell, we only backpropagated through the last 5 iterations).

Original: Loss function: We used the cross entropy between the network predictions and the labels as the loss function.

Original: the prefix-sum task is a sequence-to-sequence problem, and each forward pass consists of multiple predictions.

Original: We used the average of all of the cross entropy values computed in a single forward pass.

Original: Compute resources: All of the prefix sum networks were trained using Intel Xeon Silver CPUS and NVIDIA GPUs (both P100, T4V1 and T4V2 models were used in the experiments).

Original: Training one prefix-sum network took roughly 1 to 2 hours depending on how deep the forward pass is, and whether one is using IFT or backpropagation gradients.

Original: B.1.1 Reproducing Table 1 (Stress-Testing the AA Score)

Original: The non-path independent network used in this table is a 15 -layer weight-tied input-injected network trained with backpropagation gradients.

Original: We set the maximum number of iterations for fixed point computation to 11 as we found that this choice gives the highest test accuracy on 64 bit prefix sum strings.

Original: While selecting the optimal value for maximum iterations, we increased the value up toiterations.

Original: The path independent network is a 32 -layer weight-tied input inejcted network trained with backpropagation gradients.

Original: The maximum iterations for this network are set to 500 .

Original: We used a batch size of 1 and set the maximum number of L-BFGS [Liu and Nocedal, 1989] updates per batch tofor both the networks.

Original: We implemented this code in PyTorch and use these values of hyperparameters: lr=1 , tolerance_grad =1e-7, tolerance_change =1e-9, line_search_fn="strong_wolfe" .

Original: We report AA scores and accuracy computed on 500 examples.

Original: We provide pseudocode in Algorithm

Original: Box

Original: B.2 Details for Experiments on Mazes

Original: Unless otherwise specified, all the networks for mazes were trained under the following conditions:

Original: Training details: We used the Adam [Kingma and Ba, 2014] optimizer with an initial learning rate of 0 .

Original: 001 and batch size of 50 .

Original: We use stepwise learning rate decay with a factor of 0 .

Original: 1 applied at epoch number 100 .

Original: We trained all weight-tied input-injected networks for a total of 150 epochs.

Original: We observed that implicitly trained equilibrium models needed a longer training time for their validation accuracy to improve.

Original: All implicitly trained equilibrium models were pretrained for 150 epochs and trained with IFT for another 100 epochs.

Original: Model selection:

Original: We used the checkpoint corresponding to the model

Original: with the best validation accuracy during evaluation.

Original: If multiple checkpoints had same validation accuracy, we selected the most recent checkpoint.

Original: Training data: We trained on the training split of 9 × 9 mazes from the dataset described in

Original: Schwarzschild et al. [2021a]. Architecture: We used a 2D-convolutional depth-wise recurrent ResNet architecture analogous to that used by Bansal et al. [2022] in their mazes experiments.
Revised: CA is supported by NSERC Canada Graduate School - Doctorate scholarship.

Original: While pretraining implicitly trained equilibrium models, we do rollouts for 32 -iterations.

Original: For equilibrium models trained with phantom gradient in Table 2, we use a damping factor of 0 .

Original: 75 and iterate for 2 iterations.

Original: Shallow weighttied input-injected networks were trained with 8 -layers while the deeper weight-tied input-injected networks used 32 layers.

Original: Both the feedforward networks are 32 -layer deep non-weight tied networks.

Original: We ablate on input-injection for these networks.

Original: We used the original code released by [Bansal et al., 2022] to replicate the results for weight-tied input-injected networks trained with progressive training.

Original: In addition we ablate on weight normalization [Salimans and Kingma, 2016] for all the architectures— both unrolled and implicitly trained equilibrium models.

Original: We found that weight normalization helps stabilize the training, especially implicitly trained models.

Original: We found that using normalization layers (like group norm [Wu and He, 2018]) hurts generalization in mazes unlike prefix sum.

Original: Therefore, none of our architectures for mazes use any normalization layers.

Original: Forward solvers: We used either fixed point iterations (of depth 8 or 32 ) or Broyden’s method [Broyden, 1965] (with an iteration count of 40 ) depending on the experimental purpose of the experiment.

Original: Backward pass: We used the standard backpropagation algorithm, the IFT gradients, truncated backpropagation or phantom gradients to train the networks, depending on the purpose of the experiment.

Original: We used Broyden’s method to estimate the IFT gradients with an iteration budget of 40 .

Original: For experiments with truncated backpropagation, we follow the exact setting as specified in [Bansal et al., 2022].

Original: Loss function: We used the cross entropy between the network predictions and the labels as the loss function.

Original: Compute resources: All networks for mazes were trained using Intel Xeon Silver CPUs and NVIDIA

Original: GPUs (1080Ti, 2080Ti and A6000 models were used in the experiments).

Original: Training a weight-tied network took roughly around 12 - 18 hours depending on how deep the forward pass is, whereas the implicitly trained models took around 2 days of training time on a single GPU.

Original: B.3 Reproducing Figure

Original: Both models were trained using 32 fixed point iterations and backpropagation.

Original: The path independent model is input injected, the path dependent model is not input injected.

Original: We run the network forfixed point iterations on five different validation examples of length 32 to obtain different hidden layer initializations.

Original: We then run the models for 128 steps using these initializations to obtain five different trajectories.

Original: For the input injected network, we use the same input injection for all trajectories, obtained from a different example of length 32.

Original: We project the resulting hidden states onto two random orthogonal directions in the hidden layer state space.

Original: B.4 Experimental Details for Section 3 (Upwards Generalization with Equilibrium Models) B.4.1 Reproducing Figure 2a (OOD Generalization with Mazes)

Original: The training details for all the weight-tied input-injected networks used to produce this plot can be found in Section B.2.

Original: The feedforward network used in this plot is a 32 -layer non-weight tied network without any input injection that was trained with regular backpropagation.

Original: The maximum number of test-time iterations (both unroll and Broyden solver) for mazes of size 9 , 13 , 25 , 31 and 59 were 50 , 100 , 400 , 600 , and 800 respectively.

Original: The test accuracy for mazes of a given dimension was computed on the entire split of 10 , 000 mazes as provided in the original dataset [Schwarzschild et al., 2021b].

Original: B.4.2 Reproducing Figure 2b (

Original: PI networks

Original: get

Original: better

Original: with

Original: more

Original: iterations)

Original: The training details for the 12 networks used to produce this plot can be found in Section B.1.

Original: We varied (1) the depth ( 6 vs. 32 ), whether weight normalization[Salimans and Kingma, 2016] was used or not, and the initial learning rate ( 0 . 001 , 0 . 0001 or 0 . 00001 ).

Original: We evaluated each model on 128 examples of length 16, 32, 64, and 128 each, using fixed point iterations and measuring the average number of correct bits at each iteration.

Original: For each model and iteration step, we plot average accuracies over all splits and examples.

Original: The Fixed Point Alignment

Original: Scores for each model are calculated as in Figure 4a (see Appendix B.5.1).

Original: B.4.3 Reproducing Table 1 (Stress-Testing the AA Score)

Original: The non-path independent network used in this table is a 8 -layer weight-tied input-injected network trained with backpropagation gradients.

Original: We set the maximum number of iterations for fixed point computation to 12 as we found that this choice gives the highest test accuracy on 13 × 13 mazes.

Original: While selecting the optimal value for maximum iterations, we increased the value up to 60 iterations.

Original: The path independent network is a 32 -layer weight-tied input injected network trained with backpropagation gradients.

Original: The maximum iterations for this network are set to 500 .

Original: We used a batch size ofand set the maximum number of L-BFGS [Liu and Nocedal, 1989] updates per batch to 50 for both the networks.

Original: We implemented this code in PyTorch and use these values of hyperparameters: lr=1 , tolerance_grad =1e-7, tolerance_change =1e-9, line_search_fn="strong_wolfe" .

Original: We report AA scores and accuracy computed on 500 examples.

Original: We provide pseudocode in Algorithm Box

Original: B.5 Experimental Details for Section 5 (Path Independence Correlates with Upward Generalization) B.5.1 Reproducing Figure 4a (Path Independence - Accuracy Correlation on Prefix Sum)

Original: The training details for the 12 networks used to produce this plot can be found in Section B.1.

Original: We trained 12 networks under each category reported in the plot, where (1) the depth ( 6 vs. 32 ), whether weight normalization[Salimans and Kingma, 2016] was used or not, and the initial learning rate ( 0 . 001 , 0 . 0001 or 0 . 00001 ) were varied.

Original: Each value on the plots were computed using 300 examples from the validation splits corresponding to lengths 16 , 32 , 64 , 128 and 256 (i.e. 5 ∗ 300 = 1500 samples per point).

Original: The average cosine similarities involved in computing the Fixed Point Alignment scores were computed using 10 different initialization per example.

Original: For the test-time forward pass, we used 512 iterations of Anderson

Original: Acceleration [Kolter et al., 2020] (note that this is using a more sophisticated solver using 16 times more forward iterations than the deepest of the trained networks).

Original: The trend lines on the plots were computed using the line that minimizes the least squares error.

Original: B.5.2 Reproducing Figure 4b (Path Independence - Accuracy Correlation on Mazes)

Original: The architectural and training details for all the networks in this plot can be found in Section B.2.

Original: Each value in this plot was computed using 500 examples from the validation splot corresponding to mazes of size 13 and 25.

Original: The average cosine similarities involved in computing the Fixed Point

Original: Alignment scores were computed using 10 different initialization per example.

Original: Unrolled + backprop networks use 100 test-time iterations for the forward pass for mazes of size 13 × 13 , and 400 testtime iterations for mazes of size 25 × 25 .

Original: We set the maximum number of permissible iterations for the Broyden solver [Broyden, 1965], and for non-weight tied feedforward network to the same values.

Original: The shallow unrolled + backprop network was trained using 8 layers.

Original: Empirically, we obtained the highest test accuracies on 11 th and 14 th iterations for mazes of size 13 × 13 and 25 × 25 respectively, for networks that use weight normalization.

Original: For networks that do not use weight normalization, the corresponding numbers are reported on 12 th and 18 th iteration.

Original: The trend lines on the plots were computed using the line that minimizes the least squares error.

Original: B.6 Experimental Details for Section 6 (Experimental Manipulations of Path Independence) B.6.1 Reproducing Figure 3 (Manipulations of PI)

Original: The training details for the 12 networks that formed the baseline for the interventions are described in

Original: Section B.1.

Original: We trained 12 networks under each category reported in the plot, where (1) the depth (vs. 32 ), whether weight normalization[Salimans and Kingma, 2016] was used or not, and the initial learning rate ( 0 . 001 , 0 . 0001 or 0 . 00001 ) were varied.

Original: For each intervention, we kept all of the aforementioned training conditions except for the experimental manipulations that are intended to promote or hurt path independence.

Original: • Mixed initialization: During each training forward pass, each sample was assigned with either zero initialization (i

Original: e. the fixed point was initialized with the 0 vector) or standard normal distribution (i.e. ...) using a Bernoulli random variable of probability 0 .

Original: 5 (i.e. the examples that were run with zero vs. normal initializations were roughly half-half.

Original: • Randomized depth: Randomized training was done by first specifying the minimum and maximum possible depths, then uniformly sampling a depth from that range during each forward pass.

Original: We used a range of 3 to 63 forward iterations for training 6 of the networks, and used a range of 32 to 64 for training the other 6 . • Alignment penalty:

Original: The alignment penalty, which was added to the original training loss in form of a regularized was computed as follows: (1) For each example in the batch, compute the fixed points k times using normal initialization, (2) Compute the dot product between each pair of fixed points (which yields k 2 − k dot products, excluding each fixed point’s dot product with itself), return the sum of the dot products, normalized by k 2 − k .

Original: In our experiments, we picked k = 3 .

Original: B.7 Experimental Details for Section 7 (Disambiguating Convergence and Path Independence) B.7.1 Reproducing Table 2 (Training Convergence vs. Accuracy)

Original: The architectural and training details for equilibrium models trained with phantom gradients and with

Original: IFT have been provided in Section B.2.

Original: Both the networks use Broyden solver for the forward pass.

Original: We set the number of maximum permissible iterations to 40 for in-distribution examples ( i.e., 9 ×mazes) and to 400 for 25 × 25 mazes.

Original: All the values reported in this table have been computed on 500 examples from the respective splits for in-distribution and out-of-distribution data.

Original: B.7.2 Reproducing Figure 5a (Test Time Convergence vs. Path-Independence)

Original: The architectural and training details of weight-tied input-injected networks used in this figure have been provided in Section B.2.

Original: The reported numbers have been averaged over 10 different runs.

Original: Each run uses a batch size of 100.

Original: The maximum number of forward iterations is set to 200 for both fixedpoint and Broyden solvers

Original: . B.7.3 Reproducing Figure 5b (Per Instance Path Independence vs. Accuracy)

Original: To compute the histogram, we took the 12 networks trained with the training condition that yielded the strongest OOD accuracy (mixed initialization) and computed both the accuracy and AA scores on 300 examples for the prefix-sum dataset validation splits of lengths 16 , 32 , 64 , 128 and 256 .

Original: We grouped each AA score based on whether the prediction it’s paired with was correct or not, then produced the histogram.

Original: C Limitations

Original: Currently, the experimental results are obtained only using two tasks: prefix-sum and mazes.

Original: Testing the validity of our findings on a more diverse set of tasks would increase our confidence in the fact that the phenomenon reported in this paper is universal.

Original: Moreover, it is not clear from our analysis which tasks would benefit from more test-time compute.

Original: While the necessity of additional test-time compute is obvious for prefix sum and mazes, it’s less clear for perception-like tasks like image classification.

Original: Furthermore, we believe that a more systematic study of how the choice of solver (i.e.

Original: Anderson (first order) vs. Broyden (second order) solvers) impacts the empirical AA scores would strengthen our findings.

Original: We observe, for example, that networks trained with Broyden struggle to find fixed points at test time when run with fixed-point iterations.

Original: Expanding on these findings would improve our analysis and let us make more concrete suggestions to practitioners aiming to use the path-independence related concepts from this paper.

Original: D Path Independence vs. Accuracy Plots Different Difficulty Levels

Original: The path-independence vs. accuracy plots in Figure 4b uses averaged values over different problem difficulties.

Original: In Figure 7, we show how these correlations look like at different lengths on the prefix sum task.

Original: The same figure also shows how the accuracy differs if one does inference using the training time forward pass budget.

Original: The main takeaways from this plot are:

Original: • Out-of-distribution problem lengths (right top of the Figure) do indeed require larger inference depths - all models perform very poorly when inference is run with training-time forward pass budget.

Original: • The correlation between path independence and accuracy - as the definition implies - is more apparent in the very large forward pass limit.

Original: In this regime, the correlation persists for all lengths.

Original: • While networks with low PI values using the training-time budget occasionally perform worse when the test-time compute is increased, that doesn’t happen with path-independent networks.

Original: E Intervention Results on Different Difficulty Levels

Original: Mirroring Section D, we provide how the accuracy correlates with path independence on 1) different problem difficulties ( 32 is in-distribution) and different forward pass budgets in Figure 8.

Original: We especially emphasize the mixed initialization results (mixed initialization is one of the interventions that promotes path independence).

Original: Note that some of these networks (half of which are as shallow as 6 layers) do poorly in and out-of-distribution when the forward pass uses the training forward pass budget.

Original: Unsurprisingly, the AA scores for these networks are low in this condition.

Original: However, when additional test-time compute is provided, these networks reach perfect in-distribution accuracy, and near-perfect OOD accuracy (especially on length = 64 split).

Original: This demonstrates that training time convergence is not required for path independence.

Original: We find it surprising that with the help of a pathindependence-promoting regularization procedure, it’s possible to train very shallow (i.e. less thanlayers)

Original: networks such that they can exploit additional test-time compute to achieve very high accuracy.

Original: F Per-Instance Path Independence Analyses - Convergence vs. Path Independence

Original: In addition to the analysis in Section 8, we also analyzed how test-time convergence correlates with path independence on a per-instance basis.

Original: The per-instance AA score vs. convergence (measured in terms of the L2 distance between the last two root solver iterates) can be seen in Figure 9.

Original: We used the 12 networks trained using the mixed initialization strategy (since these networks yielded the best in and OOD performance), and used 300 samples from the length splits 16 , 32 , 64 , 128 and 256 .

Original: This analysis reaffirms our past observation that test-time convergence is not needed for path independence.

Original: The plot contains data from three different regimes:

Original: • Full convergence to a fixed point: These datapoints, observed on the right bottom part of Figure 9, corresponds to samples where the solver nearly converges to a fixed point.

Original: This is almost always associated with high FPA scores, and good per-instance accuracy.

Original: • Limit cycles: These datapoints, corresponding to the right-top part of Figure 9, correspond to cases where the solver doesn’t converge to a fixed points, but enters an orbit around it.

Original: AA scores, as well as the per-instance accuracy values remain high in this regime as well, indicating that test-time convergence (to a fixed point) is not necessary for path independence, and good accuracy.

Original: Note that the boundary between the full-convergence regime and the limit cycle regime is gradual.

Original: Divergence:

Original: On a number of the samples, the solver diverges.

Original: These are associated with high residuals, low AA scores and low per-instance accuracies.

Original: This shows that the main source of per-instance lack of path independence on networks that are otherwise overwhelmingly path independent (on other samples) is almost always solver divergence.

Original: We used the 12 networks trained using the mixed initialization strategy (since these networks yielded the best in and OOD performance), and used 300 samples from the length splits

Original: While good convergence is almost always associated with high AA score values, the converse is not necessarily true: there are many samples on which convergence is poor, yet the AA score is high

Original: G Test Time Convergence and Path independence

Original: We provide a per-instance analysis of test-time convergence in Figure 10.

Original: We display plots for perinstance residuals i.e. ∥ f ( x, z ) − z ∥ 2 for Broyden’s method and naive fixed point iterations over solver steps.

Original: In addition, we also plot L2 norm between the fixed points of these solvers at every step.

Original: We provide plots for both in-distribution data i.e. 9 × 9 mazes and on more difficult problem instances i.e. 13 × 13 and 25 × 25 mazes.

Original: Across all the mazes, we observe that there are problem instances where both the solvers converge to the same limiting behavior (as indicated by low values of residuals and L2 norms).

Original: However, there are a considerable number of points where naive fixed point iterations converge to a limit cycle but still output correct predictions.

Original: This behavior can be seen on both in-distribution data and on harder problem instances.

Original: We find that the absolute residuals between the points in the limit cycles are a small percentage of the Euclidean norms of the points: 0.9% for 9x9 mazes, 0.49% for 13x13 mazes and 1.5% for 25x25 mazes which indicates that these limit cycles are localized.

Original: These examples of problem instances where both the solvers converge to different limiting behavior reaffirm our observation that convergence to the same fixed point at test-time is not a necessary condition for path independence.

Original: H Results on the Blurry MNIST Task

Original: We tested our path-independence hypothesis on task we call BlurryMNIST [Liang et al.].

Original: This task involves training an image classifier on MNIST digits corrupted with small degrees of Gaussian blur, and testing the performance on significantly more corrupted ones.

Original: The blur corruption is implemented by convolving each image with Gaussian filters of differing standard deviations.

Original: We used standard

Original: [Figure]

Original: deviations from 2 to 5 .

Original: 5 (in increments of 0 . 5 ) to generate each split.

Original: This dataset allows for testing in and out-of-distribution performance by way of training on a subset of the splits (i.e. the four lowest blur splits) and testing performance on all splits.

Original: We trained a number of fully connected equilibrium models on the BlurryMNIST task and inspected whether path independence still correlates with accuracy.

Original: The results can be found in Figure 11.

Original: The correlation we reported in the main body of the paper also holds in the BlurryMNIST dataset.

Original: Note that the in-distribution error rates of the trained models vary between 1 and 5 percent.

Original: I Results on Matrix Inversion Task

Original: We also tested the connection between path independence and generalization on the Matrix Inversion task proposed by Du et al.

Original: This task is concerned with learning to invert 20x20 matrices.

Original: Success is defined by how well the trained model works on matrices with worse condition numbers than those observed during training.

Original: Note that this task is qualitatively very different from all the others we considered before.

Original: We have summarized the results in Figure 12.

Original: Using a fully-connected ResNet block as the equilibrium model cell of width 512 , we trained a number of equilibrium models where we varied 1) the forward pass (fixed point iterations vs. solver) 2) the backwards pass (backprop gradients or implicit gradients) 3) learning rate ( 0 . 001 , 0 . 0001 and 0 . 00001 ), 4) learning rate schedule (step decays of magnitude 0 .

Original: 5 at different points during training) and 5) whether layer normalization [Ba et al., 2016] was used or not.

Original: The results can be seen in

Original: Figure 12.

Original: We see a similar pattern that we saw on earlier tasks: lack of path independence correlates with poor generalization.

Original: Note that our best model matches the performance of the energy based model approach proposed by Du et al.

Original: [2022] in the matrix inversion task and significantly beats their baselines.

Original: J Results on the Edge Copy Task

Original: We also test our path independence hypothesis on tasks that take in a graph as an input.

Original: We consider the following the edge copy task proposed by Du et al.

Original: [2022], where the goal is to learn to simply output the input edge features, in a way that generalizes to larger graph sizes.

Original: We used an equilibrium model cell that’s compatible with graph tasks in the Edge Copy experiments, whose structure is shown in Figure 14.

Original: This cell is especially suited for edge regression tasks, since each application of the cell refines the edge features.

Original: We trained a number of equilibrium models where we varied 1) the forward pass (fixed point iterations vs. solver) 2) the backwards pass (backprop gradients or implicit gradients) 3) learning rate ( 0 . 0001 , , 0 . 000333 and 0 . 0001 ) and 4) whether layer normalization [Ba et al., 2016] was used or not.

Original: Each hyperparameter configuration was run twice with different seeds.

Original: The results can be seen in Figure 13.

Original: We see a similar pattern that we saw on earlier tasks: lack of path independence correlates with poor generalization.

Original: Note that our best equilibrium model outperforms the the energy based model approach proposed by Du et al.

Original: [2022] in this task.