Towards training digitally-tied analog blocks via hybrid gradient computation

Authors: timothy  nest
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 19:28:03

Abstract: Power efficiency is plateauing in the standard digital electronics realm such that new hardware, models, and algorithms are needed to reduce the costs of AI training. The combination of energy-based analog circuits and the Equilibrium Propagation (EP) algorithm constitutes a compelling alternative compute paradigm for gradient-based optimization of neural nets. Existing analog hardware accelerators, however, typically incorporate digital circuitry to sustain auxiliary non-weight-stationary operations, mitigate analog device imperfections, and leverage existing digital platforms. Such heterogeneous hardware lacks a supporting theoretical framework. In this work, we introduce \emph{Feedforward-tied Energy-based Models} (ff-EBMs), a hybrid model comprised of feedforward and energy-based blocks housed on digital and analog circuits. We derive a novel algorithm to compute gradients end-to-end in ff-EBMs by backpropagating and ``eq-propagating'' through feedforward and energy-based parts respectively, enabling EP to be applied flexibly on realistic architectures. We experimentally demonstrate the effectiveness of this approach on ff-EBMs using Deep Hopfield Networks (DHNs)  as energy-based blocks, and show that a standard DHN can be arbitrarily split into any uniform size while maintaining or improving performance with increases in simulation speed of up to four times. We then train ff-EBMs on ImageNet32 where we establish a new state-of-the-art performance for the EP literature (46 top-1 \%). Our approach offers a principled, scalable, and incremental roadmap for the gradual integration of self-trainable analog computational primitives into existing digital accelerators.

Main paper: 
 introduction: 
 Gradient-based optimization, the cornerstone and most energy greedy component of deep learning, fundamentally relies upon three factors: i) highly parallel digital hardware such as GPUs, ii) feedforward models and iii) backprop (BP) . With skyrocketing demands of AI compute, reducing the energy consumption of AI systems has become a matter of great economic, societal and environmental urgency [Strubell et al., 2020] , calling for the exploration of novel compute paradigms [Thompson et al., 2020 , Scellier, 2021 , Stern and Murugan, 2023] .One promising path towards this goal is analog in-memory computing [Sebastian et al., 2020] : by mapping weights onto a crossbar of resistive devices, Kirchoff current and voltage laws inherently perform matrix-vector multiplications with constant time complexity [Cosemans et al., 2019] . By stacking multiple such crossbars, an entire neural network can be mapped onto a physical system. An important formalism for such a system is that of energy-based (EB) analog circuits [Kendall et al., 2020 , Stern et al., 2023 , Dillavou et al., 2023 , Scellier, 2024] , which are "selflearning" systems that can compute loss gradients through two relaxations to equilibrium (i.e. two "forward passes"). Such a procedure falls under the umbrella of energy-based learning (EBL) algorithms [Scellier et al., 2024] . One such algorithm, Equilibrium Propagation (EP) [Scellier and Bengio, 2017] , particularly stands out for its strong theoretical guarantees, relative scalability in the realm of backprop alternatives [Laborieux and Zenke, 2022, 2023] and proven application on small analog systems with 10, 000× greater energy-efficiency and substantial speedups compared to its GPU-based counterpart [Yi et al., 2023] . This suggests a new alternative compute paradigm for gradient-based optimization consisting of: i) analog hardware, ii) EBMs, and iii) EP. In this paper, we propose a theoretical framework for extending end-to-end gradient computation to a realistic setting where the system in question may or may not be fully analog. Such a setting is plausible in the near term, for two major reasons. First, analog circuits exhibit non-ideal physical behaviors which affect both the inference pathway [Wang et al., 2023 , Ambrogio et al., 2023] and parameter optimization [Nandakumar et al., 2020 , Spoon et al., 2021 , Lammie et al., 2024] , compromising performance. Second, owing to the latency and energy-consumption of resistive devices' write operations, analog circuits should be fully weight stationary -weights must be written before the inference procedure begins -which excludes many operations used conventionally in machine learning including activation functions, normalization, and attention [Spoon et al., 2021 , Jain et al., 2022 , Liu et al., 2023 , Li et al., 2023] . Therefore, analog systems are likely to be used in combination with auxiliary digital circuitry, resulting in hybrid mixed precision systems [Haensch et al., 2018] . While the design of purely inferential engines made up of analog and digital parts is nearing commercial maturity [Ambrogio et al., 2023] , in-situ learning of such systems has barely been explored. An important challenge remains in proving EBL algorithms can scale in a manner comparable to backprop, given the requirement of simulating EB systems on GPUs. Because of the necessity of convergence, this amounts in practice in performing lengthy root finding algorithms to simulate physical equilibrium, limiting proof-of-concepts thereof to relatively shallow (5-6 layer) models [Scellier et al., 2024 , Scellier, 2024] .Our work contends that the best of both worlds can be achieved with the following triad: i) hybrid digital and analog hardware, ii) feedforward and EB models, iii) BP and EP. Namely, by modeling digital and analog parts as feedforward and EB modules respectively, we show how backprop and EP error signals can be chained end-to-end via feedforward and EB blocks respectively in a principled fashion. Rather than opposing digital and analog, or backprop and "alternative" learning algorithms, as is often done in the literature, we propose a novel hardware-aware building block which can, in principle, leverage advances from both digital and analog hardware in the near-term. More specifically:• We propose Feedforward-tied Energy-based Models (ff-EBMs, Section 3.1) as high-level models of mixed precision systems whose inference pathway read as the composition of feedforward and EB modules (Eq. (5), Alg. 1).• We show that gradients in ff-EBMs can be computed in an end-to-end fashion (Section 3.3), backpropagating through feedforward blocks and "eq-propagating" through EB blocks (Theorem 3.1, Alg. 2) and that this procedure is rooted in a deeply-nested optimization problem (Section 3.2).• Finally, we experimentally demonstrate the effectiveness of our algorithm on ff-EBMs where EBM blocks are Deep Hopfield Networks (DHNs) (Section 4).In particular we show that i) final and transient gradient estimates computed by our algorithm (Alg. 2) near perfectly match gradients computed by end-to-end automatic differentiation (Section 4.2), which we also prove mathematically (Theorem 4.1), ii) a standard DHN model can be arbitrarily split into a ff-DHN with the equivalent layers and architectural layers while maintaining or improving performance, remaining on par with automatic differentiation and being up to four times faster to simulate depending on the convergence criterion at use to compute equilibrium (Section 4.3), iii) the proposed approach yields 46 % top-1 (70% top-5) validation accuracy on ImageNet32 when training a ff-EBM of 15 layers, beating current state-of-the-art for EP by a large margin, without relying on holomorphic transformations inside EBM blocks [Laborieux and Zenke, 2022, 2023] 2 Background Notations. Given A : R n → R m a differentiable mapping, we denote its total derivative with respect to s j as d sj A(s) := dA(s)/ds j ∈ R m , its partial derivative with respect to s j as ∂ j A(s) := ∂A(s)/∂s j ∈ R m . When A takes scalar values (m = 1), its gradient with respect to s j is denoted as∇ j A(s) := ∂ j A(s) ⊤ . 
