A cookbook for hardware-friendly implicit learning on static data

Venue: NeurIPS_cc_2024_Workshop_MLNCP
Published: 2024-09-12 13:53:27

Abstract: The following aims to be a pragmatic introduction to hardware-friendly learning of implicit models, which encompass a broad class of models from feedforward nets to physical systems, taking static data as inputs. Starting from first principles, we present a minimal hierarchy of independent concepts to circumvent some problems inherent to the hardware implementation of standard differentiation. This way, we avoid entangling essential ingredients with arbitrary design choices by naively listing existing algorithms and instead propose the draft of a “cookbook” to help the exploration of many possible combinations of these independent mechanisms.

Main paper: 
 other tricks which independently apply: 
 Greedy learning & loss design. It appears from Eq. ( 4) that in general, an error signal λ k+1 must be passed backwards to "unlock" the computation of λ k . In order to parallelize learning across blocks entirely, a heuristic consists in shutting off the top-down error signal (λ k+1 = 0) and recreate an error signal through a locally-defined (supervised (8; 71; 9; 88; 27) or self-supervised (50; 91; 31; 83)) loss, which amounts to choosingO = O 1 (s 1 , θ 1 ) + • • • + O K (s K , θ K ).In this case, the resulting block satisfies the exact same adjoint equation as that of the original learning problem (Eq. ( 2)) so that all gradient computation techniques herein presented may also apply block-wise. Alternatively, another solution to backward locking is to estimate λ k+1 with auxiliary modules (38) .Checkpointing & reversible models. One way to mitigate the memory problem mentioned above (at the expense of compute) is to simultaneously compute s and λ, which can be viewed as activation checkpointing (17) . In models with an explicit layer hierarchy (P 2 ), another special instantiation of activation checkpointing is possibly when using reversible models (19; 26; 16; 52): a given constraint C k+1 (s k+1 , s k ) = 0 can be explicitely inverted as C -1 k+1 (s k , s k+1 ) = 0 such that s k can be recomputed backward from layer s k+1 instead of being stored, or recomputed forward from s 0 = x.This can be achieved for instance by splitting each block state s k into two,s k = (s k ⊤ a , s k ⊤ b ) ⊤ , and defining C k+1 = (C a k+1 , C b k+1), with dedicated transformations f a , f b , as (26) :EQUATIONPipelining. When dealing with a hierarchical model (P 2 ), the block C k may become idle or "forward-locked" after passing s k to C k+1 until next input comes in. A solution to this is to push multiple inputs in sequence through the blocks, allowing them to process different inputs in parallel, e.g. C k processes input x p while C k+1 processes input x p+1 , the same strategy applying backwards for the computation of the λ k 's for different inputs (37) . As naive pipelining may still maintain idleness "bubbles", more elaborate schemes have been proposed (22) , for instance by allowing each block to alternate between forward pass and backward passes for different inputs ( 67), yet at the cost of gradient stalenessλ is computed with a different θ that the one used to compute s. This problem can be mitigated for instances by by maintaining different weight versions at each circuit (68).3 Forward-only learning beyond zeroth order s ±β : ∇ s F(±β, s ±β , θ) = 0, g(θ) = 1 2β (∇ θ F(β, s β , θ) -∇ θ F(-β, s -β , θ)) + O(β 2 ) (8)Eq. ( 8) is in stark contrast with Eq. ( 2), though being mathematically equivalent through λ = d β s β | β=0 (78; 69): instead of estimating s and λ on two circuits, s β and s -β are estimated on a single circuit through "energy minimization", which is why C is said to be energy based (EB).The core intuition behind the magic of EB learning is that while error signals are usually carried "backward" through ∂ s C ⊤ (Eq. ( 2)), since we have∂ s C ⊤ = ∇ 2 s K = ∂ s C, error signals can in thiscase be equivalently carried "forward" through ∂ s C through a small perturbation of s along ∇ s O of sufficiently small β. Finally note that Eq. ( 8) is only one many variants (80) to estimate d β (∇ θ F)| β=0 (78) , where there is trade-off between the number of s β being evaluated and the resulting bias (94; 42). 
