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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

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 ofdepth-wise recurrent networks.
Revised: Recent efforts have shown promising results in this direction by making use ofdepth-wise recurrent networks.

Original: In this work, we reproduce the performance of theprior art using a broader class of architectures called equilibrium models, and findthat stronger generalization performance on harder examples (which require moreiterations of inference to get correct) strongly correlates with the path independence of the system—its ability to converge to the same attractor (or limit cycle)regardless of initialization, given enough computation.
Revised: In this work, we reproduce the performance of theprior art using a broader class of architectures called equilibrium models, and findthat stronger generalization performance on harder examples (which require moreiterations of inference to get correct) strongly correlates with the path independence of the system—its ability to converge to the same attractor (or limit cycle)regardless of initialization, given enough computation.

Original: Experimental interventionsmade to promote path independence result in improved generalization on harder(and thus more compute-hungry) problem instances, while those that penalize itdegrade this ability.
Revised: Experimental interventionsmade to promote path independence result in improved generalization on harder(and thus more compute-hungry) problem instances, while those that penalize itdegrade this ability.

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

Original: Thus, considering equilibrium models and path independence jointly leads to avaluable new viewpoint under which we can study the generalization performanceof these networks on hard problem instances.
Revised: Thus,considering equilibrium models and path independence jointly leads to a valuablenew viewpoint under which we can study the generalization performance of thesenetworks on hard problem instances.

Original:  Introduction
Revised:  Introduction

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

Original: One particularly importanttype of out-of-distribution  generalization is upwards generalization , or the ability to generalize tomore difficult problem instances than those encountered at training time [Selsam et al., 2018, Bansalet al., 2022, Schwarzschild et al., 2021b, Nye et al., 2021].
Revised: One particularly importanttype of out-of-distribution (OOD) generalization is upwards generalization , or the ability to generalizeto 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 difficultinstances will require a larger amount of test-time computation, so a natural question arises: howcan we design neural net architectures which can reliably exploit additional test-time computation toachieve better accuracy?
Revised: Often, good performance onmore difficult instances will require a larger amount of test-time computation, so a natural questionarises: how can we design neural net architectures which can reliably exploit additional test-timecomputation to achieve better accuracy?

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 globally stable ; that is, whetherthe learned model converges to the same limit point (or limit cycle), regardless of the initializationof the system or the number of times the components of the model are applied.
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 globally stable ; that is, whetherthe learned model converges to the same limit point (or limit cycle), regardless of the initialization ofthe system or the number of times the components of the model are applied.

Original: For example, asimple integrator x t +1 = x t + 1 is clearly not globally stable, as it the final state depends on theinitial state and the number of iterations run; conversely, the system x t +1 =
Revised: For example, a simpleintegrator x t +1 = x t + 1 is clearly not globally stable, as it the final state depends on the initial stateand the number of iterations run; conversely, the system x t +1 =

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

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