energy-based models (ebms): 
 For a given static input and set of weights, Energy-based models (EBMs) implicitly yield a prediction through the minimization of an energy function. As such they are a particular kind of implicit model. Namely, an EBM is defined by a (scalar) energy function E : s, θ, x → E(s, θ, x) ∈ R where x, s, and θ respectively denote a static input, hidden and output neurons and model parameters, and each such tuple defines a configuration with an associated scalar energy value. Among all configurations for a given input x and some model parameters θ, the model prediction s ⋆ is implicitly given as an equilibrium state which minimizes the energy function:s ⋆ := arg min s E(s, θ, x).(1)Assuming that ∇ 2 s E(x, s ⋆ , θ) is invertible, note that the equilibrium state s ⋆ implicitly depends on x and θ by virtue of the implicit function theorem [Dontchev et al., 2009] . Therefore our goal when training an EBM-in a supervised setting, for instance -is to adjust the model parameters θ such that s ⋆ (x, θ) minimizes some cost function ℓ : s, y → ℓ(s, y) ∈ R where y is some ground-truth label associated to x. More formally, this learning objective can be stated with the following bilevel optimization problem [Zucchet and Sacramento, 2022] :EQUATIONSolving Eq. ( 2) in practice amounts to computing the gradient of its outer objective C(x, θ) with respect to θ (d θ C(x, θ)) and then performing gradient descent over θ.An algorithm used to train an EBM model in the sense of Eq. ( 2) may be called an EBL algorithm [Scellier et al., 2024] . Equilibrium Propagation (EP) [Scellier and Bengio, 2017] is an EBL algorithm which computes an estimate of d θ C(x, θ) with at least two phases. During the first phase, the model is allowed to evolve freely to s ⋆ = arg min s E(s, θ, x). Then, the model is slightly nudged towards decreasing values of cost ℓ and settles to a second equilibrium state s β . This amounts to augmenting the energy function E by an additional term βℓ(s, y) where β ∈ R ⋆ is called the nudging factor. Next, the weights are updated to increase the energy of s ⋆ and decrease that of s β , thereby "contrasting" these two states. More formally, Scellier and Bengio [2017] show in the seminal EP paper:EQUATIONwhere α denotes some learning rate. EP comes in different flavors depending on the sign of β inside Eq. (3) or on whether two nudged states of opposite nudging strengths (±β) are contrasted, a variant called Centered EP (C-EP) which was shown to work best in practice [Laborieux et al., 2021 , Scellier et al., 2024] and reads as:EQUATION3 Tying energy-based models with feedforward blocksIn the present section we introduce a new model, Feedforward-tied EBMs (ff-EBMs, section 3.1), which read as composition of feedforward and EB transformations (Alg. 1). We show how optimizing ff-EBMs amounts to solving a multi-level optimization problem (Section 3.2) and propose a BP-EP gradient chaining algorithm as a solution(Section 3.3, Theorem 3.1, Alg. 2). We highlight as an edge case the fact that ff-EBMs reduce to standard feedforward nets (Lemma A.1) and the proposed BP-EP gradient chaining algorithm to standard BP (Corollary A.1) when each EB block comprises a single hidden layer. We highlight in red and blue the parts of the model and associated algorithms performed inside feedforward (digital) and EB (analog) blocks respectively.3.1 Feedforward-tied Energy-based Models (ff-EBMs)Inference procedure. We define Feedforward-tied Energy-based Models (ff-EBMs) as compositions of feedforward and EB transformations. Namely, an data sample x is fed into the first feedforward transformation F 1 parametrized by some weights ω 1 , which yields an output x 1 ⋆ . Then,x 1⋆ is fed as a static input into the first EB block E 1 with parameters θ 1 , which relaxes to an equilibrium state s 1 ⋆ . s 1 ⋆ is in turn fed into the next feedforward transformation F 1 with weights ω 1 and the above procedure repeats until reaching the output layer ô. More formally, denoting F k and E k the k th feedforward and EB blocks parametrized by the weights ω k and θ k respectively, the inference pathway of a ff-EBM reads as:     s 0 := x x k ⋆ := F k (s k-1 ⋆ , ω k ), s k ⋆ := arg min s E k (s, θ k , x k ⋆ ) ∀k = 1 • • • N -1 ô⋆ := F N (s N -1 ⋆ , ω N ) (5)The ff-EBM inference procedure is depicted more compactly inside Fig. 2 (left) and Alg. 1. Algorithm 1 ff-EBM inference (Eq. ( 5))1: s ← x 2: for k = 1 • • • N -1 do 3: x ← F k s, ω k 4: s ← Optim s E k (s, θ k , x) 5: end for 6: ô ← F N (s, ω N )Form of the energy functions. We further specify the form of the energy of the k th EB block of a ff-EBM as defined per Eq. ( 5). The associated energy function E k takes some static input x k from the output of the preceding feedforward transformation, has hidden neurons s k and is parametrized by weights θ k . More precisely: 6) reveals three different contributions to the energy. The first term determines the non-linearity applied inside the EB block [Zhang and Brand, 2017, Høier et al., 2023] : for a given invertible and continuous activation function σ, G is defined such that ∇G = σ -1 (see Appendix A.1.3).E k (s k , θ k , x k ) := G k (s k ) -s k ⊤ • x k + U k (s k , θ k ) (6) Eq. (The second term inside Eq. ( 6) accounts for a purely feedforward contribution from the previous feedforward block F k . Finally, the third term accounts for internal interactions within the layers of the EB block.Recovering a feedforward net. When taking the gradient of E k as defined in Eq. ( 6) with respect to s k and zeroing it out, it can be seen that s k ⋆ is implicitly defined as:EQUATIONA noteworthy edge case highlighted by Eq. ( 7) is when U k = 0 for all k's, i.e. when there are no intra-block layer interactions, or equivalently when the EB block comprises a single layer only. In this case, s k ⋆ is simply a feedforward mapping x k through σ and in turn the ff-EBM is simply a standard feedforward architecture (see Lemma A.1 inside Appendix A.1.1). 
multi-level optimization of ff-ebms: 
 Just as learning EBMs can be naturally cast as a bilevel optimization problem, learning ff-EBMs equates to a multi-level optimization problem where the variables being optimized in the inner subproblems are comprised of EB block variables s 1 , • • • , s N -1 . To make this clearer, we re-write the energy function of the k th block E k from Eq. ( 6) to highlight the dependence between two consecutive EB block states:E k (s k , θ k , s k-1 ⋆ , ω k ) := E k s k , θ k , F k s k-1 ⋆ , ω k-1 (8)It can be seen from Eq. ( 8) that the equilibrium state s k ⋆ obtained by minimizing E k will be dependent upon the equilibrium state s k-1 ⋆ of the previous EB block, which propagates back through prior EB blocks. DenotingW := {θ 1 , • • • , θ N -1 , ω 1 , • • • , ω N }, the learning problem for a ff-EBM can therefore be written as:min W C(x, W, y) := ℓ(ô ⋆ = F N (s N -1 ⋆ , ω N ), y) (9) s.t. s N -1 ⋆ = arg min s E N -1 (s, θ N -1 , s N -2 ⋆ , ω N -1 ) • • • s.t. s 1 ⋆ = arg min s E 1 (s, θ 1 , x, ω 1 )Here again and similarly to bilevel optimization, solving Eq. ( 9) in practice amounts to computing g θ k := d θ k C and g ω k := d ω k C and performing gradient descent on θ k and ω k .Main result: explicit BP-EP chaining. Based on the multilevel optimization formulation of ff-EBMs learning in Eq. ( 9), we state the main theoretical result of this paper in Theorem 3.1 (see proof in Appendix A.2.1).Theorem 3.1 (Informal). Assuming a model of the form Eq. ( 5), we denotes 1 ⋆ , x 1 ⋆ , • • • , s N -1 ⋆, ô⋆ the states computed during the forward pass as depicted in Alg. 1. We define the nudged state of block k, denoted as s k β , implicitly through ∇ 1 F k (s k β , θ k , x k ⋆ , δs k , β) = 0 with:F k (s k , θ k , x k ⋆ , δs k , β) := E k (s k , θ k , x k ⋆ ) + βs k ⊤ • δs k (10)Denoting δs k and ∆x k the error signals computed at the input of the feedforward block F k and of the EB block E k respectively, then the following chain rule applies:δs N -1 := ∇ s N -1 ℓ(ô ⋆ , y), g ω N = ∇ ω N ℓ(ô ⋆ , y) (11) ∀k = 2 • • • N -1 :    ∆x k = d β ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0 , g θ k = d β ∇ 2 E k (s k β , θ k , x k ⋆ ) β=0 δs k-1 = ∂ 1 F k s k-1 ⋆ , ω k ⊤ • ∆x k , g ω k = ∂ 2 F k s k-1 ⋆ , ω k ⊤ • ∆x k (12)Proposed algorithm: implicit BP-EP chaining. Theorem 3.1 reads intuitively. It prescribes an explicit chaining of EP error signals passing backward through E k (δs k → ∆x k ) and BP error signals passing backward through ∂F k ⊤ (∆x k → δs k-1 ), which directly mirrors the ff-EBM inference pathway as depicted in Fig. 2 . Yet noticing that:   δs k-1 = ∂ 1 F k s k-1 ⋆ , ω k ⊤ • ∆x k = d β ∇ 3 E k s k β , θ k , s k-1 ⋆ , ω k β=0 , g ω k = ∂ 2 F k s k-1 ⋆ , ω k ⊤ • ∆x k = d β ∇ 4 E k s k β , θ k , s k-1 ⋆ , ω k β=0 ,the same error signal can by passed through E k (δs k → δs k-1 ) where BP and EP are implicitly chained inside E k (see Appendix A.2.1). This insight, along with a centered scheme to estimate derivatives with respect to β around 0 as done for the C-EP algorithm (Eq. ( 4)), motivates the implicit BP-EP gradient chaining algorithm in Alg. 2 we used for our experiments (see Alg. 5 inside Appendix A.3.1 for its explicit counterpart). Given that the proposed algorithm appears as a a generalization of EP, we refer to Alg. 2as "EP" in the experimental section, for simplicity.Algorithm 2 Implicit BP-EP gradient chaining (Theorem (3.1)) 1: δs, g ω N ← ∇ s N -1 ℓ(ô ⋆ , y), ∇ ω N ℓ(ô ⋆ , y) ▷ Single BP step 2: for k = N -1 • • • 1 do 3: s β ← Optim s E k (s, θ k , s k-1 ⋆ , ω k ) + βs ⊤ • δs ▷ EP in E k 4: s -β ← Optim s E k (s, θ k , s k-1 ⋆ , ω k ) -βs ⊤ • δs 5: g θ k ← 1 2β ∇ 2 E k (s β , θ k , s k-1 ⋆ , ω k ) -∇ 2 E k (s -β , θ k , s k-1 ⋆ , ω k ) 6: g ω k ← 1 2β ∇ 4 E k (s β , θ k , s k-1 ⋆ , ω k ) -∇ 4 E k (s -β , θ k , s k-1 ⋆ , ω k ) ▷ Implicit BP in F k 7: δs ← 1 2β ∇ 3 E k (s β , θ k , s k-1 ⋆ , ω k ) -∇ 3 E k (s -β , θ k , s k-1 ⋆ , ω 
experiments: 
 In this section, we present the ff-EBMs used in our experiments (Section 4.1) and carry out static gradient analysis -computing and analyzing ff-EBM parameter gradients for some x and y (Section 4.2). We extend the observation made by Ernoult et al. [2019] -that transient EP parameter gradients obtained during the second phase match those computed by automatic differentiation through equilibrium and across blocks-to ff-EBMs (Fig.(3)-(4), Theorem 4.1), showing that gradient estimates of automatic differentiation and EP in our framework, are near perfectly aligned (Fig. 5 ). We then show on the CIFAR-10 task that performance of ff-EBMs can be maintained or improved across various block splits maintaining the same number of layers, while remaining on par with automatic differentiation(Section 4.3). We show furthermore that blocks of smaller size are up to four times faster to simulate depending on the convergence criterion at use for computing equilibrium inside EB blocks. Finally, we perform further ff-EBM training experiments on CIFAR-100 and ImageNet32 where we establish a new state-of-the-art performance in the EP literature (Section 4.4).Models. Using the same notations as in Eq. ( 6), the ff-EBMs at use in this section are defined:U k FC (s k , θ k ) := -1 2 s k⊤ • θ k • s k , U k CONV (s k , θ k ) := -1 2 s k • θ k ⋆ s k , F k BN (s k-1 , ω k ) := BN P ω k CONV ⋆ s k-1 L ; ω k α , ω k β , F k ID (s k-1 ) := s k-1 L (13) with BN(•; ω k α , ω k β), P and ⋆ the batchnorm, pooling and convolution operations, • the generalized dot product for tensors ands k := s k ⊤ 1 , • • • s k ⊤ L ⊤the state of block k comprising L layers. Such EBM blocks are known as Deep Hopfield Networks (DHNs). DHNs are comprised of fully connected (U k FC ) and convolutional operations (U k CONV ) forming a symmetric weight matrix θ k with a sparse, block-wise structure such that each layer s k ℓ is bidirectionally connected to its neighboring layers s k ℓ-1 and s k ℓ+1 through connections θ k ℓ-1 and θ k ⊤ ℓ respectively (see Appendix A.1.3). To empirically ensure convergence, the non-linearity σ applied within EB blocks is σ α (x) := min (max (αx, 0) , 1) with α ∈ (0, 1). Finally, two design choices were instrumental to the success of ff-EBM gradient computation and subsequent training. First, we initialized the weights of U k FC and U k CONV usingGaussian Orthogonal Ensembles (GOE) [Agarwala and Schoenholz, 2022] to enable faster equilibrium computation (see next paragraph). Second, while the last layer of a given block was simply passed as an input to the next block (i.e. using F k ID in Eq. ( 13)) for small enough models (L = 6 inside the experiment depicted in Section 4.3), the use of batchnorm layers in between blocks (i.e. using F k BN in Eq. ( 13)) becomes essential for deeper models.Equilibrium computation. As depicted in Alg. 2, the steady states s ±β may be computed with any loss minimization algorithm. Here, as in past works on EP [Ernoult et al., 2019 , Laborieux et al., 2021 , Laborieux and Zenke, 2022 , Scellier et al., 2024] , we employ a fixed-point iteration scheme to compute the EB blocks steady states. Namely, we iterate Eq. ( 7) until reaching equilibrium (the same scheme is used for ff-EBM inference, Alg. 1, with β = 0.):s k ±β,t+1 ← σ x k -∇ 1 U k (s k ±β,t , θ k ) ∓ βδs k (14)Important details about how Eq. ( 14) is executed in practice have to be highlighted. First, we employ a scheme to asynchronously update even (s k 2ℓ ′ ) and odd (s k 2ℓ ′ +1 ) layers [Scellier et al., 2024] -see Appendix A.1.3. Second, Eq. ( 14) were either executed for a fixed and predetermined number of steps as done in the aforementioned EP literature, or using an ϵ-tolerance-based convergence criterion (TOL) which stops executing Eq. ( 14) when (s i t+1 -s i t )/s i t ≤ ϵ on average -see Appendix A.5.3 for details.Algorithm baseline. As an algorithmic baseline, we simply use automatic differentiation (AD) backward through the fixed-point iteration scheme Eq. ( 14) with β = 0 and directly initializing s k t=0 = s ⋆ (Fig. 4 ). This version of AD, where we backpropagate through equilibrium, is known as "Recurrent Backpropagation" [Almeida, 1987 , Pineda, 1987] or Implicit Differentiation (ID).In order to study the transient dynamics of ID and EP, we define, with W k := {θ k , ω k }:EQUATIONwhere s k ±β,t is computed from Eq. ( 14) with the nudging error current δs k computed with Alg. 2, and T is the total number of iterations used for both ID and EP in the gradient computation phase.Figure 3 : EP and ID partially computed gradients (( g EP w (t)) t≥0 in black dotted curves and ( g ID w (t)) t≥0 in plain colored curves) going backward through equilibrium for ID and forward through the nudging phase for EP [Ernoult et al., 2019] for a random sample x and associated label y. The ff-EBM comprises 6 blocks and 15 layers in total, with block sizes of either 2 or 3 layers. Each sub-panel represents a layer (labeled on the y-axis) with each curve corresponding to a randomly selected weight. "Backward" time is indexed from t = 0 to T = 120, starting from block 6 backward to block 1, with 20 fixed-point iteration dynamics (Eq. ( 14)) being used for both EP and ID within each block. 4 depicts the computational graph that is differentiated through when using ID and shows where g ID W k (t) are obtained correspondingly. Similarly, g EP W k (t) is an EP gradient truncated at t steps forward through the nudged phase. When T is sufficiently large,k, d W k (T -k) C(x, W, y) is the "sensitivity" of the loss C to parameter W k at timestep T -k so that g ID W k (t) is a ID gradient truncated at T -t. Fig.g ID W k (T ) and g EP W k (T ) converge to d W k C(x, W, y). Fig. 3 displays ( g ID W k (t)) t≥0 and ( g EP W k (t))t≥0 on an heterogeneous ff-EBM of 6 blocks and 15 layers (16 if counting the last linear "readout" layer computing the logits) with blocks comprising 2 or 3 layers for a randomly selected sample x and its associated label y -see caption for a detailed description. It can be seen EP and ID error weight gradients qualitatively match very well throughout time, across layers and blocks. We also display the cosine similarity between the final EP and ID weight gradient estimate g ID W k (T ) and g EP W k (T ) for each layer and observe that EP and ID weight gradients are near perfectly aligned. Theorem 4.1 generalizes the equivalence between EP and ID to ff-EBMs [Ernoult et al., 2019] .EQUATION... ...Light grey: computational graph associated with ff-EBM inference (Alg. 1) when applying fixed-point iteration to compute equilibrium states within each block (Eq. ( 14)) where the node s k t denotes the state of block k (comprising several layers) at timestep t. Blue arrows: backward automatic differentiation (AD) through the computational graph where ĝID W k (t) is the partially computed gradient truncated at T -t. Since the states which are differentiated through are taken at equilibrium(s k t = s k ⋆ ∀t = 0 • • • τ )this instantiation of AD can be viewed as Implicit Differentiation (ID).For a given (standard, single block) EBM with a fixed number of layers, we ask how block splitting of this EBM into a ff-EBM with multiple EB blocks affects training performance and Wall Clock (simulation) Time (WCT). We address this question with two different depths (L = 6 and L = 12 layers in total) and various block sizes (bs), maintaining a fixed total number of layers (e.g. for L = 6, 1 block of 6 layers, 2 blocks of 3 layers, etc.). Additionally, to ensure the fairest comparison in terms of WCTs across different splits, we adopt the aforementioned TOL approach to execute the fixed-point dynamics Eq. ( 14) within each EB block. We display the results obtained on the CIFAR-10 task inside Table 1 . We observe that EP performance improves with smaller block sizes Figure 5 : Cosine similarity between EP and ID weight gradients on a randomly selected sample x and associated label y in the same setting as Fig. 3 using the same color code to label the layers. We observe near-perfect alignment between EP and ID gradients. See Fig. 7 for a precise depiction of the model at use. (reaching 90.1% and 92.5% for L = 6 and L = 12 respectively with bs = 2 and bs = 3) with overall WCT reduction (up to ≈ ×4) while remaining on par with ID. This significant reduction in WCT is due to the fact that inference time for ff-EBMs with DHN blocks by construction scales linearly with the number of blocks rather than supralinearly with the number of layers as has been empirically observed in the EP literature [Ernoult et al., 2019] . When instead using a fixed number of iterations to execute Eq. ( 14) inside EB blocks (Table 5 in Appendix A.5.3), EP performance is maintained across all splits (90.1% and 92.5% for L = 6 and L = 12 resp.) and is still on par with ID. However, there is no advantage in terms of WCTs in this case as the number of iterations is kept the same across all block splits and is much larger than necessary for smaller block sizes. Results for L = 6 are consistent with the existing literature and those for L = 12 surpass EP state-of-the-art on CIFAR-10 [ Scellier et al., 2024, Laborieux and Zenke, 2022] . Overall these results suggest that: i) ff-EBM performance is agnostic to EB block sizes and are therefore flexible in design, ii) ff-EBMs are much faster to simulate that EBM counterparts of equivalent depth.We now consider ff-EBMs of fixed block size 2, and relatively small number of iterations . We train two models of depth (L = 12 and L = 15) on CIFAR-100 and ImageNet32 by EP and ID and show the results obtained in Table 2 . Here again we observe that EP matches ID performance on all models and tasks, ff-EBMs benefit from depth, and the performance obtained by training the 15-layer deep ff-EBM by EP exceeds state-of-the-art performance on ImageNet32 by around 10% top-1 validation accuracy [Laborieux and Zenke, 2022] and by around 5% the best performance reported on this benchmark among all backprop alternatives [Høier et al., 2023] . Zach [2021] where feedforward model learning is cast as a deeply nested optimization in which consecutive layers are tied by elemental pair-wise energy functions. This work more recently inspired the Dual Propagation algorithm [Høier et al., 2023] . Such a setting can be construed as a particular case of ff-EBM learning by EP where each EB block comprises a single layer (U k = 0) inside Eq. ( 6)-which, as we have shown, is tantamount to standard BP(see last paragraph of Section 3.3).Forward-only learning beyond EP. Given that zeroth-order (ZO) optimization and "forwardforward" (FF) algorithms [Dellaferrera and Kreiman, 2022, Hinton, 2022] can be applied to any model, and-like EP-compute a learning rule through multiple inference steps, one may wonder why it is important that our models should be energy-based. While mechanistically appealing for analog hardware [Oguz et al., 2023 , Momeni et al., 2023 , 2024 , Xue et al., 2024] , these forward-only approaches do not match the performance of automatic differentiation on equivalent models, even if they are roughly the same size as those studied in our work. On the one hand, weight perturbation [Fiete et al., 2007] (WP or "SPSA" [Spall, 1998] ), yields unbiased yet noisy gradient estimates with variance scaling cubically with the model dimensionality [Ren et al., 2022] , resulting in a significant gap in model performance compared to backprop, that can only be partially mitigated when using heuristics [Silver et al., 2021 , Ren et al., 2022 , Fournier et al., 2023 , Chen et al., 2023] . On the other hand FF algorithms, as learning heuristics, suffer from a lack of theoretical guarantees which may impact the resulting model performance.Limitations and future work. Since our recipe advocates EP-BP chaining by construction, it is fair to say that ff-EBM learning partially inherits the pitfalls of BP. Fortunately, nothing prevents feedforward modules inside ff-EBMs from being trained by any BP alternative to mitigate specific issues. For instance: BP can be parameterized by feedback weights to obviate weight transport from the inference circuit to the gradient computation circuit [Akrout et al., 2019] ; BP gradients can be approximated as finite differences of feedback operators [Ernoult et al., 2022] ; or computed via implicit forward-mode differentiation [Hiratani et al., 2022 , Fournier et al., 2023 , Malladi et al., 2023] ; local layer-wise self-supervised or supervised loss functions can be used to prevent "backward locking" [Belilovsky et al., 2019 , Ren et al., 2022 , Hinton, 2022] . This insight may help exploring many variants of ff-EBM training.Figure 6 : ff-EBMs as hierarchical systems implementing EP at chip scale (adapted from [Yi et al., 2023] ) using energy-based analog processors made up of resistors (green edges), diodes (in blue), voltage sources (in purple), ADCs and DACs (adapted from [Scellier, 2024] ), digital processors, memory buffers, all of these being connected by digital buses (red lines).Pursuing the core motivation of this work, one natural extension of this study is to incorporate more hardware realism into ff-EBMs. Beyond Deep Hopfield networks, Deep Resistive Nets (DRNs) -developed by Scellier [2024] and strongly inspired by Kendall et al. [2020] -are exact models of idealized analog circuits, trainable by EP, promising fast simulation times. As such, using DRNs as EB blocks inside ff-EBMs is an exciting research direction -see Fig. 6 . Still, further work in this direction presents new challenges especially given device non-idealities which may affect the inference pathway, such as analog-todigital and digital-to-analog noise [Rasch et al., 2023 , Lammie et al., 2024] . Finally, considerable work is needed to prove ff-EBM further at scale on more difficult tasks (e.g. standard ImageNet), considerably deeper architectures, and moving beyond vision tasks. One other exciting research direction would be the design of ff-EBM based transformers, with attention layers being chained with energy-based fully connected layers inside attention blocks.Concluding remarks and broader impact. We show that ff-EBMs constitute a novel framework for deep-learning in heterogeneous hardware settings. We hope that the proposed algorithm can help to overcome the typical division between digital versus analog or BP versus BP-free algorithms and that the greater energy-efficiency afforded by this framework provides a pragmatic, near-term blueprint to mitigating the dramatic carbon footprint of AI training [Strubell et al., 2020] . While we are still a long way from fully analog training accelerators at commercial maturity, we believe this work offers an incremental and sustainable roadmap to gradually integrate analog, energy-based computational primitives as they are developed into existing digital accelerators.A.1.1 Feedforward-tied EBMs (ff-EBMs)We first formally define Feedforward-tied Energy-based Models (ff-EBMs) with precise assumptions on the energy-based and feedforward blocks.Definition A.1 (ff-EBMs). A Feedforward-tied Energy-based Model (ff-EBM) of size N comprises N twice differentiable feedforward mapping F 1 , • • • , F N and N -1 twice differentiable energy functions E 1 , • • • , E N -1with respect to all their variables. For a given x, the inference procedure reads as:     s 0 := x x k ⋆ := F k (s k-1 ⋆ , ω k ), s k ⋆ := arg min s E k (s, θ k , x k ⋆ ) ∀k = 1 • • • N -1 ô⋆ := F N (s N -1 ⋆ , ω N ) (17)Finally, we assume that ∀k = 1• • • N -1, ∇ 2 1 E k (s k ⋆ , θ k , ω k ) is invertible.A.1.2 Feedforward nets as a special caseWe show that when energy-based blocks comprise a single layer only, the ff-EBM becomes purely feedforward.Lemma A.1. We consider ff-EBM per Def. (A.1) where the energy functions E k have the form:EQUATIONWe assume thatU k = 0 for k = 1 • • • N -1, s → ∇G(s) is invertible and we denote σ := ∇G -1 .Then, the resulting model is a feedforward model described by the following recursive equations:EQUATIONTherefore Eq. ( 19) is immediately obtained from Eq. ( 20) with U k = 0.For a single block. As mentioned in Section 3.1, the energy function of the k th EB block has the form:EQUATIONwhere x k is the output of the preceding feedforward block. For a given choice of a continuously invertible activation function, G k σ is defined as:EQUATIONTo be more explicit and as we did previously, we re-write the augmented energy-function which encompasses both the k th EB block and the feedforward module that precedes it:EQUATION)Deep Hopfield Nets (DHNs) as EB blocks. In our experiments, we used weight matrices of the form:EQUATIONwhereby each layer ℓ is only connected to its adjacent neighbors. Therefore, fully connected and convolutional DHNs with L layers have an energy function of the form:EQUATION)Synchronous fixed-point iteration. We showed that when G is chosen such that ∇G = σ -1 for some activation function σ, then the steady state of the k th block reads:EQUATIONwhich justifies the following fixed-point iteration scheme, when the block is influenced by some error signal δs with nudging strength β:EQUATION)The dynamics prescribed by Eq. 28 are also used for the inference phase with β = 0. To further refine Eq. ( 28), let us re-write Eq. ( 28) with a layer index ℓ where ℓ ∈ [1, L k ] with L k being the number of layers in the k th block, and replacing x k by its explicit expression:EQUATION)As done in past EP works [Ernoult et al., 2019 , Laborieux et al., 2021 , Laborieux and Zenke, 2022 , 2023 , Scellier et al., 2024] and for notational convenience, we introduce the primitive function of the k th block as:Φ k s k , θ k , s k-1 ⋆ , ω k := s k ⊤ • F k s k-1 ⋆ , ω k -U k (s k , θ k ) (30)such that Eq. ( 29) re-writes:EQUATIONEq. ( 31) depicts a synchronous scheme where all layers are simultaneously updated at each timestep.Asynchronous fixed-point iteration. Another possible scheme, employed by Scellier et al. [2024] , instead prescribes to asynchronously update odd and even layers and was shown to speed up convergence in practice:EQUATION)We formally depict this procedure as the subroutine Asynchronous inside Alg. 3. In practice, we observe that it was more practical to use a fixed number of iterations rather than using a convergence criterion with a fixed threshold.Algorithm 3 Asynchronous (for all blocks until penultimate)Input: T , θ k , ω k , s k-1 ⋆ , β, δs k Output: s k β 1: s k ← 0 2: for t = 1 • • • T do 3: ∀ odd ℓ ∈ {1, • • • , L k } : s k ℓ,β ← σ ∇ s k ℓ Φ s k β , θ k , s k-1 ⋆ , ω k -βδs k 4: ∀ even ℓ ∈ {1, • • • , L k } : s k ℓ,β ← σ ∇ s k ℓ Φ s k β , θ k , s k-1 ⋆ , ω k -βδs k 5: end forResulting ff-EBM inference algorithm. With the aforementioned details in hand, we re-write the inference algorithm Alg. 1 presented in the main as a Forward subroutine.Algorithm 4 ForwardInput: T , x, W = {θ 1 , ω 1 , • • • ω N } Output: s 1 , • • • , s N -1 or ô depending on the context 1: s 0 ← x 2: for k = 1 • • • N -1 do 3: s k ← Asynchronous T, θ k , ω k , s k-1 ▷ Alg. 3 4: end for 5: ô ← F N s, ω N A.2 Main theoretical result A.2.1 Proof of Theorem 3.1The proof of Theorem 3.1 is structured as follows:• We directly solve the multilevel problem optimization defined inside Eq. ( 9) using a Lagrangian-based approach (Lemma A.2), yielding optimal Lagrangian multipliers, block states and loss gradients.• We show that by properly nudging the blocks, EP implicitly estimates the previously derived Lagrangian multipliers (Lemma A.3).• We demonstrate Theorem 3.1 by combining Lemma A.2 and Lemma A.3.• Finally, we highlight that when a ff-EBM is a feedforward net (Lemma A.1), then the proposed algorithm reduces to BP (Corollary A.1).Lemma A.2 (Lagrangian-based approach). Assuming a ff-EBM (Def. A.1), we denotes 1 ⋆ , x 1 ⋆ , • • • , s N -1 ⋆, ô⋆ the states computed during the forward pass as prescribed by Eq. ( 17). Then, the gradients of the objective function C := ℓ(ô(s N -1 ⋆ ), y) as defined in the multilevel optimization problem (Eq. ( 9)), where it is assumed that ℓ is differentiable, read:EQUATIONwhereλ 1 ⋆ , • • • , λ N -1⋆ satisfy the following conditions:EQUATIONProof of Lemma A.2. Denoting s := (s 1 , • • • , s N -1 ) ⊤ the state variables of the energy-based blocks, λ := (λ 1 , • • • , λ N -1 ) ⊤ the Lagrangian multipliers associated with each of these variables,W := {θ 1 , ω 1 , • • • , θ N -1 , ω N -1} the energy-based and feedforward parameters and ô(s N -1 ) := F N s N -1 , ω N -1 the logits, the Lagrangian of the multilevel optimization problem as defined in Eq. ( 9) reads:L(s, λ, W ) := ℓ ô(s N -1 ), y + N -1 k=1 λ k ⊤ • ∇ 1 E k (s k , θ k , s k-1 , ω k ), s 0 := x (35)Writing the associated Karush-Kuhn-Tucker (KKT) conditions ∂ 1,2 L(s ⋆ , λ ⋆ , W ) := 0 satisfied by optimal states and Lagrangian multipliers s ⋆ and λ ⋆ , we get : 37)-( 38))EQUATION)) +d W λ ⊤ ⋆ • ∂ 2 L(s ⋆ , λ ⋆ , W ) =0 (Eq. (EQUATIONMore precisely, applying Eq. ( 39) to the feedforward and energy-based block parameters yields:d ω N ℓ(ô ⋆ , y) = ∂ 2 F N (s N -1 ⋆ , ω N ) ⊤ • ∇ 1 ℓ(ô ⋆ , y), d θ k ℓ(ô ⋆ , y) = ∇ 2 1,2 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • λ k ⋆ ∀k = 1 • • • N -1 d ω k ℓ(ô ⋆ , y) = ∇ 2 1,4 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • λ k ⋆ ∀k = 1 • • • N -1Lemma A.3 (Computing Lagrangian multipliers by EP). Under the same hypothesis as Lemma A.2, we define the nudged state of block k, denoted asEQUATION)Defining (δs k ) k=1•••N -1 recursively as:δs N -1 := ∇ s N -1 ℓ(ô ⋆ , y), δs k := d β ∇ 3 E k+1 s k+1 β , θ k+1 , s k ⋆ , ω k+1 β=0 ∀k = 1 • • • N -2,(41) then we have:EQUATIONwhere (λ k ) k=1•••N -1 are the Lagrangian multipliers associated to the multilevel optimization problem defined in Eq. ( 9).Proof of Lemma A.3. We prove this result by backward induction on k.Initialization (k = N -1). By definition, s N -1 β satisfies :β∇ s N -1 ℓ (ô ⋆ , y) + ∇ 1 E N -1 s N -1 β , θ N -1 , s N -2 ⋆ , ω N -1 = 0 (43)Differentiating Eq. ( 43) with respect to β and evaluating the resulting expression at β = 0, we obtain:∇ s N -1 ℓ (ô ⋆ , y) + ∇ 2 1 E N -1 s N -1 ⋆ , θ N -1 , s N -2 ⋆ , ω N -1 • d β s N -1 β | β=0 = 0 (44)Substracting out Eq. ( 37) defining the Lagrangian multiplier λ N -1 ⋆ and Eq. ( 44), we obtain:∇ 2 1 E N -1 s N -1 ⋆ , θ N -1 , s N -2 ⋆ , ω N -1 • d β s N -1 β | β=0 -λ N -1 ⋆ = 0 (45) By invertibility of ∇ 2 1 E N -1 s N -1 ⋆ , θ N -1 , s N -2 ⋆, ω N -1 , we therefore have that:λ N -1 ⋆ = d β s N -1 β | β=0 (46) Backward induction step (k + 1 → k). Let us assume that λ k+1 ⋆ = d β s k+1 β | β=0 . We want to prove that λ k ⋆ = d β s k β | β=0 . Again, s k+1 βsatisfies by definition:EQUATION)On the one hand, proceeding as for the initialization step, differentiating Eq. ( 47) with respect to β and taking β = 0 yields:EQUATION)On the other hand, note that δs k rewrites :EQUATION)where we used at the last step the recursion hypothesis. Therefore combining Eq. ( 48) and Eq. ( 49), we get:EQUATIONFinally, we substract out Eq. ( 38) and Eq. ( 50) to obtain:EQUATION)We conclude again by invertibility of∇ 2 1 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) that λ k ⋆ = d β s k β | β=0 .Theorem A.1 (Formal). Assuming a model of the form Eq. ( 5), we denotes 1 ⋆ , x 1 ⋆ , • • • , s N -1 ⋆, ô⋆ the states computed during the forward pass as prescribed by Alg. 1. We define the nudged state of block k, denoted as s k β , implicitly throughEQUATION)Denoting δs k and ∆x k the error signals computed at the input of the feedforward block F k and of the energy-based block E k respectively, g θ k and g ω k the gradients of the loss function:EQUATION) then the following chain rule applies:δs N -1 := ∇ s N -1 ℓ(ô ⋆ , y), g ω N = ∂ 2 F N s N -1 ⋆ , ω N ⊤ • ∇ 1 ℓ(ô ⋆ , y) (54) ∀k = 1 • • • N -1 :    ∆x k = d β ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0 , g θ k = d β ∇ 2 E k (s k β , θ k , x k ⋆ ) β=0 δs k-1 = ∂ 1 F k s k-1 ⋆ , ω k ⊤ • ∆x k , g ω k = ∂ 2 F k s k-1 ⋆ , ω k ⊤ • ∆x k (55)Proof of Theorem A.1. Combining Lemma A.2 and Lemma A.3, the following chain rule computes loss gradients correctly:δs N -1 := ∇ s N -1 ℓ(ô ⋆ , y), g ω N = ∂ 2 F N s N -1 ⋆ , ω N ⊤ • ∇ 1 ℓ(ô ⋆ , y) (56) ∀k = 1 • • • N -1 :    ∆s k-1 = d β ∇ 3 E k s k β , θ k , s k-1 ⋆ , ω k β=0 , g θ k = ∇ 2 1,2 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • d β s k β | β=0 g ω k = ∇ 2 1,4 E k (s k ⋆ , θ k , s k-1⋆,ω k ) • d β s k β | β=0 (57)Therefore to conclude the proof, we need to show that ∀k = 1, • • • , N -1:d β ∇ 3 E k s k β , θ k , s k-1 ⋆ , ω k β=0 = ∂ 1 F k s k-1 ⋆ , ω k ⊤ • d β ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0 (58) ∇ 2 1,2 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • d β s k β | β=0 = d β ∇ 2 E k (s k β , θ k , x k ⋆ ) β=0 (59) ∇ 2 1,4 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • d β s k β | β=0 = ∂ 2 F k s k-1 ⋆ , ω k ⊤ • d β ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0 (60) Let k ∈ [1, N -1].We prove Eq. ( 58) as:d β ∇ 3 E k s k β , θ k , s k-1 ⋆ , ω k β=0 = d β ∇ s k-1 E k s k β , θ k , F k s k-1 ⋆ , ω k β=0 = ∂ 1 F k s k-1 ⋆ , ω k ⊤ • d β ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0Eq. ( 59) can be obtained as:∇ 2 1,2 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • d β s k β | β=0 = d β ∇ 2 E k (s k β , θ k , s k-1 ⋆ , ω k ) β=0 = d β ∇ 2 E k (s k β , θ k , x k ⋆ ) β=0Finally and similarly, we have:∇ 2 1,4 E k (s k ⋆ , θ k , s k-1 ⋆ , ω k ) • d β s k β | β=0 = d β ∇ 4 E k (s k β , θ k , s k-1 ⋆ , ω k ) β=0 = d β ∇ ω k E k (s k β , θ k , F k s k-1 ⋆ , ω k ) β=0 = d β ∂ 2 F s k-1 ⋆ , ω k ⊤ • ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0 = ∂ 2 F s k-1 ⋆ , ω k ⊤ • d β ∇ 3 E k (s k β , θ k , x k ⋆ ) β=0A.2.2 An alternative proof of Theorem 3.1An energy function for ff-EBMs? While it is clear that the energy function of a ff-EBM is not E = N -1 k=1 E k (which would correspond in this case to the "single block" standard case), one may wonder if:• ff-EBM inference (Alg. 1) can still be described as the minimization of some energy function?• Therefore, if Theorem 3.1 can be derived by directly applying EP to this energy function?We show below that this is indeed the case. We follow Zach [2021] , denoting s := (s1 ⊤ , • • • , s N -1 ⊤ ) ⊤ and W = {W 1 , • • • , W N -1 }, by picking the following energy function:EQUATIONwhere we locally redefine x as the concatenation of all block inputs, i.e.x ←(x ⊤ , s 1 ⊤ ⋆ , • • • , s N -2 ⊤ ⋆) ⊤ , and with s ⋆ := (s1 ⊤ ⋆ , • • • , s N -1 ⊤ ⋆) implicitly defined through ∇ 1 F(s ⋆ , W, x, β = 0) = 0. In Lemma A.4, we show that the free steady state of the above energy function (s ⋆ obtained with β = 0 inside Eq. ( 61)) indeed corresponds to the states computed by the ff-EBM inference scheme (Alg. 1).Lemma A.4. Let E 1 , • • • , E N -1be the block-wise energy functions of a ff-EBM defined per Def. A.1. Assume s ⋆ implicitly defined through ∇ 1 F(s ⋆ , W, β = 0) = 0 where F is defined by Eq. ( 61). Then:EQUATIONProof of Lemma A.4. For k = N -1, the stationarity condition ∇ s N -1 F(s ⋆ , W, x, β) yields:EQUATIONThen, for any1 ≤ k < N -1, ∇ s k F(s ⋆ , W, x, β) = 0 yields: ∇ 1 E k (s k ⋆ , W k , s k-1 ⋆ ) + ∇ 3 E k+1 s k+1 , W k+1 , s k ⋆ -∇ 3 E k+1 s k+1 ⋆ , W k+1 , s k ⋆ =0 = 0 (64)Eq. ( 63) and Eq. ( 63) indeed correspond to ff-EBM inference as depicted inside Alg. 1.The EP fundamental Lemma. For self-containedness of this paper, we restate the fundamental EP result below inside Lemma A.5. Lemma A.5 ( [Scellier, 2021] ). Let F(s, W, x, β) be a twice differentiable function of the three variables s, W and β. For fixed W , x and β, let s β be a point that satisfies the stationarity condition:EQUATIONand suppose that ∇ 2 1 F(s β , W, x, β) is invertible. Then, in the neighborhood of this point, we can define a continuously differentiable function (x, W, β) → s β such that Eq. ( 65) holds for any (x, W, β) in this neighborhood. Furthermore, we have the following identity:EQUATIONIn particular, Eq. ( 66) may be evaluated with F = E + βℓ at β = 0 to yield the EP learning rule, denoting C := ℓ(s ⋆ , y) [Scellier and Bengio, 2017] :EQUATIONFor previous blocks, i.e. k < N -1, we have:EQUATIONAltogether, Eq. ( 73) and Eq. ( 74) finishes to prove Eq. (71).We presented in Alg. 2 a "pure" EP algorithm where the BP-EP gradient chaining is implicit. We show below, inside Alg. 5, an alternative implementation (equivalent in the limit β → 0) where the use of BP through feedforward modules is explicit and which is the direct implementation of Theorem A.1. We also show the resulting algorithm when the ff-EBM reduces to a feedforward net (Lemma A.1) inside Alg. 7, highlight in blue the statements which differ from the general case presented inside Alg. 5.Algorithm 5 Explicit BP-EP gradient chaining (Theorem (3.1))1: δs, g ω N ← ∇ s N -1 ℓ(ô ⋆ , y), ∇ ω N ℓ(ô ⋆ , y) ▷ Single backprop step 2: for k = N -1 • • • 1 do 3: s β ← Optim s E k (s, θ k , x k ⋆ ) + βs ⊤ • δs ▷ EP through E k 4: s -β ← Optim s E k (s, θ k , x k ⋆ ) -βs ⊤ • δs 5: g θ k ← 1 2β ∇ 2 E k (s β , θ k , x k ⋆ ) -∇ 2 E k (s -β , θ k , x k ⋆ ) 6: ∆x ← 1 2β ∇ 3 E k (s β , θ k , x k ⋆ ) -∇ 3 E k (s -β , θ k , x k ⋆ ) 7: g ω k ← ∂ 2 F k s k-1 ⋆ , ω k ⊤ • ∆x ▷ Explicit BP through F k 8: δs ← ∂ 1 F k s k-1 ⋆, ω k ⊤ • ∆x 9: end for A.3.2 Recovering backprop through feedforward nets as a special case Corollary A.1. Under the same hypothesis as Theorem A.1 and Lemma A.1, then the following chain rule applies to compute error signals backward from the output layer:   δs N -1 := ∇ s N -1 ℓ(ô ⋆ , y), g ω N = ∇ ω N ℓ(ô ⋆ , y) ∆x k = σ ′ (x k ⋆ ) ⊙ δs k δs k-1 = ∂ 1 F k s k-1 ⋆ , ω k ⊤ • ∆x k , g ω k = ∂ 2 F k s k-1 ⋆ , ω k ⊤ • ∆x k (75) Proof of Corollary A.1. Let k ∈ [1, N -1].As we can directly apply Theorem A.1 here, proving the result simply boils down to showing that:EQUATIONFirst, we notice that when E k is of the form of Eq. ( 18), then ∆x k reads as:EQUATIONs k β satisfies, by definition and when U k = 0:EQUATIONCombining Eq. ( 77) and Eq. ( 78) yields Eq. ( 76), and therefore, along with Theorem A.1, the chain-rule Eq. ( 75).We showcase in Alg. 6 and Alg. 7 the resulting algorithms implicit and explicit BP-EP chaining respectively, with lines in blue highlighting differences with the general algorithm Alg. 2.Algorithm 6 Implicit BP-EP gradient chaining with U k = 01: δs, g ω N ← ∇ s N -1 ℓ(ô ⋆ , y), ∇ ω N ℓ(ô ⋆ , y) ▷ Single backprop step 2: for k = N -1 • • • 1 do 3: s β , s -β ← σ x k ⋆ -βδs k , σ x k ⋆ + βδs k ▷ EP through E k 4: g ω k ← 1 2β ∇ 4 E k (s β , θ k , s k-1 ⋆ , ω k ) -∇ 4 E k (s -β , θ k , s k-1 ⋆ , ω k ) ▷ i-BP through F k 5: δs ← 1 2β ∇ 3 E k (s β , θ k , s k-1 ⋆ , ω k ) -∇ 3 E k (s -β , θ k , s k-1 ⋆ , ω k ) 6: end for Algorithm 7 Explicit BP-EP gradient chaining with U k = 0 1: δs, g ω N ← ∇ s N -1 ℓ(ô ⋆ , y), ∇ ω N ℓ(ô ⋆ , y) ▷ Single backprop step 2: for k = N -1 • • • 1 do 3: ∆x ← -1 2β σ x k ⋆ -βδs k -σ x k ⋆ + βδs k 4: g ω k ← ∂ 2 F k s k-1 ⋆ , ω k ⊤ • ∆x ▷ Explicit BP through F k 5: δs ← ∂ 1 F k s k-1 ⋆ , ω k ⊤ • ∆x 6: end for A.3.3 Detailed implementation of the implicit BP-EP chaining algorithm (Alg. 2)Nudging the last block. From looking at the procedure prescribed by Theorem 3.1 and algorithms thereof (Alg. 2, Alg. 5), all the error signals used to nudge the EB blocks are stationary, including the top-most block where the loss error signal is fed in. Namely, the augmented energy function of the last block reads as:EQUATIONwhere ô⋆ := F N s N -1 ⋆ , ω N is constant. Up to a constant, Eq. ( 80) uses the cost function linearized around s N -1 ⋆ instead of the cost function itself. This is, however, in contrast with most EP implementations where the nudging force acting upon the EB block is usually elastic, i.e. the nudging depends on the current state of the EB block. More precisely, instead of using Eq. ( 79), we instead use:F N -1 (s N -1 , θ N -1 , x N -1 ⋆ , β) := E N -1 (s N -1 , θ N -1 , x N -1 ⋆ ) + βℓ(F N (s N -1 , ω N ), y), (80)This results in the following asynchronous fixed-point dynamics for the last block:   ∀ odd ℓ ∈ {1, • • • , L k } : s k ℓ,±β,t+ 1 2 ← σ ∇ s k ℓ Φ s k ±β,t , θ k , s k-1 ⋆ , ω k ∓ β∇ s k ℓ(s k ±β,t , y) , ∀ even ℓ ∈ {1, • • • , L k } : s k ℓ,±β,t+1 ← σ ∇ s k ℓ Φ s k ±β,t+ 1 2 , θ k , s k-1 ⋆ , ω k ∓ β∇ s k ℓ(s k ±β,t , y) .The resulting Asynchronous subroutine, applying for the last block, is depicted inside Alg. 8.Readout. Laborieux et al. [2021] introduced the idea of the "readout" whereby the last linear layer computing the loss logits is not part of the EB free block dynamics but simply "reads out" the state of the penultimate block. In all our experiments we use such a readout in combination with the cross entropy loss function. Using our formalism, our readout is simply the last feedforward transformation used inside ℓ, namely F N (•, ω N ).Detailed implicit EP-BP chaining algorithm. We provide a detailed implementation of our algorithm presented in the main (Alg. 2) in Alg. 11. As usually done for EP experiments, we always perform a "free phase" to initialize the block states (Forward subroutine, Alg. 4). Then, two Algorithm 8 Asynchronous (for last block)Input: T , θ N -1 , ω N -1 , ω N , s k-1 ⋆ , β, ℓ (cost function), y Output: s N -1 β 1: s N -1 ← 0 2: for t = 1 • • • T do 3: ∀ odd ℓ ∈ {1, • • • , L N }: 4: s N -1 ℓ,β ← σ ∇ s N -1 ℓ Φ s N -1 β , θ N -1 , s N -2 ⋆ , ω N -1 -β∇ s N -1 ℓ ℓ(F N (s N -1 , ω N ), y) 5: ∀ even ℓ ∈ {1, • • • , L N }: 6: s N -1 ℓ,β ← σ ∇ s N -1 ℓ Φ s N -1 β , θ N -1 , s N -2 ⋆ , ω N -1 -β∇ s N -1 ℓ ℓ(F N (s N -1 , ω N ), y)7: end for nudged phases are applied to the last block and parameter gradients subsequently computed, as done classically (BlockGradient subroutine for the last block, Alg. 9), with an extra computation to compute the error current to be applied to the penultimate block (δs N -2 ). Then, the same procedure is recursed backward through blocks (Alg. 10), until reaching first block.Algorithm 9 BlockGradient (for last block)Input: T , s N -2 ⋆ , θ N -1 , ω N -1 , ω N , β, ℓ, y Output:δs N -2 1: s N -1 β ← Asynchronous T, θ N -1 , ω N -1 , ω N , β, ℓ, y ▷ Alg. 2: s N -1 -β ← Asynchronous T, θ N -1 , ω N -1 , ω N , -β, ℓ, y 3: g ω N ← 1 2 ∇ s N -1 ℓ(F N s N -1 β , ω N ) + ∇ s N -1 ℓ(F N s N -1 -β , ω N ) 4: g θ N -1 ← 1 2β ∇ 2 E N -1 (s N -1 β , θ N -1 , s N -2 ⋆ , ω N -1 ) -∇ 2 E N -1 (s N -1 -β , θ N -1 , s N -2 ⋆ , ω N -2 ) 5: g ω N -1 ← 1 2β ∇ 4 E N -1 (s N -1 β , θ N -1 , s N -2 ⋆ , ω N -1 ) -∇ 4 E N -1 (s N -1 -β , θ N -1 , s N -2 ⋆ , ω N -2 ) 6: δs N -2 ← 1 2β ∇ 3 E N -1 (s N -1 β , θ N -1 , s N -2 ⋆ , ω N -1 ) -∇ 3 E N -1 (s N -1 -β , θ N -1 , s N -2 ⋆ , ω N -2 )Algorithm 10 BlockGradient (for all blocks until penultimate)Input: T , s k-1 ⋆ , θ k , ω k , β, δs Output:δs k-1 1: s k β ← Asynchronous T, θ k , ω k , β, δs ▷ Alg. 2: s k -β ← Asynchronous T, θ k , ω k , -β, δs 3: g θ k ← 1 2β ∇ 2 E k (s k β , θ k , s k-1 ⋆ , ω k ) -∇ 2 E k (s k -β , θ k , s k-1 ⋆ , ω k ) 4: g ω k ← 1 2β ∇ 4 E k (s k β , θ k , s k-1 ⋆ , ω k ) -∇ 4 E k (s k -β , θ k , s k-1 ⋆ , ω k ) 5: δs k-1 ← 1 2β ∇ 3 E k (s k β , θ k , s k-1 ⋆ , ω k ) -∇ 3 E k (s k -β , θ k , s k-1 ⋆ , ω k )Algorithm 11 Detailed implicit BP-EP gradient chaining1: s 1 ⋆ , • • • , s N -1 ⋆ ← Forward (T free , x, W ) ▷ Alg. 2: δs ← BlockGradient T nudge , s N -2 ⋆ , θ N -1 , ω N -1 , ω N , β, ℓ, y ▷ Alg. 3: for k = N -2 • • • 1 do 4: δs ← BlockGradient T nudge , s k-1 ⋆ , θ k , ω k , β, δs▷ Alg. 5: end for A.4 Static gradient analysis Important foreword. The whole subsection is dedicated to an important tool when developing code for EP research. While EP is agnostic to how the steady states are obtained -the EP theory only prescribes they are energy minimizers -they can be obtained in practice (i.e. in simulations) through fixed-point iteration schemes (see Appendix A.1.1). The below formally defines the computational graph spanned by these schemes and abstract them away into a transition function K and defines three different techniques to compute gradients on this graph: Automatic Differentiation (AD, Prop. A.1), Implicit Differentiation (ID, Def. A.3) or Equilibrium Propagation (EP, Def. A.4). After defining each of these algorithms formally, we will state and demonstrate an equivalence between EP and ID (Theorem A.3) which we test numerically and relied upon for the development of our codebase.Definition of the computational graph being optimized. We abstract fixed-point iteration dynamics away into a kernel function K which, given some block state s k t yields s k t+1 . Definition A.2 (Form of the computational graph through equilibrium).EQUATIONNote that we emphasize, through the W k t-1 = W k notation, that the parameters W k are shared across all timesteps t = 1, • • • , τ . This will help us define loss gradient with respect to W k t-1 further below, i.e. how much W k contributes at time t -1 to changing the loss C. The total contribution of W k reads as the sum of the elemental contributions of all W k t . This intuition is more precisely illustrated further below. Given the computational graph defined in Def. A.2, we can now formally define the Automatic Differentiation (AD) baseline.Automatic Differentiation (AD). Our goal is to compute:EQUATIONIn plain words, ĝAD W k (t) denotes the loss gradient for parameter W k truncated at the t th step moving backward in time. We formally define below Automatic Differentiation (AD).Proposition A.1 (Automatic Differentiation (AD)). The gradients ĝAD W k (t) can be computed using the following recursive equations:EQUATIONProof of Prop. A.1. This is a straightforward application of the chain rule applied to Eq. ( 81).Implicit Differentiation (ID). We define the steady state of block k, which we denote s k ⋆ , as the fixed point of Eq. ( 81). With this notation in hand, we can define Implicit Differentiation (ID) in this setting.⋆ the fixed point of Eq. ( 81) inside block k, we define Implicit Differentiation (ID) through the following recursive equations:EQUATIONWe are now ready to state a simple algorithmic equivalence between ID and AD, which we built upon for our implementation of Alg. 12.Corollary A.2 (Equivalence of ID and AD). Assuming that:EQUATIONwhere s k ⋆ denotes the fixed-point of Eq. 81, then automatic differentiation (Prop. A.1) and implicit differentiation (Def. A.3) are equivalent, namely:EQUATIONProof of Corollary A.2. This is a straightforward application of the definition of AD (Prop. A.1 along with the hypothesis made inside Corollary A.2.Resulting implementation of ID. We describe our implementation of ID inside Alg. 12. First, we relax all blocks sequentially to equilibrium following Alg. 4 and we do not track gradients throughout this first phase, using T free fixed-point iteration steps per block. Then, initializing the block states with those computed at the previous step, we re-execute the same procedure (still with Alg. 4), this time tracking gradients and using T nudge steps fixed-point iteration steps for each block. Then, we use automatic differentiation to backpropagate through the last T nudge steps for each block, namely backpropagating, backward in time, through equilibrium.Algorithm 12 Our implementation of ID 1: Without tracking gradients: ▷ e.g. with torch.no_grad()2: s 1 ⋆ , • • • , s N -1 ⋆ ← Forward (T free , x, W ) ▷ Alg. 4 3: Initialize block states at s 1 ⋆ , • • • , s N -1 ⋆ 4: ô⋆ ← Forward (T nudge , x, W )▷ This time gradients are tracked 5: C ← ℓ(ô ⋆ , y) 6: Backpropagate C backward through the last T nudge steps for each block ▷ e.g. C.backward() An important note about this implementation of ID. Note that this is not a standard implementation of ID and it may be surprising at first glance to implement ID as AD, thereby loosing the constant O(1) memory cost of ID with respect to the length of the computational graph. Instead, the memory cost of Alg. 12 is O((N -1)τ )4 . However, our goal is not so much to optimize for memory usage (as in the context of Deep Equilibrium Models [Bai et al., 2019] ) but to code an algorithmic baseline which we know to be equivalent to EP. Lastly, note that this implementation of ID is also known as Recurrent Backprop (RBP, [Almeida, 1987 , Pineda, 1987] ) or Von-Neumann RBP [Liao et al., 2018] , and that ID generally comes in many more algorithmic flavors [Blondel et al., 2022] .W k (t) rewrites:EQUATIONLikewise, we can show that:EQUATIONAltogether, Eq. ( 91), Eq. ( 92) Eq. ( 93) and Eq. (94) yield, denoting ĝEPEQUATIONStarting from k = N -1, (-∆s N -1 t ) t∈ 1,τ and (δs N -1 t ) t∈ 1,τ , (∆x N -1 t ) t∈ 1,τ and (δx N -1 ) t∈ 1,τ , (ĝ EP W N -1 (t)) t∈ 1,τ and (ĝ ID W N -1 (t)) t∈ 1,τ satisfy the same initial conditions and recursive equations, therefore there are all (pair-wise) equal for t = 1, • • • , τ . Therefore in particular, ∆x N -1 τ = δx N -1 τ such that (-∆s N -2 t ) t∈ 1,τ and (δs N -2 t ) t∈ 1,τ from the previous (N -2) th block satisfy the same initial conditions, such that the reasonning applying to k = N -1 recurses for k < N -1, which yields Eq. (88).Description Source ernoult σ(x) = max(min(x, 1), 0) [Ernoult et al., 2019 ] laborieux σ(x) = max(min(0.5 × x, 1), 0) [Laborieux et al., 2021] nest σ(x) = max(min(α × x, 1), 0) This work enhances gradient stability and speeds convergence in deep multi-block settings. In the interest of multi-scale uniformity and consistency with previous literature [Laborieux et al., 2021] Ernoult et al.[2019], we apply clamping activations on all layers in our 6-layer architecture.For the scaling experiments, we apply the "laborieux" activation at every layer except the output of each block. For the 12-layer splitting experiment, we do the same, omitting clamping from the output of the final layer of each block in the block-size=4 and block-size=3 experiments. However, in the block-size=2 case we clamp the output of the second and fourth blocks to preserve dynamics of the block-size=4 split. Such consistency is not possible for the block-size=3 experiment, constituting a possible discrepancy in dynamics.Cross-entropy loss and softmax readout. Following [Laborieux et al., 2021] , all models were implemented such that the output y is removed from the system (e.g. not included in the relaxation dynamics) but is instead the function of a weight matrix: W out of size dim(y) × dim(s), where s is the state of the final layer. For each time-step t, ŷt = softmax(W out s t ).The cross-entropy cost function associated with the softmax readout is then:l(s, y, W out ) = - C c=1 y c log(softmax c (W out • s)).Convention to count layers. It is important to note that by convention we refer to architectures throughout this text to the exclusion of the softmax readout, which is technically an additional layer, though not involved in the relaxation process.Architecture. All convolutional layers used in experiments are of kernel size 3 and stride and padding 1. Max-pooling was applied with a window of 2 × 2 and stride of 2. For the 6-layer model used in Table 1 , batchnorm was applied after the first layer convolution and pooling operation. All other models in both experiments use batch-normalization on the first layer of each block after convolution and pooling (where applied). These details exclude the linear softmax readout of size n classes.Hyperparameters. Detailed hyperparameters for to reproduce Table 1 and Table 2 are given inside the configuration files of our repository. Note that all architectural details for the 12-layer models are identical across splitting and scaling experiments. Additionally, identical hyperparameters were used for CIFAR100 and Imagenet experiments of Table 2 . Unlike previous literature, the use of GOE intialization eliminates the need for separate layerwise learning rates and initialization parameters. One noteworthy detail is that only 100 epochs were used for the larger model for Table 2 compared with 200 epochs for the smaller 12-layer model. This was due to prohibitively long run-time of training the larger model. Noteably, performance still significantly improves with decreased overall runtime.Root-finding algorithms. While in principle any root-finding algorithm may be used for the two relaxation phases of our EP implementation (for inference and gradient computation), our implementation utilized a simple fixed-point iteration procedure, in which neuron states are initialized as zero vectors with values updated asynchronously on each iteration to that of the gradient of the total system energy with respect to current state. An approximate illustration of this procedure is found in Alg. 3. As indicated in Section 4.3, two variants of the convergence procedure were employed, one in which the average value of current state is compared to that of the previous state for each layer, and relaxation is truncated when values for all layers have a difference of less than 1e-4. This was known as the tolerance-based (TOL) procedure. Notably, tolerance-based convergence criteria were applied on the free phase only, with nudging computed with a fixed value of iterations. This was to ensure consistency between ID and EP, though in practice a tolerance can be applied equally to the nudging phase.Algorithm 13 Asynchronous with Tolerance (for all blocks until penultimate)Input: T , θ k , ω k , s k-1 ⋆ , β, δs k Output: s k β 1: s k ← 0 2: c ← ∞ 3: for t = 1 • • • T do 4: ∀ odd ℓ ∈ {1, • • • , L k } : s k ℓ,β,temp ← σ ∇ s k ℓ Φ s k β , θ k , s k-1 ⋆ , ω k -βδs k 5: ∀ even ℓ ∈ {1, • • • , L k } : s k ℓ,β,temp ← σ ∇ s k ℓ Φ s k β , θ k , s k-1 ⋆ , ω k -βδs k 6:if t ≥ 2 then 7: 13: end for Supplementary results with a fixed number of iterations. In addition to the TOL-based procedure, we obtained results for 4.3 using the more conventional approach of [Scellier and Bengio, 2019 ][Laborieux et al., 2021 ][Ernoult et al., 2019] , applying fixed number of iterations on the first and second relaxation phases (see 1). This approach was also the default used for our scaling experiments in 4.4. Importantly, with the TOL procedure described above Alg 3 becomes Alg 13. Results using a fixed iteration root-finding scheme are shown in 5 Other details. All experiments were run using Adam optimizer [Kingma and Ba, 2014] and Cosine Annealing scheduler [Loshchilov and Hutter, 2017] , specifying some minimum learning rates and setting maximum T equal to epochs (i.e. no warm restarts). Code was implemented in Pytorch 2.0 and all simulations were run on NVIDIA A100 SXM4 40GB GPUs.∀ ℓ ∈ {1, • • • , L k } : c k ℓ ← s k ℓ,∇ s N -1 ℓ(ô(s N -1 ⋆ ), y) + ∇ 2 1 E N -1 (s N -1 ⋆ , θ N -1 , s N -2 ⋆ , ω N -1 ) • λ N -1 ⋆ = 0 (37) ∇ 2 1,3 E k+1 s k+1 ⋆ , θ k+1 , s k ⋆ , ω k+1 • λ k+1 ⋆ + ∇ 2 1 E k s k ⋆ , θ k , s k-1 ⋆ , ω k • λ k ⋆ = 0 ∀k = N -2, • • • , 1(38)Eq. (36) governs the bottom-up block-wise relaxation procedure (as depicted in Alg. 1), while Eq. (37) and Eq. (38) governs error propagation in the last block and previous blocks respectively. Given s ⋆ and λ ⋆ by Eq. (36) -Eq. (38), the total derivative of the loss function with respect to the model parameters read:We are not accounting for the spatial depth (L) of the computational graph in this cost. In this case, standard ID would have memory cost O(L) and our implementation inside Alg. 12 O(L(N -1)τ ). 
 
 Appendix:  


Decision: The authors present feedforward tied Energy Based Models (ff-EBMs) a new class of models that combine traditional feedforward layers with EBM layers. The paper also describes an end-to-end training methodology for ff-EBMs that chains the gradient estimations via Backprop (for ff layers) and Eq Prop (for EBM layers). This training methodology is supported by a detailed proof that confirms its correctness and shows that the ff-EBM training methodology works for any arbitrary split of a network into ff/EBM layers.

In the experiments section, the authors apply ff-EBM models for CIFAR10/CIFAR100/ImageNet32 datasets and show competitive accuracy. They also validate the equivalence of the gradients estimated using the proposed training methodology against traditional backprop based gradients and the validity of their training methodology for a range of ff/EBM splits for target model. 

While the paper presents a strong theoretical explanation of ff-EBMs training and inference, many reviewers raised interesting questions about why ff-EBMs should be of interest and what are the benefits of Energy based models. The Authors have provided a detailed response to every question and provided useful figures and explanations of ff-EBMs usefulness and the HW system they have in mind to deploy these models. 

A couple of points that stood out to me in terms of this work’s importance are as follows. Digital AIML HW implementations show high accuracy but suffer from high power cost, Analog HW implementation promises to counter some of this high-power cost. However, the analog HW suffers from reliability concerns and slow weight updates, they aren’t suitable for backprop based training and non-weight stationery operations (e.g. Batch Normalization). While EBMs offer a solution that enables reliable training on analog HW, it still depends on Digital HW for implementing non-weight stationery operations. The proposed ff-EBMs are a pioneering solution to combine the best of both worlds, by using analog HW wherever it enables low energy operations and outsourcing the rest to digital HW. Thanks to the training methodology demonstrated in this paper, any given model can be arbitrarily split onto analog/digital HW. 

Further, EBMs currently suffer from long simulation times due to the delays involved in the root-finding step for each layer of EBMs. By combining ff-EBMs the authors speed up the training steps significantly (~4x as shown in Table 1 of rebuttal pdf). This in turn enables deploying networks with deeper EBM layers interspersed with ff layers to achieve superior accuracy compared to existing EBMs. 

The authors also add a detailed timeline (Figure 3 rebuttal pdf) of their future work indicating extension to other datasets (ImageNet244) and model architectures (Transformer). In conclusion, ff-EBMs and their proposed training algorithm is a pioneering step in the direction of hybrid (digital/analog) HW systems implementing efficient AIML algorithms. This paper should be accepted for NeurIPS.

Edit: Updated my recommendation to spotlight (prev Oral) post discussion with SAC.

Reviews: review: paper_summary: paper_summary: Analog in-memory computing is gaining traction as an energy-efficient platform for deep learning. However, fully analog-based accelerators are challenging to construct, necessitating a training solution for digital-analog hybrid accelerators. This paper introduces Feedforward-tied Energy-based Models (ff-EBMs), a hybrid model that integrates feedforward components, typical in digital systems, with energy-based blocks suitable for analog circuits. An algorithm is derived to compute gradients end-to-end for training ff-EBMs. Experimental results show that ff-EBMs achieve superior accuracy on ImageNet32 classification compared to conventional equilibrium propagation literature.
strength_weakness: strength_weakness: 1. The paper lacks examples of situations where ff-EBM is needed. It does not adequately explain which parts of the overall system are digital and which are analog when using ff-EBM. Additionally, it does not clarify whether the existing training techniques can be used in these scenarios or what advantages ff-EBM offers compared to traditional methods.
questions: questions: As I understand it, analog in-memory computing naturally finds the node voltage that minimizes the energy function through Kirchhoff's laws, classifying it as an EP problem. Therefore, if we disregard the natural minimization of the energy function, the training process with analog in-memory computing appears very similar to the conventional backpropagation-based training procedure. Hence, the derived algorithm for calculating end-to-end gradients of the ff-EBM seems trivial. In this context, I don't see the difference between this research and previous studies that have conducted training on analog in-memory systems using EP or backpropagation. Could you explain the differences between the previous approaches and the proposed approach in more detail? Providing examples of situations where ff-EBM is needed, as discussed in the Weakness section of this review, would be helpful.
limitations: limitations: Please check the weakness.
paper_summary: paper_summary: This paper presents Feedforward-tied Energy-based Models (ff-EBMs), a hybrid model that integrates feedforward and energy-based components, accounting for both digital and analog circuits. A novel algorithm is proposed to compute gradients end-to-end in ff-EBMs by backpropagating and "eq-propagating" through feedforward and energy-based sections, respectively, allowing EP to be applied to more flexible and realistic architectures. It has been shown that ff-EBMs can be trained on ImageNet32, achieving new state-of-the-art performance in the EP literature with a top-1 accuracy of 46%.
strength_weakness: strength_weakness: -- The primary limitation of this work is its accuracy performance on the ImageNet dataset. Although the paper has set a new state-of-the-art accuracy, it still lags behind the results achieved by Transformers and CNNs. There is a lack of compelling reasons to use this method given its comparatively lower performance.

-- Additionally, it is crucial to measure the energy consumption of this training method and compare it with traditional methods. How much energy savings does your method offer?
questions: questions: See weaknesses.
limitations: limitations: NA
paper_summary: paper_summary: This paper proposes a new building model block with analog forward circuits and energy-based blocks built on a digital-analog hybrid setup.  A novel algorithm is further proposed to train the new block model. Experiments show the SOTA accuracy in the EP literature.
strength_weakness: strength_weakness: * The paper's motivation is vague to me, as I am new to this topic. Why must we combine an analog forward circuit with a digital one for EP-based training? For the analog accelerator, we can use several other methods to do training, e.g., forward-forward only and zeroth-order optimization. What are the benefits of incorporating an energy-based model?
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: This paper discusses a novel approach to improving the power efficiency of AI training by integrating analog and digital hardware. The authors introduce a hybrid model called Feedforward-tied Energy-based Models (ff-EBMs) that combines feedforward neural networks with energy-based models (EBMs). This model aims to leverage the benefits of both analog and digital systems to create a more efficient computational framework for neural network training.

The paper introduces a new model, ff-EBMs, which integrates feedforward and EB modules. These models are designed to perform end-to-end gradient computation, combining backpropagation through feedforward blocks and "eq-propagation" through EB blocks. And bilevel and multi-level optimization are concepts used to frame the learning process of the Feedforward-tied Energy-based Models (ff-EBMs).

The paper promises to demonstrate the effectiveness of the proposed approach experimentally, using DHNs as EB blocks and achieving new state-of-the-art performance on the ImageNet32 dataset.
strength_weakness: strength_weakness: While the paper introduces an innovative approach to hybrid digital-analog neural network training, there are a few areas where it shows some limitations. 
- The paper could benefit from clearer articulation regarding the energy efficiency claims, providing more detailed comparisons with existing systems to substantiate these claims. 
- Additionally, the reliance on simple datasets, ImageNet32, for experimental validation, while common, might limit the generalizability of the findings. A broader range of datasets with varying characteristics could strengthen the paper's conclusions. The choice of ImageNet32 should be justified with respect to its relevance to the research goals and its limitations for testing the model's robustness.
- Lastly, while the authors acknowledge the need for further research, a more explicit discussion on the current limitations and a roadmap for future work would provide a more comprehensive view of the research's trajectory.
questions: questions: In the proposed implicit BP-EP chaining algorithm, is it necessary to ensure the satisfaction of the first-order optimality conditions derived from the implicit theory, or can the algorithm be robust and effective without strictly guaranteeing these conditions?
limitations: limitations: A notable limitation of our study is the constrained scope of network architectures and datasets utilized for validation. While the implicit BP-EP chaining algorithm has shown success with Deep Hopfield Networks on ImageNet32, its performance on the standard ImageNet and transformer-based models, which are more complex and widely used, is yet to be established.
paper_summary: paper_summary: State-of-the-art (SOTA) analog hardware accelerators consist of both analog and digital components supporting major and auxiliary operations. Moreover, they typically suffer from device imperfections on their analog parts. In this paper, the authors propose feedforward-tied energy-based models (ff-EBMs) for digital and analog circuits. ff-EBMs compute gradients by both backpropagating on feedforward parts and eq-propagating on energy-based parts. In this paper, ff-EBMs use a Deep Hopfield Network (DHN) as one example, which can be arbitrarily partitioned into uniform size. ff-EBMs achieves a SOTA performance on ImageNet32.
strength_weakness: strength_weakness: 1. The paper does NOT consider a detailed analog computing device model. It will be interesting to see how significant variations on the analog devices the ff-EBMs can tolerate. And what is the relationship between the convergence speed and the device variations?

2. The paper does not include energy-related data. What is the energy saving for the energy-based models? What is the energy saving of skipping the accurate gradient computations?
questions: questions: Please comment on the points in the weakness section.
limitations: limitations: no limitations.

Comments: title: Respond to rebuttal.
Source of improvement & hardware implementation of our algorithm
Clarifying Scellier-Bengio's EP paper & scalability of ff-EBMs
Thank you for your rebuttal
Answering additional questions on the novelty of the paper
Additional questions on the novelty of this paper.
comment: Thank you for your reply. The reviewer has no additional questions.
Dear Reviewer 9JF8, thank you so much for taking the time to engage in a discussion with us and for raising our score! Your remarks will greatly help improve our manuscript for readers interested in ZO optimization. Thank you again!
Thank you so much for your rebuttal and for successfully addressing my concerns.
I am happy to raise my rating!
We are glad that Reviewer 9JF8 appreciated our rebuttal, and are indeed happy to clarify their questions. 

- *you propose to make the analog part to an energy model and set the digital part to the normal feedforward layer*. Yes: analog and digital parts account for energy-based and feedforward models respectively.

- *you derive the automatic differentiation method to train this hybrid setup*. To be certain to be on the same page as Reviewer 9JF8, we slightly reformulate this sentence:
  + we apply the *Lagrangian method* to derive optimality conditions, i.e. “KKT” conditions, to “train this hybrid setup” (Appendix A.2).
  + the application of this method yields an algorithm which is itself a **hybrid** differentiation method which uses standard backprop (i.e. “automatic differentiation”) inside feedforward / digital parts and **equilibrium propagation inside energy-based / analog parts**.

Therefore, our algorithm isn’t “pure” automatic differentiation end-to-end, but *only within feedforward parts of the model*.

- *Which is the main source of the accuracy improvement?* The main source of accuracy improvement, *with respect to the past EP works*, is simply the *depth* of the model trained: the ff-EBMs trained here are twice as deep as the largest EBMs trained in previous EP works. Reviewer 9JF8’s question may then translate to: why can you train deeper networks? The most important reason, which is highlighted in Table 1 of our PDF and in our rebuttal to k7Rs, is that **ff-EBMs are faster to simulate than their fully EBM counterparts** (up to 4 times faster per Table 1). Instead of the superlinear scaling of the convergence time of DHNs (with respect to the number of layers) empirically observed in past EP works, the convergence time ff-EBMs made up of DHNs are EB blocks are guaranteed, by construction, to scale linearly with the number of blocks, the convergence time of a single block decreasing with its size. On the algorithmic side: since i) our algorithm performs on par with end-to-end automatic differentiation on the same ff-EBMs on the one hand, and that ii) ZO techniques are expected to perform less well than automatic differentiation on equivalent architectures (as highlighted in our rebuttal to Reviewer 9JF8), we can conjecture that, on *equal ff-EBMs*, our algorithm may likely perform better than ZO *when applied end-to-end*.