the importance of nudging: 
 The weak & strong nudging limits. The most current EBL setting is when C describes the implicit model itself and O the cost function: (C, O) = (F = ∇ s E, ℓ) where we have denoted K = E. In this case, the condition on the nudged state reads (78):EQUATIONWe call the nudging factor the scalar controlling the strength of the cost function ℓ in the definition of s β . For this choice of (C, O), β is the nudging factor and since theoretical guarantees of Eq. ( 8) hold for β → 0, it corresponds to small ∇ s ℓ perturbations, which is therefore called the weak nudging limit. Conversely, when swapping the objective and the constraint (59; 60), i.e. (C, O) = (∇ s ℓ, E), which amounts to swap E and ℓ inside Eq. ( 9), one can see that β -1 now gates the error signal ∇ s ℓ.Therefore in this situation, β → 0 corresponds to large ∇ s ℓ perturbations, which is therefore called the strong nudging limit. While it seems counter-intuitive that strong nudging solutions are relevantthe main goal of of Eq. ( 1) should be to minimize the loss subject to constraints on the network energy and not the other way around -they are global minimizers of θ → ℓ(θ, s(θ)) for certain choices of E (60). Strong nudging can enable learning in situations where the physical system at use is noisy by improving the signal to noise ratio for the teaching signal (59; 42).Nudging through adiabatic oscillations. EBL classically operates in the weak nudging regime and estimates the error signal λ = d β s β | β=0 through a discrete, two-steps finite difference procedure, as highlighted in Eq. ( 8). Another way to view the error signal is as the contour integration in the complex plane λ = 1T 0 e i2πt/T s β(t) with β(t) := |β|e i2π/T . Then, provided T is sufficiently large compared to the characteristic time constant which governs the relaxation of C to equilibrium, the error signal can be computed through slow adiabatic oscillations of the system (5; 42; 3).Nudging heuristics. The error signal β∇ s ℓ may be transmitted instantaneously to the rest of the network, or with a delay. Denoting h(θ, s) the model logits, some output controller dynamics u may integrate the error signal inside the output layer as τ u u + u = -β∇ o ℓ(h(s, θ)) and feed it back to the rest of the circuit (60) . There also exists other types of nudging among the broader family of Contrastive Learning (CL) algorithms. Denoting y some label, the "nudged" hidden state h y can be defined by hard-clamping y to the output units, ∇ s E(h y , o = y, θ) = 0, while prescribing a weight update of the same contrastive form as Eq. 8 (66) . An hybrid between weak nudging and target clamping consists of clamping the output to a weakly nudged target, resulting in a nudged hidden state h β defined as:∇ s E(h β , o = (1 -β)o ⋆ + βy) = 0where o ⋆ corresponds to the free state (β = 0) of the output (86) .Nudging synapses. One may wonder whether neurons and parameters could play a symmetrical role, by having them both evolve throughout the gradient computation phase or satisfy an equilibrium condition. One common post hoc solution is to play on characteristic time constants by having parameters slowly integrate the instantaneous contrastive updates prescribed by Eq. ( 8) (20; 87; 42; 18) , which however requires the detailed knowledge of the energy function at use. Another solution circumventing this constraint is to let both s and θ equilibrate in the same energy landscape, with control variables u acting upon θ with strength α (79). When computing s -β (for instance), u is adjusted to keep θ at its current value θ t such that u t = θ t + α∇ s F(-β, s -β , θ t ). Therefore, keeping u t constant when computing s β , θ equilibrates at: θ t+1 = u t -α∇ s F(β, s β , θ t+1 ) ≈ θ t -αβg(θ).Lagrangian reparametrization. One may usually regard the contrastive update Eq. ( 8) as a procedure performed on the same physical system whose state is measured at two different times with two different nudging values (e.g. -β, β), or a continuum of such. Yet, another view of Eq. ( 8) is to assume that the two nudged states are computed simultaneoulsy by two different circuits sharing the same parameters. An advantage of this implementation, which has been experimentally demonstrated with transistor-based synapses (18) , is that it requires a single phase only instead of two, yet at the expense of area and complex engineer to enable parameter sharing. A different way to approach the same question is through the concept of dyadic neurons (34; 35): two variables s + and s -are such that their (convex) sum encodes the model prediction and their difference the error signal if they are critical points of the following reparametrized Lagrangian (35) :EQUATIONwhere L defined inside Eq. ( 2) and α ∈ [0, 1]. While the resulting gradient computation algorithm is generally a reparametrization of implicit differentiation by construction, when the constraint at use inside Eq. ( 1) is energy based and β is sufficiently small, s ± can be simply construed as s ±β of Eq. ( 8). This can be shown by performing this reparametrization on the relaxed Lagrangian L 3 (Eq. ( 5)) with fixed β (35), or equivalently approximated from Eq. ( 10) in the limit β → 0.Weak nudging. One of the ways to cast an implicit model F , which does not explicitly derive from an energy function, into an EB model is simply to employ the energy functionK = E = 1 2 ∥F ∥ 2 (60).An important case, as one means to train feedforward nets by EP, is whenF (s, θ) = s -f (θ, s)where θ is typically lower block diagonal (25; 89; 62; 64) . Eq. ( 8) for some β, in the weak nudging regime ((O, K) = (ℓ, E)), yields in this case:EQUATIONA nice feature of Eq. ( 11) is that by construction, s β = f (s β , θ) when β = 0 such that the gradient can be estimated up to O(β) with a single nudged state. However, the presence of ∂ s f (s β , θ) ⊤ signals the potential use analytical derivatives. To obviate this, one may instead estimate the required derivatives through finite differences of feedback operators (46; 76; 57; 58; 21; 56) as:s β ≈ f (s β , θ) + g(s β , ω) -g(f (s β ), ω)) -β∇ s ℓ(s β , θ)where the feedback parameters may have to be learned so that ∂ s g ≈ ∂ s f ⊤ (21). However a problem remaining is that both s β and their "feedforward" prediction f (s β ) need to be simultaneously maintained to implement Eq. (11) . A solution is to use auxiliary variables ψ which track the feedforward activity such that ψ β ≈ f ψ (s β , θ)EQUATION.Recovering a least control problem in the strong nudging regime. When in the strong nudging regime, the optimization problem Eq. ( 1) using E = 1 2 ∥F ∥ 2 = O along with ℓ = K acquires a very intriguing meaning, though the loss is not picked as the objective. Defining ψ β := -F (s β , θ), the learning problem amounts to find the set of parameters which minimizes the amount of control O = 1 2 ∥ψ β ∥ 2 such that the controlled equilibrium ψ β + F (s β , θ) = 0 minimizes the loss ℓ, which works in practice on feedforward and implicit models more broadly (60) . However, this interpretation is tied to the choice E = 1 2 ∥F ∥ 2 and further investigation is needed to extend the applicability of strong nudging and its connection to control theory beyond this choice of energy function.When not handled cautiously, the application of Eq. 8 to implicit models yields activation derivatives (Eq. ( 11)). While neuroscientists, seeking for biologically plausible learning theories, gave birth to the aforementioned approximations of Eq. ( 11) and many others, these remain impractical for hardware-friendly learning. One ad hoc solution used in practice, when the inverse of the activation function is continuously invertible, is to separate the linear and nonlinear contribution inside the energy function as E = G + U where G is defined such that ∇ s G = σ -1 (36; 92; 34; 80) . A more principled solution to handle σ is to instead treat it as a constraint C σ on the set of feasible configurations over which energy minimization is performed, which descent steps are hard projected onto (77) . A "relaxed" version of this approach, somewhat bridging with the former solution, is through the lens of mirror descent using the Bregman divergence associated with G (2):s t+1 = arg min{s ⊤ • ∇ s E(s) + 1 η ∆ G (s, s t )} with ∆ G (x, y) := G(x) -G(y) + ∇G(y) ⊤ • (x -y).Note that many other fixed point algorithms could be use for the same purpose (13).An important example. Another way to cast feedforward nets training by EBL algorithms is to break them into multiple, hierarchical energy-based constraints based on Eq. (3) (92). The following example, which shows how to do so, is also a good illustration of why the problems P 1 (Eq. ( 1))and P 2 (Eq. ( 3)) are generally not equivalent. Taking E :=K k=1 E k with E k (s k , s k-1 , θ k ) := G(s k ) -s k ⊤ • θ k • s k-1 , we consider P 1 with C = ∇ s E and P 2 with C k = ∇ s k E k . In this case, s kis a single layer and not a block of several layers. Setting the constraint to zero in these two cases yields:EQUATIONIn one case, the model retained at inference time is a Hopfield model and in the other one a feedforward model. This difference is instrumental to understand why works revolving around "contrastive learning" may actually apply to feedforward models through the P 2 problem formulation (34) .Likewise, the nudged state defined by Eq. ( 8) reads for any layer until penultimate the same as Eq. ( 12) for problem P 1 and as follows for P 2 , taking β sufficiently small (92; 69):EQUATIONNote that many other choices of energy functions are possible in order to learn feedforward nets as a hierarchical energy-based model (15; 47; 29; 92) Breaking the forward-backward synchrony. As seen from Eq. ( 13), computing s k ±β still requires the storage of s k-1 . A trick is to notice that for sufficiently small β,s k ≈ (s k β + s k -β )/2 := sβ such that s k β ≈ σ θ k • sβ ± θ k+1 ⊤ • ∆s k+1 β with ∆s k+1 β := s k+1 β -s k+1 -β. Implemented this way, the gradient corresponding to problem P 2 with the above choice of energy function can be computed by simultaneously solving for s k β 's and s k -β 's (34) . This implementation can also be derived by picking hierarchical constraints (Eq. ( 3)), relaxing these (Eq. ( 5)) and reparametrizing them (Eq. ( 10)) ( 35). 