Original: However, the term “global stability” is slightly overloaded, as it sometimesrefers to the actualconvergence, sometimes requires convergence to a single point, and sometimes just implies the systemis convergent everywhere, even if to different points.
Revised: However,the term “global stability” is slightly overloaded, as it sometimes requires convergence to a singlepoint [Slotine et al., 1991], and sometimes just implies the system is convergent everywhere, even ifto 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 (whateverthat 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 limitingbehavior (whatever that might be) regardless of the initial state of the system.

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

Original: For instance, gradient descent applied to a convex objective ispath 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 ispath 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 ispath dependent, and extending the simulation won’t yield more accurate predictions of a given day’sweather.
Revised: Conversely, a weather simulation ispath dependent, and extending the simulation won’t yield more accurate predictions of a given day’sweather.

Original: Based on this intuition, we hypothesize that path independence of a learned model is a keydeterminer of whether it can take advantage of an increased test-time iteration budget.
Revised: Based on this intuition, we hypothesize that path independence of a learned model is a keydeterminer of whether it can take advantage of an increased test-time iteration budget.

Original: In principle, neural nets have the capacity to represent many different approaches to solving a problem;for instance, even vanilla RNNs are Turing-complete [Siegelmann and Sontag, 1995].
Revised: In principle, neural nets have the capacity to represent many different approaches to solving a problem;for instance, even vanilla RNNs are Turing-complete [Siegelmann and Sontag, 1995].

Original: However, anyparticular combination of architecture and training regime may exhibit a strong inductive bias towardsa particular approach.
Revised: However, anyparticular combination of architecture and training regime may exhibit a strong inductive bias towardsa particular approach.

Original: For instance, deep equilibrium models [Bai et al., 2019] use implicit weightupdates which depend only on the final hidden state (not the path), and hence ought to encouragepath independence, despite the ability to represent path-dependent functions.
Revised: For instance, deep equilibrium models [Bai et al., 2019] use implicit weightupdates which depend only on the final hidden state (not the path), and hence ought to encouragepath independence, despite the ability to represent path-dependent functions.

Original: We hypothesize that thedegree of path independence has a fundamental impact on generalization: 
Revised: We hypothesize that thedegree of path independence has a fundamental impact on generalization: Path Independence Hypothesis

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

Revised: exploit more

Revised: test

Revised: time

Revised: iterations

Revised: to

Revised: achieve higher accuracy

Revised: , compared to

Revised: those which fit

Revised: the

Original:  We first introduce a metric for measuring the path independence of a learned model, the Fixed Point
Revised: training dis tribution with a path-dependent function We first introduce a metric for measuring the path independence of a learned model, the Fixed Point

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

Original: On two algorithmic generalization benchmarks introduced by Schwarzschildet al.
Revised: On two algorithmic generalization benchmarks introduced by Schwarzschildet al.

Original: [2021a,b], we show that a model’s FPA score is strongly correlated with its accuracy when farmore iterations are used at test time than at training time.
Revised: [2021a,b], we show that a model’s FPA score is strongly correlated with its accuracy when farmore iterations are used at test time than at training time.

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

Original: We find that previously introduced architectural features which improve generalization, includinginput injection and implicit updates which are explicitly captured by equilibrium fixed point models,improve both the path independence and the accuracy.
Revised: We find that previously introduced architectural features which improve generalization, includinginput injection and implicit updates which are explicitly captured by equilibrium fixed point models,improve both the path independence and the accuracy.

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

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

Original:  Background
Revised:  Background

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

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

Original: Concretely, letting x ∈ R n x and f w :71 R n x × n z (cid:55)→ R n z stand for an input and the equilibrium model function (or “cell") parametrized byw ∈ R n w respectively, equilibrium models aim to solve for the fixed point z ∗ ∈ R n z that satisfiesz ∗ =
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 “cell") parametrized by w ∈ R n w respectively, equilibrium models aim to solve for the fixed pointz ∗ ∈ 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 transformerblock [Vaswani et al., 2017].
Revised: The cell f w might represent anything from a fully connectedlayer to a transformer block [Vaswani et al., 2017].

Original: As f w is applied repeatedly, these networks are weight-tied in nature. Further, the input x (or, its projection) is injected into every layer of this weight-tied network.
Revised: We emphasize that f w directly depends on theinput x – following existing literature, we refer to this as input injection .

Original: Werefer to this as input-injection.
Revised: The outcome of the fixedpoint finding process might depend on the initial fixed point guess. 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 fixedpoint z 0 to an output that satisfies the fixed point equation FIX f w ( x , z 0 ) : = f w ( x , FIX f w ( x , z 0 )) . 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, whichrecursively applies the function f w on the internal representations z (i.e. z t +1 = f w ( x , z t ) ) until itconverges to a fixed point: z ∗ = f w ( x , z ∗ ) .
Revised: The most straightforward approach to solve for fixed-points is the fixed point iteration method, whichrecursively applies the function f w on the internal representations z (i.e. z t +1 = f w ( x , z t ) ) until itconverges to a fixed point: z ∗ = f w ( x , z ∗ ) .

Original: As solving for fixed points exactly is expensive, fixedpoint iterations are often terminated after a fixed number of steps or when the norm of the differencebetween subsequent iterates falls below a pre-selected threshold.
Revised: As solving for fixed points exactly is expensive, fixedpoint iterations are often terminated after a fixed number of steps or when the norm of the differencebetween subsequent iterates falls below a pre-selected threshold.

Revised: The model weights can be updatedusing 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 modelf ∞ 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 modelf ∞ w ( x , z ∗ ) =

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

Original: Deep Equilibrium (DEQ)
Revised: Deep Equilibrium (DEQ)

Original: Models Deep equilibrium models [Bai et al., 2019] solve for fixed pointsexplicitly using black-box root finders, such as Broyden’s method [Broyden, 1965] or Andersonacceleration [Anderson, 1965].
Revised: Models Deep equilibrium models [Bai et al., 2019] solve for fixed pointsexplicitly using black-box root finders, such as Broyden’s method [Broyden, 1965] or Andersonacceleration [Anderson, 1965].

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

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

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

Original: [Equation]
Revised: [Equation]

Original: Computing the inverse of a Jacobian can become computationally expensive.
Revised: Computing the inverse of a Jacobian can become computationally expensive.

Original: Recent works [Genget al., 2021a,b, Fung et al., 2021] have shown that the inverse-Jacobian term in Eq.
Revised: Recent works [Genget al., 2021a, Fung et al., 2021] have shown that the inverse-Jacobian term in Eq.

Original: (1) can be replacedwith an identity matrix without affecting the final performance of DEQs.
Revised: (1) can be replacedwith an identity matrix i.e. Jacobian-free or an approximate inverse-Jacobian [Geng et al., 2021b]without affecting the final performance of DEQs.

Original: This approximation makesthe backward pass of DEQs inexpensive and lightweight.
Revised: This approximation makes the backward passof DEQs inexpensive and lightweight.

Original: DEQ models have been shown to achievestate-of-the-art performance on various tasks including language modelling [Bai et al., 2019], imagerecognition, semantic segmentation [Bai et al., 2020], object detection [Wang et al., 2020], and graphmodeling [Gu et al., 2020, Liu et al., 2021a].
Revised: DEQ models have been shown to achieve state-of-the-artperformance 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: DEQs vs. Depthwise Recurrent Models Both DEQs and depthwise recurrent (i.e. weight-tied,fixed-depth) networks leverage weight-tying i.e., they apply the same transformation at each layer,f
Revised: DEQs vs. Depthwise Recurrent Models Both DEQs and 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: However, the key difference between the two is that weight-tied fixed depth networksrequire backpropagation through an explicit stack of layers but DEQs directly solve for the infinitedepth limit of the forward pass and use implicit differentiation to compute gradients.
Revised: However, the key difference between the two is that weight-tied fixed depth networksrequire backpropagation through an explicit stack of layers but DEQs directly solve for fixed pointsusing black-box solvers during the forward pass and use implicit differentiation to compute gradients.

Original: Upwards Generalization with Equilibrium Models
Revised: Convergence In order to guarantee convergence to a unique fixed point, it suffices for the cell of theequilibrium model to be contractive over its input domain 2 (i.e. the singular values of its Jacobian alllie below 1 ). Previous work have leveraged Lipschitz constrained cells to ensure contractivity[Revayet al., 2020]. Other approaches for ensuring global convergence exist: the monotone DEQ architectureguarantees global convergence by utilizing a DEQ parametrization that bears similarities to solutionsto a particular form of monotone operator splitting problem [Winston and Kolter, 2020]. Unrestricted

Revised: DEQs aren’t constrained enough to guarantee convergence: they can easily express globally divergentvector fields that prohibit the existence of fixed points.

Revised: It is, therefore, interesting that they canand often learn path independent solutions.

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

Revised: [Figure]

Revised:  Upwards Generalization with Equilibrium Models

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

Original: We focus in particular on two algorithmic generalizationbenchmarks introduced by Schwarzschild et al.
Revised: We focus in particular on two algorithmic generalizationbenchmarks introduced by Schwarzschild et al.

Original: [2021a,b]: prefix sum and mazes .
Revised: [2021a,b]: prefix sum and mazes .

Original: In this section, wefirst replicate their findings on these problems, and then extend them to a wider variety of architecturesconsidered in this paper, most notably deep equilibrium models.
Revised: In this section, wefirst replicate their findings on these problems, and then extend them to a wider variety of architecturesconsidered in this paper, most notably deep equilibrium models.

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

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

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

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

Original: The output is the optimal path in the maze that connectsthese two points without passing through the walls.
Revised: The output is the optimal path in the maze that connectsthese 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 areprovided in the Supplementary Material.
Revised: Instances of each of these problems areprovided in the Supplementary Material.

Original: Strong Upward Generalization Figure 2 shows that equilibrium models demonstrate very strongupward generalization performance compared to non weight-tied fixed-depth models.
Revised: Strong Upward Generalization Figure 2a shows that equilibrium models demonstrate very strongupward generalization performance compared to non weight-tied fixed-depth models.

Original: Moreover,
Revised: Moreover,

Original: Figure 3 shows that increasing inference depth consistently improves performance—especially onharder problem instances.
Revised: Figure 2b shows that increasing inference depth consistently improves performance—especially onharder problem instances.

Original: Terminology and Abbreviations We use the term ‘solver’ to refer to the use of black-box rootfinders like Anderson acceleration to find fixed points of a DEQ.
Revised: Terminology and Abbreviations We use the term ‘solver’ to refer to the use of black-box root finderslike Anderson acceleration to find fixed points of a DEQ.

Original: These networks use implicitgradients (IFT) for the backward pass.
Revised: These networks use implicit gradients (IFT)for the backward pass.

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

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

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

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

Original:  Path Independence
Revised:  Path Independence

Original: [Figure] Having intuitively motivated the idea of pathindependence in Sec. 1, we now formulate aprecise definition: we say that the computationperformed by a recurrent operator computingfunction f w on an input x is path independentif 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 formulate a precisedefinition: we say that the computation performed by a recurrent operator computing function f won an input x is path independent if it converges to the same limiting behavior regardless of thecurrent state z t .

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

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

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

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

Original: For instance, some classes ofequilibrium models have unique fixed points(such as Monotone Operator Equilibrium Networks (MON) [Winston and Kolter, 2020]), orany equilibrium model with a contractive recurrent cell.
Revised: However, most common

Original: However, most common DEQarchitectures—and the ones we will be using throughout this paper—have the expressivepower to easily learn multiple fixed points perinput.
Revised: DEQ architectures—and the ones we will be using throughout this paper—have the expressive power to easily learn multiple fixed points per input.

Original: Since it is unclear whether architecturesenforcing the contraction property lose expressiveness [Bai et al., 2021], we focus our investigation on architectures not guaranteed to bepath independent.
Revised: Since it is unclear whether architectures enforcing thecontraction property lose expressiveness [Bai et al., 2021], we focus our investigation on architecturesnot guaranteed to be path independent.

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

Original: We dedicate the restof the paper on exploring the Path Independence Hypothesis —the idea that equilibrium modelsshow strong upward generalization from easyto hard problem instances of algorithmic tasksbecause they learn path independent solutions.
Revised: We dedicate the rest of the paper on exploring the Path Independence Hypothesis —the idea thatpath dependent networks can better exploit test-time computation and hence show strong upwardgeneralization from easy to hard problem instances of algorithmic tasks.

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, then proceed to describe a metric toquantify how path independent an equilibrium model is on differing data distributions.
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 pathindependence, then proceed to describe a metric to quantify degrees of path independence.

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 fora learned model to be PI.
Revised: We observe that both architectural components are also necessary for alearned model to be PI 3 .

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

Original: Input injection, whichmimics skip connections in ResNets

Original: [He et al., 2016], ensures that the equilibrium point dependsthat on the input despite having an “infinite depth" network.
Revised: Input injection ensuresthat the equilibrium point depends on the input despite having an “infinite depth".

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

Original: Interestingly, both architectural motifs are also key components of deepequilibrium models [Bai et al., 2019]; in that work, the motivation was to enable efficient gradientestimation via the implicit function theorem (IFT)—this is closely related to path independence, sincethe premise of the IFT  is that only the final hidden state matters.
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 functiontheorem (IFT)— a concept closely related to path independence, since the premise of the IFT gradientestimator is that only the final hidden state matters, not the path taken to get there.

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

Original: We believe this is not a coincidence.
Revised: We believe this is not a coincidence.

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

Original: [2022], in Figure 2 we show upward generalization performance without input injection and weight tying.
Revised: [2022], in Figure 2a we show upward generalization performance using both equilibrium models and progressive nets [Bansal et al., 2022] – and the lackthereof using non-input injected networks.

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

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

Original: Input: A batched input (cid:20) xx(cid:21) , a network f w Initialize: (cid:20) zz(cid:21) =
Revised: We propose a simple metric to quantify path independence based on the directional alignment ofthe fixed points computed with the same input, but different initializations.

Original: Define: h ( y 1 , y 2 ) = y∥ y 1 ∥ 2 · y∥ y 2 ∥
Revised: We name this metric the Fixed Point Alignment (FPA) score .

Original: Compute (cid:20) z ′z ′(cid:21) =
Revised: Pseudocode to compute the metric is given in Algorithm Box1.

Original: f w (cid:18)(cid:20) zz(cid:21) , (cid:20) xx(cid:21)(cid:19) ▷ Interchange and reinitialize fixed points Compute z ′′ 2 = f w ( z ′ 1 , x 1 ) return h ( z ′′ 2 , z ′ 1 ) average cosine  the  (often simply the 
Revised: The FPA score is the average cosine similarity between the fixed point obtained with the trainingtime initialization (often simply the zero vector) and the fixed points obtained when one initializesthe solver using the fixed points computed on different inputs .

Original: ) and  the  obtained one the the computed on the
Revised: Higher FPA scores (with 1 being thehighest value) imply higher degrees of path independence. In Section 5, we show a strong correlationbetween path independence and accuracy using the FPA score. FPA score is cheap to compute, is a reliable indicator of path independence (see below) and is unitless,meaning networks obtained from different training runs can be compared on equal footing. See SM forother metrics we’ve considered for quantifying path independence and why we settled on FPA score.

Original: Stress-testing the FPA score To stress-testthe extent to which the FPA score really measures path independence, we search for adversarial initializations that are optimized to resultin distinct fixed points, hence low FPA values.
Revised: Stress-testing the FPA score To stress-test the extent to which the FPA score really measures pathindependence, we search for adversarial initializations that are optimized to result in distinct fixedpoints, hence low FPA 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 searchmultiple times starting from different  initializations.
Revised: [Liu and Nocedal, 1989] optimizer, and repeat the search multiple times startingfrom different fixed point initializations.

Original: We include pseudocode in the Supplementarymaterial.
Revised: We include pseudocode in the Supplementary material.

Original: The results of the adversarial stress test can be seen in Table 1.
Revised: The results of the adversarial stress test can be seen in Table 1.

Original: The results corroborate that the FPAscore is indeed a reliable measure of path independence; while it is not possible to find adversarialinitializations 
Revised: The results corroborate that the FPAscore is indeed a reliable measure of path independence; while it is not possible to find adversarialinitializations for high FPA networks (indicating high path independence), low FPA networks caneasily be adversarially initialized to be steered away from the original fixed point estimate. Input: A batched input (cid:20) xx(cid:21) , an operator f w Initialize: (cid:20) zz(cid:21) = Define: h ( y 1 , y 2 ) = y∥ y 1 ∥ 2 · y∥ y 2 ∥

Original: [Table]
Revised: Compute (cid:20) z ′z ′(cid:21) = FIX f w (cid:18)(cid:20) xx(cid:21) , (cid:20) zz(cid:21)(cid:19) # Interchange and reinitialize fixed points Compute (cid:20) z ′′z ′′(cid:21) =

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: [Figure]
Revised: [Table] [Table]

Revised: on 13 × 13 mazes and 64 bit prefix sum.

Revised: Attacked FPA refers to the cosine similarity between the fixed point from zero initialization and an adversarial initialization.

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

Original: To investigate whether path independence (as measured by the FPA score) is a strong predictor ofupwards generalization, we took the trained networks from Section 3, computed their average FPA Scores on in and out-of-distribution splits and inspected whether the FPA scores are correlated withupward generalization.
Revised: Is path independence (as measured by the FPA score) is a strong predictor of upwards generalization? We took the trained networks from Section 3, computed their average FPA Scores on in- and out-ofdistribution splits and inspected whether the FPA scores are correlated with upward generalization.

Original: On prefix sum experiments, we varied 1) network depth, 2) whether  weight norm (wnorm) [Salimansand Kingma, 2016] was used or not 3 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: On prefix sum experiments, we varied 1) network depth, 2) whether or not weight norm (wnorm)[Salimans and Kingma, 2016] was used or not 4 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 gradientestimator (backprop or implicit gradients) 5 .