- *“Is your proposed method hardware-friendly? Please discuss the hardware implementation of your method”*. First, as highlighted in our rebuttal, our method extends EP-based training to scaled hardware systems which: i) may not fit a single analog core, ii) may still require operations which cannot be supported on analog hardware and may instead require digital, high precision hardware. Therefore, Fig. 1 of the PDF attached to our rebuttal is a plausible depiction of the hardware implementation of our method and taking these constraints into account, our algorithm is **hardware plausible**. Second, we insist that EP-based training inside each of the analog / energy-based block inherits the “hardware-friendly” / FF-like features of EP training: in our algorithm, **gradients inside analog / energy-based blocks are computed using only forward passes / relaxations to equilibrium**. A plausible hardware implementation of this analog blocks are *deep resistive networks* [1, 2]. Lastly, feedforward blocks, which are maintained in digital, could also directly leverage quantization algorithms [3] and even **ZO algorithms** (as mentioned in L.290-291 of our paper) to facilitate their implementation on memory-constrained, low-power hardware.
 
To summarize: our algorithm makes the best of analog and digital worlds, by: i) being hardware plausible at the system level (Fig .1 of PDF), ii) preserving the “hardware-friendliness” of FF-like learning inside EB/analog blocks, iii) possessing the ability to leverage any quantization algorithm inside feedforward blocks and iv) possessing the ability **to apply ZO algorithms** instead of backprop (as currently done with our algorithm) **inside feedforward blocks**.