looking ahead & take-aways: 
 Forward-only learning beyond first order. One may wonder if (and how) second-order optimization could be emulated using a single circuit. A second-order update generally prescribes EQUATIONTherefore, moving to second order optimization here amounts to route the error signal ∇ s O backward through the system's inverse rather than its adjoint. While this viewpoint has motivated several implementations of self-invertible (rather than self-adjoint, or equivalently energy-based) circuits (46; 57; 11; 58) , it may not straightforwardly generalize to the mini-batch regime (|D| > 1).Temporal data? One way to leverage equilibrium techniques towards learning from sequences is to treat the neuron velocity ṡ as a variable of its own inside the energy function (30; 82) .This "cookbook" is an attempt to complement recent algorithm unification papers (92; 93; 63; 72) using pragmatic, ingredient-based logics. For instance, "Predictive coding" (89) amounts to apply Eq. ( 8) with a specific choice of energy function (3.4) and nudging scheme (3.2).Likewise, the "Forward-forward" algorithm (31) is a greedy learning strategy, with each block comprising a single layer, using local self-supervised loss (2.3). Conversely, one may want to split an architecture into blocks (Eq. ( 3)), apply standard implicit differentiation in some (2.2, §1), ZO in others (2.2, §1), along with a pipelining mechanism (2.3) depending on the constraints at hand, which does not correspond to any known algorithm in the literature. 