Original: On the maze experiments, we varied 1) network depth, 2) use ofweight norm 3) forward solver (fixed point iterations or Broyden solver [Broyden, 1965]), and 5) thegradient 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]), and5) the gradient estimator (backprop or implicit gradients)

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

Original: We evaluated performance on in- and out-of-distribution validationsplits (different columns), using both small and large inference budgets (different rows). We matchedthe small-inference-budget conditions exactly with the training conditions to additionally measurethe effect of training time vs. test time path independence.
Revised: We evaluated performance on a mixture of in and OOD validationdata - results on individual data splits can be found in the Appendix.

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

Original: This shows that PI networks that allow for scaling test-time compute improve testtime accuracy.
Revised: This showsthat PI networks that allow for scaling test-time compute improve test-time accuracy (also see Figure2b).

Original: The in-distribution validation performance of non-PI networks degrade with deeperinference depths.
Revised: The in-distribution validation performance of non-PI networks degrade with deeper inferencedepths.

Original: Unsurprisingly, these networks generalize poorly on harder problem instances thatprovably require deeper inference depths. 
Revised: Unsurprisingly, these networks generalize poorly on harder problem instances that provablyrequire deeper inference depths. [Figure]

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

Original: Previous sections demonstrated a strong correlation between path independence and the ability toexploit test-time iterations.
Revised: Previous sections demonstrated a strong correlation between path independence and the ability toexploit test-time iterations.