Finally, we would like to thank reviewer 9JF8 for drawing our attention to some critical points where our paper could better communicate some of the high-level motivations and details of our algorithm. In addition to what was proposed in our original rebuttal, we would like to propose adding a pseudo algorithm in the appendix showing how ZO could be applied within feedforward blocks (instead of backprop currently) such that gradients would be computed **everywhere with forward passes only**, i.e. both inside analog blocks (by EP) and feedforward blocks (by ZO) (as first suggested in L. 290-291 of our paper).

[1] Kendall et al (2020). Training end-to-end analog neural networks with equilibrium propagation.

[2] Scellier, B. (2024). A Fast Algorithm to Simulate Nonlinear Resistive Networks. ICML 2024.

[3] Lin et al (2022). On-device training under 256kb memory. NeurIPS 2022
We are very grateful to Reviewer Jt5z for spending time to understand the value of our paper, subsequently increasing our score and engaging in this discussion which will tremendously benefit the presentation of our work. Thank you so much!
 
In the light of Reviewer’s Jt5z last answer, we would like to clarify some essential points they raised:

-  *"EP is designed to be applicable to any network architecture, including those with feedforward blocks"*. As mentioned in Section 2.3 of our paper, **EP only applies to energy-based models**. Fig. 1 of the seminal EP paper [1] is indeed misleading: when writing “Equilibrium Propagation applies to any architecture”, one should understand “any architecture **topology** so long as it derives from an energy function”. Indeed, as indicated by the title of the section 3 of this paper, EP really is a “Machine Learning Framework **for Energy-Based models**”: their Figure 1 is only meant to emphasize that EP applies to *any* energy-based models, not necessarily *layered* energy-based models. In this context: **layered does not meant feedforward**. As Scellier & Bengio write themselves: “*In particular, the [EP learning rule] holds for any architecture and **not just a layered architecture** (Figure 1) like the one considered by Bengio and Fischer (2015)* [which is also an EB model]”.