: 
 1 Problem statement Learning at equilibrium. Given an input x, we want to find a set of model parameters θ which minimizes a given objective O defined over the model (hidden and output) variables s(x, θ), such that the model variables abides by some constraint C. Namely: P 1 : min θ J(x, θ) := O(s(x, θ), θ) s.t. C(x, s(x, θ), θ) = 0.(Note that s may implicitly contain several layers, e.g. s = (s1 ⊤ , s 2 ⊤ , • • • , s N ⊤ ) ⊤ . Classically, Ois some cost function ℓ measuring the discrepancy between the model prediction and some groundtruth data and C the "physical laws" that governs the substrate sustaining the model prediction at equilibrium -note that this implicit formalism encompasses both feedforward models (60) , deep equilibrium models (4), as well as resistive networks governed by Kirchoff's laws (39; 86; 77) .The problem defined in Eq. ( 1) is traditionally solved via first-order optimization by estimating the gradient over a minibatch B (93): g(θ) := E x∼D [d θ J(x, θ)]. Therefore, solving P 1 boils down to how to compute these gradients. We herein call a "circuit" the abstract physical system that may realize a given equilibrium condition.The Lagrangian method. One way to solve P 1 is by writing the associated Lagrangian functional L 1 (s, λ, θ) := O(s, θ) + λ ⊤ • C(s, θ) and solving for the Karush-Kuhn-Tucker (KKT) conditions (81; 12) . Loosely speaking, the primal feasibility condition ∂ λ L 1 = 0 and dual feasibility condition ∂ s L 1 = 0 generalize the notions of "forward pass" and "backward pass" respectively, and can be viewed as two circuits determining s and λ whose equilibrium satisfy:EQUATIONwith the resulting gradient estimate reading as g(θ) = ∇ θ O(s, θ) + ∂ θ C(s, θ) ⊤ • λ. We briefly mention that this can also be derived using the Implicit Function Theorem (IFT) (93) . Here as well, Submitted to the Second Workshop on Machine Learning with New Compute Paradigms at NeurIPS (MLNCP 2024). Do not distribute.note that feedforward nets trained by backprop are a special case of Eq. ( 2). Yet, it appears from Eq. ( 2) that naively solving for s and λ comes with some challenges from the (high-level) hardware viewpoint: i) the transport problem: s and θ need to be transported from the inference circuit to the error circuit; ii) the memory problem: s needs to be stored; iii) analytical derivatives: ∂ s Cpotentially contains analytical derivatives which may have to be computed with high precision; iv) forward locking: in a hierarchical model, the circuit computing s k becomes idle when computing s k+1 (Eq. ( 3)); v) backward locking: symmetrically, λ k cannot be solved before λ k+1 (Eq. ( 4)); vi)forward-backward synchrony: λ is computed after s. P 2 : min θ J(x, θ) := O(s, θ) s.t. C 1 (s 1 , s 0 := x, θ 1 ) = 0, • • • , C K (s K , s K-1 , θ K ) = 0 (3)Here, the model is split into a hierarchy of K subcircuits (or "blocks") with parameters θ k , state s k , which may comprise one or multiple layers (K ≤ L), subject to the influence of the previous circuit through s k-1 . Note that for a given set of constraints, P 1 is generally not equivalent to P 2 -see § 3.6for an example. In this case, solving for the KKT conditions of P 2 yields:EQUATIONNote that O(s, θ) = ℓ(s K , y) corresponds to the classical "end-to-end" supervised learning setting in which case the second term of Eq. ( 4) vanishes except for the last block s K .Relaxed constraints. When C = ∇ s K (see Section 3), P 1 can be relaxed through an optimal value reformulation with a fixed Lagrangian multiplier β -1 as (73; 92; 35):EQUATIONThe resulting Lagrangian L 3 , sometimes called "surrogate", is intimately tied to energy-based learning (80) (Section 3) and is especially convenient to express quantities as finite differences in β which would otherwise appear as exact derivatives in β starting from P 1 (which amounts to send β → 0).Importantly, this relaxation could also well be applied at each level of the hierarchy of P 2 (35).2.2 Picking an algorithm to estimate λ's: standard approaches Implicit differentiation. In spite of the aforementioned problems inherent to Eq. ( 2), one may still directly solve for λ in combination with other tricks (see section 2.3). For feedforward models,∂ s C ⊤ may be explicitly inverted with λ = -∂ s C -⊤ • O reducing to backprop.In other cases, λ may be computed by implicit differentiation (41; 10; 28), which comes in many different flavors depending on the constraints at hand (13) . Note that these techniques can be slightly adjusted to accommodate some constraints. For instance, to avoid weight transport between the inference and error circuits, we can equip the error circuit with its own "feedback" parameters ω such that the error signal satisfies: Forward-only learning with zeroth order optimization. One technique to solve the memory and transport problems, which has recently regained some popularity, is zeroth order optimization (ZO) (49) . ZO techniques estimate a projection of the gradient g(θ) along some direction u in the weight space by performing multiple forward passes:∂ s C(s, ω) ⊤ • λ + ∇ s O(s, ω) = 0EQUATIONAn unbiased estimate of the gradient can be obtained by averaging multiple such derivatives:g(θ) = E u∼N (0,σ 2 ) [d ϵ (J(θ + ϵu)) | ϵ=0• u] (85; 7). However, the variance of this gradient estimate scales cubically with the number of model parameters (74) , thereby restricting the applicability of ZO to the realm of models which are small enough (84; 33; 7; 24) or behaving as such, i.e. pre-trained models (51) . One way to mitigate this problem is to perturb neurons instead of synapses (23; 61; 14; 40) , namely computing projections of the error signal as:v ⊤ • λ = d ϵ O(s ϵ , θ)| ϵ=0 with s ϵ implicitly determined through C(s ϵ , θ) + ϵv = 0.Similarly, an unbiased gradient estimate is obtained asg(θ) = ∇ θ O + ∂ θ C(s, θ) ⊤ • E v∼N (0,σ 2 ) [d ϵ O(s ϵ , θ)| ϵ=0 • v].Going beyond, one can teach small auxiliary networks synthetizing such directions to output "good" weight directions u instead of randomly sampling them (24) , reducing the variance of the gradient estimate at the cost of increasing the bias (75; 84). 
 
 Appendix:  



Reviews: review: title: title: Strong review paper with a unified mathematical framework
main_review: main_review: # Summary
The presented paper is a review of hardware-friendly learning methods for implicit models. The strength of this paper is its unified presentation of these methods with a focus on solving challenges in hardware implementations such as the weight-transport problem and other related problems. 
# Discussion
The topic of the paper is highly relevant for the workshop. Novel hardware accelerators often times come with additional constraints in contrast to general purpose computers like GPUs. The paper formalizes many such important constraints in a unified optimization framework and discusses this formulation in the context of a large body of relevant literature.

The presentation is both conceptually and mathematically clear and concise. The paper will be a valuable addition to the workshop. If alternative learning methods want to become more competitive with backprop at realistic problem sizes, a framework for casting the promises and limitations of these methods  might help overcoming some of the scaling issues that these methods still have. The paper might serve as a basis of discussion for future theoretical and empirical evidence.