Original: Unfortunately, we can’t make a causal claim based on those studies: eventhough algorithmic architectural and algorithmic features which promoted path independence alsoimproved accuracy, it could be that they improved accuracy through some unrelated mechanism.
Revised: Unfortunately, we can’t make a causal claim based on those 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 directlyencourage or penalize path independence.
Revised: In this section, we intervenedirectly on path independence by imposing regularizers which directly encourage or penalize pathindependence.

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

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

Original: A straightforward way to reduce the dependency of learned networks on initialization is simplyinitializing the hidden states with random noise.
Revised: A straightforward way to encourage path independence is simply to initialize the hidden states withrandom noise during training.

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

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

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

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

Original: Figure proves generalization in the prefix sum task: interventions that are designed to promote path independence (such as initializing fixed points with random noise or running the fixed point solver using a random number of steps during training) also bring along generalization benefit.

Original: Conversely, interventions that hurt path independence (facilitated by an alignment penalty that directly penalizes fixed point alignment) lead to poorer generalization.

Original: The reported numbers have been averaged over problem instances of different dimensions.

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

Original: As can be seen in the Figure 6, the interventions lead to strong test-time path independentneural networks, while also reliably improving in and out-of-distribution validation accuracy.
Revised: As can be seen in the Figure 3, theinterventions lead to strong test-time path independent neural networks, while also reliably improvingin- and out-of-distribution validation accuracy.