-  *"Therefore, I still view Eq. (8) as a straightforward integration of FF and EBM"*. Given the clarification above, this conclusion may no longer hold, especially when noticing that Eqs.17-18 of the seminal EP paper [1] (in the section 3 mentioned by Reviewer Jt5z), i.e. the bilevel program the EP algorithm solves, **is an explicit particular case of the Eq. 8 of our paper**, i.e. the multilevel program that our algorithm solves.

-  *"I think that when adopting ff-EBM, the key concern should be scalability"*. We totally agree! Investigating the scalability of ff-EBM training by our algorithm on deeper ff-EBMs, more complex tasks and exploring new datasets and architectures is part of our research roadmap. See our detailed answer to Reviewer kDhg and associated Fig. 3 inside the PDF.

- *"The scalability of ff-EBM, as compared to using EBM alone, is a critical point that should be thoroughly explored"*. We are happy Reviewer Jt5z mentions the importance of this comparison since our “splitting experiment” (Section 4.3) goes exactly in this direction. This experiment reveals that a single EBM block performs comparably to an ff-EBM **of equal depth** with various block sizes. 

[1] Scellier, B., & Bengio, Y. (2017). Equilibrium propagation: Bridging the gap between energy-based models and backpropagation. Frontiers in computational neuroscience, 11, 24.
Thank you for your rebuttal.

