SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhicheng Sun, Yadong Mu*  
Peking University, Beijing, China  
{sunzc, myd}@pku.edu.cn
The human brain rewrites itself for neuroplasticity in the presence of new tasks. We are inspired to harness this key process in continual reinforcement learning, prioritizing adaptation to non-stationary environments. In distinction to existing rewiring approaches that rely on pruning or dynamic routing, which may limit network capacity and plasticity, this work presents a novel rewiring scheme by permuting hidden neurons. Specifically, the neuron permutation is parameterized to be end-to-end learnable and can rearrange all available synapses to explore a large span of weight space, thereby promoting adaptivity. In addition, we introduce two main designs to steer the rewiring process in continual reinforcement learning: first, a multi-mode rewiring strategy is proposed which diversifies the policy and encourages exploration when encountering new environments. Secondly, to ensure stability on history tasks, the network is devised to cache each learned wiring while subtly updating its weights, allowing for retrospective recovery of any previous state appropriate for the task. Meanwhile, an alignment mechanism is curated to achieve better plasticity-stability tradeoff by jointly optimizing cached wirings and weights. Our proposed method is comprehensively evaluated on 18 continual reinforcement learning scenarios ranging from locomotion to manipulation, demonstrating its advantages over state-of-the-art competitors in performance-efficiency tradeoffs. Code is available at https://github.com/feifeiobama/RewireNeuron.
Despite remarkable advances in artificial neural networks, they have yet to emulate the adaptivity of the human brain in realistic, non-stationary environments. A plausible explanation for this disparity from neuroscience is that the brain gains additional structural plasticity by constantly rewiring itself throughout one's lifetime [9, 35, 67]. In fact, with extensive rewiring, the brain could even confront more adverse scenarios such as stroke [11, 68]. By contrast, the wiring of artificial neural networks is usually optimized for a single domain [79, 72, 70] and remains fixed thereafter, thus exhibiting limited adaptivity in the real world. This comparison highlights the potential of rewiring neurons in enabling neural networks to continually engage and learn in dynamically evolving environments, which is referred to as the problem of continual reinforcement learning [53, 64].
A few preliminary studies have exploited this rewiring process in deep learning with other research scopes. They rewired neural networks through pruning [37, 28, 5] or dynamic routing [43, 55, 74] to allow more efficient or input-dependent inference at test time. Following this line of work, a mainstay of continual learning methods [41, 32, 23] learned task-specific masks on neuron wiring to alleviate catastrophic interference [42] between different tasks. Unfortunately, existing rewiring approaches that rely on selecting a subnetwork are typically prone to capacity shrinkage over time and therefore may not be competent for continual adaptation to new reinforcement environments.
(a) Synaptic level
(b) Network level
Figure 1: Illustration of rewiring neurons via permutation. (a) In micro, permuting neurons allows reusing existing synaptic connections to create functionally new computation units. (b) In macro, layerwise permutation of hidden neurons provides extra structural plasticity in continual learning. Compared to alternative approaches, this method enjoys the higher capacity of the full network.
To harness the adaptivity of the rewiring process as in the human brain, we propose a novel rewiring scheme by neuron permutation, which fully reuses network weights to achieve structural plasticity. A motivating example is presented in Figure 1a, where instead of modifying synaptic strengths, one can directly permute neurons to populate certain desired functionalities. In more general cases, neuron permutation may serve as a complementary mechanism to weight updates to facilitate the learning process, as it encodes numerous variations by rewiring existing synapses. For instance, a tiny neural network shown in Fig. 1b encapsulates over 10,000 wires to be induced by permutation, with the number growing exponentially w.r.t. network depth and even more dramatically w.r.t. width. And thanks to the recent differentiable sorting algorithms [22, 50], it can be efficiently parameterized for end-to-end learning and thus integrated for adaptation to new scenarios.
In the context of continual reinforcement learning, we introduce two rewiring designs to emphasize exploration and stability-plasticity, respectively, which are the two essential aspects of the problem. To enhance exploration in novel environments, a multi-mode rewiring strategy is devised. In detail, it promotes policy diversity by randomly selecting a rewiring mode at each step, while effectively sharing knowledge across different modes via a distillation loss. To enforce stability on past tasks, we let the network cache each learned wiring and regularize its weight changes, allowing previous states to be recovered at low memory cost. Meanwhile, an alignment mechanism is proposed that jointly refines cached wirings and the latest weights, complementing plasticity.
In summary, our contributions are as follows. (1) We propose a novel rewiring scheme by permuting hidden neurons, allowing for additional structural plasticity under non-stationarity. (2) To encourage exploration in new environments, a multi-mode rewiring strategy is introduced that increases policy diversity. (3) To ensure stability on history tasks, the network caches each learned wiring and aligns it with the newest weights. (4) Experiments on 18 continual reinforcement learning scenarios from Brax and Continual World clearly demonstrate the strengths of our method.
Permuting neurons. The effectiveness of neuron permutation has been demonstrated in aligning differently trained neural networks [38, 59]. Subsequent literature [6, 62, 16, 2] exploited permutation symmetries of hidden neurons [29] for finding low-loss connectors [13, 19] between loss minima. Further applications of permutation symmetry include designing equivariant neural architectures [46, 77] and facilitating continual learning [47]. Our work is orthogonal to the aforementioned studies in that we use permutation to produce functionally diverse networks, rather than simply aligning them.
The key technique behind permutation learning is differentiable sorting. A number of works [1, 44, 10, 15] approximated permutation with doubly stochastic matrices and ran the Sinkhorn algorithm [60]. Petersen et al. [48, 49] cast differentiable sorting as a series of continuous swap operations. However, both methods involve multiple forward iterations, rendering them inefficient for longer permutations. Therefore, a more lightweight alternative via unimodal row stochastic relaxation [22, 50] is adopted in our experiments. Among various efforts to embed this technique into the network structure [63, 56], we are the first to advocate it for promoting plasticity in continual learning.
Continual learning. To address the major issues in continual learning [53] (or lifelong learning [64]), such as catastrophic forgetting [42] and the stability-plasticity dilemma [7], three main approaches have been developed: memory replay [51, 58], weight regularization [33, 39] and parameter isolation [54, 41, 32, 23]. Within the last category, Mallya and Lazebnik [41] pioneered the use of pruning [37, 28] to freeze important parameters for stability, but at the cost of capacity shrinkage. Ke et al. [32] resorted to dynamic routing [43, 55] for sharing knowledge between different tasks. Gurbuz and Dovrolis [23] iteratively pruned and grew neuron connections in a sparse neural network to maintain fixed capacity. In comparison, our work leverages all model parameters with no additional pruning step, and can therefore be easily adapted to different scenarios.
Continual reinforcement learning. Since the earliest proposals [53, 64], continual learning has been intertwined with reinforcement learning [61] for control and robotics. In this intersecting area, a main stream of studies [36, 8, 71] made specific assumptions on the Markov decision process to theoretically maximize reward amidst non-stationarity. Several notable works [3, 45, 52] employed meta-learning [57] for faster adaptation to new environments. There are also counterparts [54, 31, 20] that focused on continual learning of policy networks at the parameter level. Along this direction, Rusu et al. [54] proposed to dynamically grow the policy network for incorporating new knowledge. Kaplanis et al. [31] distilled knowledge [30] from previously learned policies to ensure stability. Recently, Gaya et al. [20] suggested that learning a subspace of policies strikes a good balance between model size and performance. Following these efforts, we continue to push the performance-efficiency frontier with the proposed neuron rewiring scheme.
In this section, we begin with a brief formulation of continual reinforcement learning (Section 3.1). Then, a novel rewiring approach is introduced in Section 3.2, along with its emphasis on two aspects: exploration (Section 3.3) and stability-plasticity (Section 3.4), both deemed crucial to the problem. Finally, we remark its connection to neuroscience observations in Section 3.5.
Continual reinforcement learning concerns learning over non-stationary environments, which are characterized by a sequence of  $T$  tasks. The goal is to obtain a well-performed global policy  $\pi(\mathbf{a}|\mathbf{s},t)$  that takes a state  $\mathbf{s}$  and a task identifier  $t \in [1..T]$  as input, and outputs a distribution over actions  $\mathbf{a}$ . To achieve this, a policy network is employed. Following the current literature [69, 20], we adopt a multi-layer perceptron architecture comprising a series of alternating affine layers  $\mathbf{W} = \{\mathbf{W}_l\}_{l=1}^L$  and element-wise activation functions  $\sigma$ . The policy inference is formulated as:
$$
\boldsymbol {Y} = \boldsymbol {W} _ {L} \circ \sigma \circ \boldsymbol {W} _ {L - 1} \circ \dots \circ \sigma \circ \boldsymbol {W} _ {1} \boldsymbol {X}, \tag {1}
$$
where  $X$  and  $Y$  denote the input and the predicted policy, respectively. This problem poses significant challenges on the policy network, as it must quickly adapt to environmental changes through efficient exploration, while not catastrophically forgetting the previously learned policy.
We address the problem from a structural perspective, focusing on the design of the policy network. Unlike many existing approaches of this type that grow separate policy networks for each task at linearly increasing memory cost, our solution involves only two networks: one for the current task and one for the previous task. We will refer to the current network by default, and when necessary, use the superscripts  $t$  and  $t - 1$  to distinguish between these two networks.
3.2 Rewiring via permutation
Inspired by the brain's remarkable adaptivity by rewiring itself [9, 35, 67], we seek to incorporate a similar process into our policy network. Ideally, this would provide additional plasticity to standard weight learning in non-stationary environments. However, the implementation is not trivial due to major architectural differences between the brain and our network. Most notably, the brain exhibits sparse connectivity, whereas our network consists of fully connected layers.
<sup>1</sup>For simplicity, we will omit biases from our presentation, while it is fairly easy to extend the following discussion to include biases.
(a) Learning to rewire
(b) Multi-mode rewiring
(c) Memory of wirings
Figure 2: Illustration of our proposed rewiring framework. (a) The rewiring process, formulated as neuron permutation, is parameterized by score vectors indicating the neuron order within each layer, which can be learned via a differentiable sorting algorithm. (b) In the presence of a new task, our network simultaneously learns  $K$  different wirings to diversify the policy and promote exploration. (c) Each learned wiring is cached and then continually refined to align with the newest weights.
To this end, we propose a novel rewiring scheme that exploits the layered structure of the network. While each layer has intricate connections and rewiring within it may seem complex, we take an alternative approach by rewiring between layers through neuron permutation, as shown in Fig. 1b. This allows reusing existing synapses to redirect the information flow within the network, thereby inducing new wires with varying functionalities. Formally, it is described as inserting a series of permutation matrices  $\pmb{P} = \{\pmb{P}_l\}_{l=1}^{L-1}$  into the network architecture:
$$
\boldsymbol {Y} = \boldsymbol {W} _ {L} \circ \sigma \circ \boldsymbol {P} _ {L - 1} \boldsymbol {W} _ {L - 1} \circ \dots \circ \sigma \circ \boldsymbol {P} _ {1} \boldsymbol {W} _ {1} \boldsymbol {X}. \tag {2}
$$
Note that we permute only the hidden neurons at layer  $l \in [1..L - 1]$ , since the input and output formats are assumed to be fixed. This formulation also differs from permutation symmetry works [62, 16, 2], as our focus is on populating different wirings rather than equivalent ones.
Then, a differentiable sorting algorithm is employed to learn the permutations, as illustrated in Fig. 2a. We follow Prillo and Eisenschlos [50] to parameterize each permutation  $P_{l}$  with a score vector  $v_{l}$ . In detail, this score vector  $v_{l}$  uniquely corresponds to an order  $z_{l}$ , which can then be mapped bijectively to a permutation matrix. Thus, the permutation can be computed in a forward pass:
$$
\boldsymbol {P} _ {l} = \boldsymbol {I} [ \boldsymbol {z} _ {l},: ], \quad \boldsymbol {z} _ {l} = \operatorname {a r g s o r t} (\boldsymbol {v} _ {l}), \tag {3}
$$
where  $I$  is an identity matrix. Since both the indexing and argsort operators are non-differentiable, the following smoothed proxy is used instead during backpropagation:
$$
\hat {\boldsymbol {P}} _ {l} = \operatorname {s o f t m a x} \left(\frac {- d \left(\operatorname {s o r t} \left(\boldsymbol {v} _ {l}\right) \mathbf {1} ^ {\top} , \mathbf {1} \boldsymbol {v} _ {l} ^ {\top}\right)}{\tau}\right), \tag {4}
$$
where  $d(\cdot, \cdot)$  uses the  $L_{1}$  distance and  $\mathbf{1}$  is an all-ones vector. This way we can backpropagate loss errors through permutation and learn the score vector  $\pmb{v}_{l}$  in an end-to-end manner. For clarity, the following text will represent the induced wiring as  $P_{l}$ , omitting parameterization details involving  $\pmb{v}_{l}$ .
In addition, we would like to highlight two aspects of the proposed rewiring scheme:
Parameter efficiency. Neuron permutation is highly parameter-efficient. In a network of width  $n$ , each layer of permutation learning requires only  $O(n)$  parameters. And after training, they can be compressed into a discrete order with even smaller memory footprint. This distinguishes our method from incorporating additional linear layers, which introduces a quadratic parameter overhead that is undesirable in a resource-constrained continual learning setting.
Adaptivity. Neuron permutation exploits a variety of structural variations to facilitate adaptation. There are over  $10^{1520}$  potential wires in our adopted policy network, and the number continues to grow factorially w.r.t. network width and exponentially w.r.t. depth. Moreover, each wiring maps the current network weight to a distinct weight point, thus covering a large span of weight space. Together with the parameter efficiency, it significantly enhances the adaptivity of our network in non-stationary environments, which will be validated through experiments.
3.3 Rewiring for exploration
In applying the proposed rewiring scheme to continual reinforcement learning, there are two key challenges that need to be addressed. The first challenge is to efficiently explore the policy space when encountering a novel environment, in order to avoid being trapped in suboptimal solutions. This notion of exploration has been emphasized in many related areas, such as maximum entropy reinforcement learning [78, 24, 25], unsupervised skill discovery [17], and few-shot adaptation [34, 21].
To encourage exploration in new environments, we propose a multi-mode rewiring strategy that maintains a set of different wirings to induce diverse policies, as shown in Fig. 2b. At each step, the agent randomly samples a mode to select from the following candidate wirings:
$$
\boldsymbol {P} _ {l} \in \left\{\boldsymbol {P} _ {l, 1}, \boldsymbol {P} _ {l, 2}, \dots , \boldsymbol {P} _ {l, K} \right\}. \tag {5}
$$
where  $K$  is the total number of modes, and the mode is shared across all layers. Also, to ensure an equivalent learning rate, we multiply the learning rate of each wiring by  $K$ .
Figure 3 illustrates the policies induced by multi-mode rewiring. Interestingly, these policies naturally exhibit a diverge-and-converge pattern without any constraints on the wirings. That is, they diverge rapidly at first, but then gradually converge to a high similarity. This implies that multi-mode rewiring enables the agent to explore various policies before finally settling on some effective ones. Thus, we can safely select one of the wirings for subsequent evaluation and training.
We further distill knowledge [30] across different wirings for knowledge sharing, which resembles the idea of mutual learning [76, 27] but in a parameter-efficient manner. Kullback-Leibler (KL) divergence is used as the loss:
$$
L _ {\mathrm {K L}} (\boldsymbol {W}, \boldsymbol {P}) = \mathbb {E} _ {k ^ {\prime} \neq k} \left[ D _ {\mathrm {K L}} \left(\pi_ {k ^ {\prime}} (\cdot | \mathbf {s}) \| \pi_ {k} (\cdot | \mathbf {s})\right) \right], \tag {6}
$$
Figure 3: Evolution of Kullback-Leibler divergence between different policies induced by multi-mode rewiring at  $K = 3$
where  $\pi_{k}$  is the policy induced by the  $k$ -th wiring, and we only backpropagate through one wiring to save computational cost. As shown in Fig. 3, it significantly speeds up the convergence within modes.
3.4 Rewiring for stability-plasticity
The second major challenge in continual reinforcement learning is the stability-plasticity dilemma [7], where an agent could either suffer severe performance degradation on history tasks due to catastrophic interference [42], or fail to adapt to new environments if an overly rigid constraint is imposed to mitigate forgetting. This requires careful maintenance of old policies when acquiring new ones.
Caching each wiring. Our base solution is to cache each learned wiring and regularize the weights, as shown in Fig. 2c. Intuitively, this would allow recovery of both previous weights and wirings, thereby preserving the learned policies. On the wiring side, thanks to the parameter efficiency of permutation, it yields a very small storage overhead (less than  $0.5\%$  per wiring in our network) compared to caching the entire network. On the weight side, a typical choice would be  $L_{2}$  regularization, which has proven to mitigate forgetting well in many continual reinforcement learning scenarios [69, 20]:
$$
L _ {\text {r e g}} \left(\boldsymbol {W} ^ {t}\right) = \sum_ {l = 1} ^ {L} \left\| \boldsymbol {W} _ {l} ^ {t} - \boldsymbol {W} _ {l} ^ {t - 1} \right\| ^ {2}. \tag {7}
$$
However, such a scheme with weight regularization is usually prone to the stability-plasticity dilemma mentioned above, since the plasticity of the weights is largely reduced by the weight regularizer.
Aligning wirings with weights. To address the stability-plasticity dilemma in weight regularization, we reformulate the problem into the misalignment between the latest weights and cached wirings. For example, the forgetting since the last training stage can be reproduced by substituting the newest weights  $\pmb{W}_l^t$  and the last cached wiring  $\pmb{P}_l^{t-1}$  into the inference procedure in Eq. (2):
$$
\boldsymbol {Y} = \dots \circ \sigma \circ \boldsymbol {P} _ {l} ^ {t - 1} \boldsymbol {W} _ {l} ^ {t} \circ \dots \boldsymbol {X}. \tag {8}
$$
Instead of framing it as a shift in the weight space (as in EWC [33], SI [75], etc.), we view it as the misalignment between updated weights and fixed wiring. This provides a new perspective on the loss of plasticity: existing schemes always align the weights to the wiring, while the wiring remains fixed, resulting in an overly rigid constraint. So we ask, why not align the wiring to the weights?
Inspired by this, we propose a new alignment mechanism that jointly refines the wiring and the weights to align with each other. Specifically, the optimization of each cached wiring is achieved by inserting two wiring adapters  $P_{l}^{\prime}$  and  $P_{l}^{\prime \prime}$  before and after it. In this way, the wiring in Eq. (8) is adapted to:
$$
\boldsymbol {Y} = \dots \circ \sigma \circ \underbrace {\boldsymbol {P} _ {l} ^ {\prime} \boldsymbol {P} _ {l} ^ {t - 1} \boldsymbol {P} _ {l} ^ {\prime \prime \top}} _ {\text {a d a p t e r s o n} \boldsymbol {P} _ {l} ^ {t - 1}} \boldsymbol {W} _ {l} ^ {t} \circ \dots \boldsymbol {X}. \tag {9}
$$
Note that the element-wise activation functions can be interchanged with permutation layers, yielding an equivalent form in which the adapters could alternatively be interpreted as adapters on the weights:
$$
\boldsymbol {Y} = \dots \circ \sigma \circ \boldsymbol {P} _ {l} ^ {t - 1} \underbrace {\boldsymbol {P} _ {l} ^ {\prime \prime \top} \boldsymbol {W} _ {l} ^ {t} \boldsymbol {P} _ {l - 1} ^ {\prime}} _ {\text {a d a p t e r s o n} \boldsymbol {W} _ {l} ^ {t}} \circ \dots \boldsymbol {X}. \tag {10}
$$
From here, we can switch back from the misalignment interpretation to the weight shift interpretation. In order to mitigate forgetting, the adapted weights should be close to the previous weights. which yields the following regularizer on both the wiring and the weights:
$$
L _ {\mathrm {S P}} \left(\boldsymbol {W} ^ {t}, \boldsymbol {P} ^ {\prime}, \boldsymbol {P} ^ {\prime \prime}\right) = \sum_ {l = 1} ^ {L} \left\| \boldsymbol {W} _ {l} ^ {t} - \boldsymbol {P} _ {l} ^ {\prime \prime} \boldsymbol {W} _ {l} ^ {t - 1} \boldsymbol {P} _ {l - 1} ^ {\prime \top} \right\| ^ {2}. \tag {11}
$$
And after the wiring adapters are learned, they are used to refine all cached wirings. For instance, the last cached wiring is updated as:
$$
\boldsymbol {P} _ {l} ^ {t - 1} = \boldsymbol {P} _ {l} ^ {\prime} \boldsymbol {P} _ {l} ^ {t - 1} \boldsymbol {P} _ {l} ^ {\prime \prime \top}. \tag {12}
$$
Replacing the previous weight regularizer with this new alignment mechanism improves stability-plasticity in two aspects: (1) The weight constraint is relaxed, since the weight is not necessarily constrained near the previous one. (2) With aligned wirings to counteract weight changes, a smaller coefficient on the regularizer will suffice to mitigate forgetting.
Finally, the proposed rewiring designs are incorporated into the existing reinforcement learning framework (Soft Actor-Critic [25] in our case) using a weighted sum of the standard training objective (maximum reward plus maximum entropy) and the two newly introduced losses  $L_{\mathrm{KL}}$  and  $L_{\mathrm{SP}}$  with coefficients  $\alpha$  and  $\beta$ . The rest of the training specifications remain unchanged.
3.5 Connection to neuroscience
Since our inspiration comes from the rewiring process in the brain, it is natural to review and compare our designs with relevant neuroscience observations. Here, we would like to remark some noteworthy similarities and distinctions between our rewiring approach and the actual process.
Interdependence of weights and wiring. It is elucidated that the weight and wiring changes in the brain are not mutually exclusive, but could take place simultaneously to facilitate plasticity [9]. In our approach, a similar line of thinking is followed by updating both weights and wiring during adaptation to new environments. Moreover, we introduce an alignment mechanism (Eqs. (11) and (12)) that effectively coordinates these two updates to promote stability-plasticity.