Original: We especially emphasize that shallow networks trained with mixed initialization actuallyremain far from having high FPA scores usingthe training-time forward pass conditions due tolack of convergence.
Revised: We especially emphasize that shallow networkstrained with mixed initialization actually remain far from having high FPA scores using the trainingtime forward pass conditions due to lack of convergence.

Original: However, since the mixedinitialization strategy results in path independentnetworks, scaling up test-time compute budgetleads to high FPA scores, and therefore highupwards generalization.
Revised: However, since the mixed initializationstrategy results in path independent networks, scaling up test-time compute budget leads to high FPAscores, and therefore high upwards generalization.

Original: 6.2 Penalizing Path Independencevia 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 pathindependence 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 auxiliaryloss term that penalizes the dot product betweenthe 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 + backpropagationnetworks with an additional auxiliary loss term that penalizes the dot product between the fixed pointscomputed from the same input, but different initializations sampled from standard Gaussian noise.

Original: Figure 6 shows that this intervention succeeded in pushing the FPA scores down, while also keeping the accuracy on the same trendline.
Revised: Figure 3 shows that this intervention succeeded in pushing the FPA scores down, while also keepingthe accuracy on the same trend line. [Table]

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 question negatively and showthat neither training-time convergence nor test-time convergence is required for path independence.
Revised: We answer this statement in the negative, and showthat 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 deep equilibrium models (DEQs) trained on the
Revised: Training Time Convergence We consider two deep equilibrium models (DEQs) trained on the