First, thank you for your new Figure 2, which makes your problem statement much clearer to me, while your previous introduction is too complex and hard to get your point. So, if I understand correctly, **you propose to make the analog part to an energy model and set the digital part to the normal feedforward layer. Then, you derive the automatic differentiation method to train this hybrid setup**.

So, I have several key follow-up questions to help me decide whether to raise/lower my score.
* Which is the main source of the accuracy improvement? The hybrid setup is where the digital feedforward improves accuracy, or the beyond zeroth-order backpropagation algorithm is used.
* The reason I ask ZO/FF is not due to the accuracy; it is because ZO/FF is hardware-friendly for analog circuits where you cannot use backpropagation to do on-chip training. Therefore, ZO/FF is practical for analog hardware requiring forward passes. Is your proposed method also friendly? How can you obtain the right gradient through the analog circuit? Please discuss the real hardware implementation for your method, especially how to obtain the gradient.

thank you so much.
Thank you for your detailed explanation.

Your response has helped me better understand the proposed work, and I have made every effort to assess its value.

Firstly, I am increasing my score to 5, as the novelty of your work is now clear to me. From my understanding, EP is designed to be applicable to any network architecture, including those with feedforward blocks, as the minimization of energy can be aligned with the minimization of the objective function (as shown in Figure 1 and Chapter 3 of the BP paper [1]). Therefore, I still view Eq. (8) as a straightforward integration of FF and EBM. However, I acknowledge that the gradient calculation starting from Eq. (8) is not trivial, and the detailed derivation of these gradients is a key novelty of your work.