Original: Mazes task—one trained with implicit gradients computed via implicit function theorem (IFT), and the other trained with phantom gradients [Geng et al., 2021b].
Revised: Mazes task—one trained with implicit gradients computed via implicit function theorem (IFT), andthe other trained with phantom gradients [Geng et al., 2021b].

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

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

Original: The mixed-initialization intervention described in Section 6.1 also leads to a separation betweentraining-time convergence and path independence.
Revised: The mixed-initialization intervention described in Section 6.1 also leads to a separation betweentraining-time convergence and path independence.

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

Original: Details of this experiment are provided in the Supplementarymaterial.
Revised: Detailsare provided in the Supplementary material.

Original: [Table] Test Time Convergence From Table 2, one might conclude that test time convergence is importantfor path independence.
Revised: Test Time Convergence From Table 2, one might conclude that test time convergence is important forpath independence.

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

Original: We study test time convergenceproperties of an unrolled weight-tied input-injected network trained with backpropagation underdifferent solvers.
Revised: We study test time convergenceproperties of an unrolled weight-tied input-injected network trained with backpropagation underdifferent solvers.

Original: This network is path independent, as indicated by its high FPA scores (0.99) onboth in and out-of-distribution data.
Revised: This network is path independent, as indicated by its high FPA scores (0.99) onboth in- and out-of-distribution data.