While I believe this work is highly significant in the field of analog-based AI accelerator systems, I am hesitant to raise my score further because the presentation of your work could be improved to better highlight its true value.

In my opinion, the design of networks for analog-based computing is heavily constrained by hardware considerations, as there are significant challenges in scaling fully analog systems. Analog computing is known for its energy efficiency compared to digital computing, while digital computing offers greater scalability. To build an efficient yet scalable AI acceleration system, a hybrid design is essential. In this context, I believe the true value of your work lies in extending the scalability of EP-based models for analog computing by integrating FF modules. For this reason, I think that when adopting ff-EBM, the key concern should be scalability. The scalability of ff-EBM, as compared to using EBM alone, is a critical point that should be thoroughly explored. For example, with EBM alone, only a single analog-based unit in Figure 1 of attached PDF could be used for a single network. Obviously, ff-EBM should be much better than EBM, but I believe this paper needs detailed discussion on the accuracy of this EBM-based model and how it compares to the accuracy of a larger model designed with ff-EBM that fully utilizes the entire system depicted in Figure 1 of attached PDF, as this is an important aspect of this work.

[1] B. Scellier and Y. Bengio. Equilibrium propagation: Bridging the gap between energy-based models and backpropagation. Frontiers in computational neuroscience, 11:24, 2017.
Dear Reviewer Jt5z,
 