Original: We visualize the values of residuals observed at test time withfixed point iterations and Broyden’s method in Figure 7a.
Revised: We visualize the values of residuals observed at test time withfixed point iterations and Broyden’s method in Figure 5a.

Original: Both these solvers converge to differentfixed points but still show good upward generalization.
Revised: Both these solvers converge to differentlimit cycles but still show good upward generalization.

Original: [Figure]  The Connection between  Path Independenceand Accuracy Holds on aPer-Example Level
Revised:  Path Independence on a Per-Example Level

Original: The connection between path independence and prediction correctness also  holds on a per-instance 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 performantgroup of networks in our intervention experiments), we plotted the distribution of per-instance fixedpoint alignment scores, colored by whether the prediction on that instance was correct or not in Figure 7b.
Revised: Using the prefix-sum networks trained with the mixed-initialization strategy (themost performant group of networks in our intervention experiments), we plotted the distribution ofper-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 whethera 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-checkto determine whether a prediction is correct or not without the need for any label data, both in- andout-of-distribution.

Original: Weprovide 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 oncomplexity of task they are learning to solve—akin to the intrinsic mechanism in humans to reasonand solve problems.
Revised: There is a long line of research on neural networks that can adapt their computational budget based oncomplexity of task they are learning to solve—akin to the intrinsic mechanism in humans to reasonand solve problems.

Original: Schmidhuber [2012] introduced self-delimiting neural networks which are a typeof 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 typeof 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, 2016] also uses the output of a sigmoidal haltingunit to determine the termination condition of RNN, but it avoids long "thinking" time by explicitlypenalizing it.
Revised: [Graves, 2016] also uses the output of a sigmoidal haltingunit to determine the termination condition of RNN, but it avoids long "thinking" time by explicitlypenalizing it.

Original: Subsequent works have successfully applied variants of ACT in image classification andobject detection [Figurnov et al., 2017], visual reasoning [Eyzaguirre and Soto, 2020], Transformers[Vaswani et al., 2017] for language modelling [Dehghani et al., 2019, Elbayad et al., 2020a, Liu et al.
Revised: Subsequent works have successfully applied variants of ACT in image classification andobject detection [Figurnov et al., 2017], visual reasoning [Eyzaguirre and Soto, 2020], Transformers[Vaswani et al., 2017] for language modelling [Dehghani et al., 2019, Elbayad et al., 2020a, Liu et al.

Original: ,2021b], and recognizing textual entailment [Neumann et al., 2016].
Revised: ,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 theloss 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 theloss objective to encourage exploration.

Original: With these additions, PonderNet can extrapolate to moredifficult examples on parity task.
Revised: With these additions, PonderNet can extrapolate to more [Figure] difficult examples on parity task.

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

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

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

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

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

Original: This approach of designing adaptive networks has beenadapted both in natural language processing [Schwartz et al., 2020, Soldaini and Moschitti, 2020,
Revised: This approach of designing adaptive networks has beenadapted 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 aretrained in multiple stages, and require complex exit policies.
Revised: Most of these architectures have complex sub-modules that aretrained in multiple stages, and require complex exit policies.

Original: In contrast, equilibrium models have asimple architecture, and can use root solvers to efficiently solve for the fixed point at inference.
Revised: In contrast, equilibrium models have asimple 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 onsimple algorithmic tasks like addition [Brown et al., 2020].
Revised: More complex transformer-based language models like GPT-3 also struggle to generalize well onsimple algorithmic tasks like addition [Brown et al., 2020].

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

Original: [2021] showsthat transformers can be trained to perform well on algorithmic tasks and generalize on OOD databy emitting the intermediate steps of algorithm to a buffer called ‘scratchpad’.
Revised: [2021] showsthat transformers can be trained to perform well on algorithmic tasks and generalize on OOD databy emitting the intermediate steps of algorithm to a buffer called ‘scratchpad’.

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

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] on easy tohard generalization of RNN on algorithmic tasks.
Revised: [2021b] on easy tohard generalization of RNN on algorithmic tasks.

Original: A subsequent work [Bansal et al., 2022] proposesarchitectural choices and training mechanisms that enables weight tied networks to converge to a fixedpoint, and generalize on harder problem instances.
Revised: A subsequent work [Bansal et al., 2022] proposesarchitectural choices and training mechanisms that enables weight tied networks to converge to a fixedpoint, and generalize on harder problem instances.

Original: We identify that these networks belong a moregeneral class of networks called equilibrium models that learn solutions that are path independent,and can therefore generalize on harder problem instances.
Revised: We identify that these networks belong a moregeneral class of networks called equilibrium models that learn solutions that are path independent,and can therefore generalize on harder problem instances.

Original:  Conclusion
Revised:  Conclusion

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

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

Original: Weshow that equilibrium models are capable of displaying upwards generalization by exploiting scalabletest-time compute.
Revised: We show that equilibrium models are capable of displaying upwards generalization by exploitingscalable 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 tendencyof an equilibrium network to converge to the same limiting behavior given an input, regardless of theinitial 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 pathindependent networks indeed generalize well on harder problem instances by exploiting more testtime 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 alsoimprove 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 larger compute
Revised: Our findings suggestthat path independent equilibrium models are a promising direction towards building general purposelearning systems whose test-time performance improves with larger compute.

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .

Original: • Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.

Original: • Did you include the license to the code and datasets?

Original: [N/A]

Original: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the

Original: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes](c)
Revised: [Yes]

Original:  Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes]
Revised: [Yes]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]2.
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A](b)
Revised: [N/A]

Original:  Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]
Revised: [N/A]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?
Revised: ments multiple times)?

Original: [Yes]
Revised: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type

Original: of GPUs, internal cluster, or cloud provider)?
Revised: of GPUs, internal cluster, or cloud provider)?

Original: [Yes]
Revised: [Yes]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]
Revised: [Yes]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if 
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if [Table]

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?

Original: [N/A]
Revised: [N/A]