We thank you very much for engaging promptly in this discussion, and are happy indeed to read that the PDF file brought some satisfying clarifications about the relevance of the problem tackled.
 
If we understand you correctly, since BP and EP are well established algorithms in their own right, chaining them inside a given architecture appears straightforward “by following the chain rule”. In fact, we are very pleased that you have drawn attention to this issue, as we are increasingly recognizing that some of the terminology we use could prove misleading, and to some extent confusing for readers. Therefore, as requested, we highlight below the specific “algorithmic challenges” pertaining to this chaining, “how [they] are addressed” and more broadly the “novelty of this paper”.

- **EP-BP chaining is intuitive, but not trivial to derive**. You write that chaining EP and BP inside a given architecture “appears straightforward by following the chain rule”. We would like to focus attention on this question: what is meant by “following the chain-rule” in the context of ff-EBMs?

  + Traditionally, “chain-rule” refers to gradient chaining inside *feedforward models*. Namely, let us assume a computational graph of the form $s^1 = F(x) \to s^2 = G(s^1)$. Assuming an error signal $\delta^2 = \partial_{s^2} L$, then the “chain rule” prescribes that the error signal at $x$ reads $\partial_{x}L = \partial_x F(x)^\top \cdot \partial_{s^1}G(s^1)^\top \cdot \delta^2$.
 
  + Now let us assume instead the following “hybrid” computational path: $s^1 = F(x) \to s^2: \nabla_{s^2}E(s^2, s^1) = 0$. Given that $s^2$ in this case is an *implicit* function of $s^1$, i.e. there is no explicit G mapping as before between $s^1$ and $s^2$, how would one directly apply the above “chain-rule” here? To put it differently, how do we rigorously route error signals backward through this computational graph? This is the first “challenge” we addressed. 

  + Hence the need to derive gradient chaining inside ff-EBMs *rigorously, from first principles* by: 1) stating the learning problem as a multilevel constrained optimization problem (Eq. 8), 2) writing the associated Lagrangian (Eq. 21), 3) solving for the associated KKT conditions (Eqs. 22-34) for the primal variables (i.e. the steady states of the blocks) and associated Lagrangian multipliers (i.e. the error signals inside blocks). Therefore, this is how “we addressed” the above challenge.

We emphasize that **this derivation** (namely Theorem 3.1 about the rigorous chaining of EP and BP gradients) and the **resulting explicit and implicit chaining algorithms** (Alg. 4 inside Appendix A.3, Alg. 2 in the main) **are novel**. Also, note from the above that our algorithm comes into an *explicit* and *implicit* variant (see paragraph “Proposed algorithm: implicit BP-EP chaining” L.187 and Lemma A.4), the latter appearing as a “pure” EP implementation. As such, the fact that EP-BP chaining can be cast into a *non-trivial generalization of EP*, as appearing in Alg. 2 and Lemma A.4, is not self-evident neither.
 
- **The experimental demonstration of this algorithm is also novel**. Finally, to further highlight the “novelty of this paper”, **our algorithm was never tested in practice before this work** and having it succeed on ImageNet32 came with lots of “challenges” as well, the most important one being the simulation time. We addressed this problem in two ways:

  + As explained in the global rebuttal and in greater details inside Reviewer k7Rs’s rebuttal, splitting an EBM into several EBM blocks tied by feedforward modules results in an architecture that is not only more hardware realistic, but also **easier to simulate**. Indeed, instead of the superlinear scaling of the simulation time with respect to the number of layers observed in past EP works, our approach guarantees, by construction, that this **simulation time scales linearly with the number of blocks**, each of these blocks converging much faster than the full EBM counterpart. See our new table of results inside our PDF attached to this rebuttal. While the ff-EBMs trained are still relatively shallow, **they are twice as deeper as the deepest EBM trained by EP** in most recent related works.

   + Finally, using *Gaussian Orthogonal Ensembles* (GOE) to initialize weights inside EB blocks was instrumental in having ff-EBM training experiments work.

Finally, **we would like to propose the addition of a few sentences to our introduction highlighting this novelty**, and particularly the degree to which our algorithm is derived by exploiting an intimate theoretical connection between energy based learning and **implicit differentiation**, rather than literal "backprop" as applied in standard feedforward nets where the standard "chain rule" applies. Hence our claim that this contribution belongs to the realm of EP, and *implicit learning* more broadly. Again, we thank you for drawing our attention to this.
Thank you very much for your careful and detailed response.

Figure 1 and Figure 2 in the attached PDF file clearly illustrate the hardware set-up and the proposed layer design. I acknowledge that this paper is pioneering in addressing the training algorithm for hardware composed of multiple analog chips, and I agree that this is a realistic configuration.

However, I still have a few questions regarding the novelty of this paper. The ff-EBM model architecture comprises a chain of the FF module and the EB module. Given that the training algorithms for the FF module and EB module are established conventions (backpropagation and equilibrium propagation, respectively), I find it challenging to identify the difficulty in designing a training algorithm for ff-EBM. Since the training algorithms for the FF and EB modules are predetermined, the primary task for ff-EBM is to ensure the proper passing of gradients between the FF and EB modules. Nonetheless, passing gradients between these modules appears straightforward by following the chain rule.

Therefore, I would like the authors to highlight any algorithmic challenges related to configuring the FF and EB modules within a single network architecture. More specifically, I would appreciate it if the authors could provide more detailed explanations on the challenges involved in passing gradients between the FF and EB modules and how these challenges are addressed.