High-performance Evolutionary Algorithms for Online Neuron Control

To simulate neuronal tuning function $f$ that an optimizer might encounter in a biological setting, we used units from pre-trained CNNs as models of visual neurons (Xiao and Kreiman, 2020; Lindsay, 2020). For the in vivo Evolution experiments, we aimed for optimizers that performed well with neurons across visual areas (including V1, V4, IT) and across different levels of signal-to-noise and single-neuron isolation. Thus we designed the benchmark ”problem set” to include units from multiple CNN architectures, layers, and noise levels, testing the overall performance for each optimizer.

We chose AlexNet and a deeper, adversarially trained ResNet (ResNet50-robust) as models of ventral visual cortex. The latter was chosen because it exhibits visual representations similar to the brain (high rank on Brain-Score (Schrimpf et al., 2018)). For each network, we picked 5 layers of different depths. It has been noted that units from shallow-to-deep layers prefer features of increasing complexity (Nguyen et al., 2016), similar to that in the ventral stream cortical hierarchy (Rose et al., 2021). For detailed information about these networks and their layers, see Sec. A.1. In the context of in vivo recordings, single-presentation neuronal responses are highly noisy (Czanner et al., 2015). To simulate the Poisson-like noisy response $\tilde{r}$, we added Gaussian noise with zero mean and standard deviation proportional to the raw response $\alpha r$ (ratio $\alpha$ represented the noise level). We tested three noise levels for each objective function: no noise ($\alpha=0$), low noise ($\alpha=0.2$), and high noise ($\alpha=0.5$).

In the first round, we chose 12 gradient-free optimizers as implemented or interfaced by nevergrad(Rapin and Teytaud, 2018): NGOpt, DE, TwoPointsDE, ES, CMA(Hansen and Ostermeier, 2001), DiagonalCMA(Akimoto and Hansen, 2020), RescaledCMA, SQPCMA, PSO, OnePlusOne, TBPSA, and RandomSearch. Here we compared these algorithms by their default hyper-parameter settings. Note that RandomSearch just sampled random vectors from an isotropic Gaussian distribution, finding the vector with the highest score, which formed the naive baseline (for a short introduction to these algorithms, see (Rapin and Teytaud, 2018)). Each algorithm ran with a budget of 3000 function evaluations and three repetitions per condition.

Among the 12 optimizers, we found that Covariance Matrix Adaptation Evolution Strategy (CMA) and DiagonalCMA were the top two algorithms in achieving maximal activation (Fig. 2). Since the upper bound of the activation of a unit in CNN was arbitrary, we divided the raw activations by the empirical maximal activation achieved by that given unit, across all algorithms and repetitions. By this measure, when pooling all conditions, CMAES achieved $0.699\pm 0.004$, DiagonalCMA achieved $0.677\pm 0.004$, as a reference Random Search baseline achieved $0.139\pm 0.002$ (mean $\pm$ sem, $N=1500$, Tab. A.2, Fig. 2. This difference of CMA-driven performance vs. any other optimizer was significant per a two-sample t-test: $t>60$ for all comparisons, except for CMA vs DiagonalCMA, where $t=3.80,p=1.4\times 10^{-4}$). We found the same result held consistently for units across CNN models, layers and noise levels (see comparison in Tab. A.2). The optimization score traces and the most activating images found for an example ResNet unit are shown in Fig. A.1 and Fig. A.3.

As for time efficiency, we measured the total run time taken by optimizing the objective function with a budget of 3000 evaluations²²2These optimizers were tested on single core machine with V100 GPU without batch processing of images or activations.. With units from AlexNet as objective, CMA had a longer runtime of $104.5\pm 30.0$ s (mean $\pm$ std); DiagonalCMA accelerated the runtime by roughly five-fold ($23.0\pm 4.5$ sec), although at a slightly compromised score (6.1%). As a reference, the baseline algorithm RandomSearch had an average runtime of $17.6\pm 2.2$ sec. Same trends held for ResNet50 units, though the runtime values were generally longer because ResNet50 is deeper (Tab. A.2). In conclusion, we found CMA and DiagonalCMA algorithms were both well-suited for this domain, while DiagonalCMA achieved a good trade-off between speed and performance.

Given the general success of CMA and DiagonalCMA algorithm, we were motivated to test other types of CMA algorithms in the second round, comparing runtime values and achieved activations.

As described in (Hansen, 2016), the CMA algorithm maintains and updates a Gaussian distribution $\mathcal{N}(m,\sigma C)$ in the $\mathbb{R}^{d}$ space, with the mean vector $m$ and step size $\sigma$ initialized by the user; the covariance matrix $C$ is initialized as identity matrix $I$. In each step, the algorithm samples a population of codes from this exploration distribution, $z_{t}^{(i)}\sim\mathcal{N}(m_{t},\sigma_{t}C_{t})$. After evaluating these codes, it updates the mean vector $m_{t}$ by a weighted combination of the highest ranking codes. Note that $\sigma_{t}$ controls both the spread of samples in a generation and the average step size of the mean vector update $\|m_{t+1}-m_{t}\|$. This step size $\sigma_{t}$ is adapted based on the accumulated path length in the last few steps. Moreover, CMA updates the covariance matrix $C_{t}$ by a few low-rank matrices to adapt the shape of the exploration distribution.

In the original CMAES algorithm, after each covariance update, the covariance matrix $C$ needed to be eigen-decomposed, in order to get $C^{-1}$. For a high dimensional space $d=4096$, it is costly to compute this decomposition at each update. Using the diagonal covariance matrix approximation, the inversion step could be simplified from a $\mathcal{O}(n^{3})$ operation to $\mathcal{O}(n)$ operation, which makes the DiagonalCMA (Akimoto and Hansen, 2020) much faster (Fig. 2). This inspired us to use modified covariance update rules to accelerate the optimizer.

We found Cholesky-CMA-ES (Loshchilov, 2015) which was proposed as a large scale variant of CMAES. By storing the Cholesky factor $A$ and its inverse of the covariance matrix, it could update these factors directly, without factorizing the covariance matrix at each update. An additional parameter A_update_freq could be adjusted to tune the frequency of this update.

We implemented the $(\mu/\mu_{w},\lambda)$-Cholesky-CMA-ES algorithm (CholeskyCMA) (Loshchilov, 2015) and compared it against the CMA and DiagonalCMA implemented in pycma library (Hansen and Auger, 2014; Akimoto and Hansen, 2020) and the Genetic Algorithm (GA), classically used in this domain (Ponce et al., 2019; Srinath et al., 2021). For CholeskyCMA, the hyper-parameters $\sigma_{0}$ and update frequency of Cholesky factor $A$ were tuned and fixed $\sigma_{0}=3.0$, A_update_freq$=10$. For GA, we used the code and hyperparameters from (Ponce et al., 2019). We allowed for 3000 function evaluations per optimizer, which was 75 generations with a population size of 40.

We slightly modified the in silico benchmark: we chose 7 layers from AlexNet (conv2 to fc8), 10 units from each layer, with 3 noise levels ($\alpha=0.0,0.2,0.5$). Each optimizer were run with 5 repetitions in each condition, totaling 1050 runs. We evaluated the clean score, i.e. the highest noise-free activation of the given unit and the runtime for each optimizer. Here we also used the empirical maximal clean score achieved for each unit to get the normalized clean score and to calculate statistics.

The results were summarized in Table 1 and Figures A.4,A.5. We found that all CMA algorithms outperformed Genetic Algorithm by a large margin: if we pooled all conditions, the mean normalized score of CMA was 166.7% of that of GA optimizer. Noise deteriorated the score for all optimizers, and there the performance gap between GA and CMA type algorithm was narrowed but persisted: the mean normalized score for CMA was 196.2% of GA in noise free scenario; 151.9% for $\alpha=0.2$; 146.4% for $\alpha=0.5$ (Fig. A.4). The overall performance values of the three CMA-type optimizers were not statistically different, and all surpassed GA. (1-way ANOVA, $F_{2,3147}=1.70,p=0.18$).

When we examined the performance per layer, we found the relative performance of optimizers had a significant interaction with the source layer of the unit. DiagonalCMA was more effective than Cholesky and original CMA in the earlier layers, but performed less well in deeper layers (Tab. 1, Fig. A.5). We tested this interaction with ANOVA on a linear model: activation $\sim$ optimizer + noiselevel + layer + layer * optimizer + noiselevel * optimizer; in which, ”optimizer” was modelled as a categorical variable (CholeskyCMA, CMA, DiagonalCMA), ”noiselevel” and ”layer” were continuous variable. All the factors except optimizer had statistically significant main effect; and interaction term layer * optimizer had $F_{2,3141}=14.98,P=3.34\times 10^{-7}$ (see Tab. A.1). This curious interaction would be interpreted below in the context of different covariance matrix adaption mechanisms (Sec. 3.1).

As for runtime³³3Here, latent codes were processed in batch of 40 by generator and CNN, thus the same optimizer ran faster than the previous experiment., DiagonalCMA was still the fastest, with average runtime $6.6\pm 0.1$ sec (mean$\pm$ sem, $N=1050$ same below.), while the second fastest optimizer was GA with runtime $16.6\pm 0.2$ sec. In contrast, the classic CMA algorithm taking $97.0\pm 0.8$ was the slowest among the four, while the CholeskyCMA took $42.4\pm 0.4$ sec per run. Indeed, updating Cholesky factors made it run faster without reducing performance.

After the second round of in silico benchmarks, we were ready to compete the top candidates with previous state-of-the-art Genetic Algorithm in vivo. We chose CholeskyCMA algorithm and compared it against the classic GA, since its speed and performance were both good across noise level and visual hierarchy. For detailed methods of in vivo experiments, we refer readers to Sec. A.2, but briefly, we recorded electrophysiological activity in two animals using chronically implanted arrays, placed within V1, V4 and IT.

We conducted $14$ experiments. Two out of fourteen ($2/14$) experiments did not result in a significant increase in firing rate of the neuron for either optimizer (per criterion $P<0.001$, for t-test between firing rates in first two and last two generation), which we excluded from the further analysis. From the $12$ experiments where at least one optimizer increased the firing rate, we normalized the firing rate of the neuron to the highest generation-averaged firing rate for the CholeskyCMA optimizer (Fig. 1 B). The normalized final generation activation for CholeskyCMA was $0.908\pm 0.018$ comparing to $0.628\pm 0.045$ for GA (paired t-test, $t=6.69,p=3.4\times 10^{-5},df=11$. Raw firing rates for each experiments are shown in Fig. A.6). Thus, the maximal activation of CholeskyCMA surpassed GA by 44%, which was comparable to the activation gain in the high-noise condition (46.2%, $\alpha=0.5$) in silico.

From this, we concluded that CholeskyCMA outperformed classic GA algorithm both in vivo and in silico, becoming the preferred algorithm for conducting activation maximization.

First, we noticed that in a high-dimensional space, the covariance matrix updates were impaired in the original $(\mu/\mu_{W},\lambda)$ CMA or CholeskyCMA algorithm. In the default setting⁴⁴4See the default setting of $c_{1},c_{\mu}$ in Tab. 2 of (Hansen and Auger, 2014)., the covariance learning rates $c_{1},c_{\mu}\propto 1/d^{2}$, which were exceedingly small at $d=4096$. Thus, the covariance matrix was updated negligibly and could be well approximated by an identity matrix. Recently, this was also pointed out in (Akimoto and Hansen, 2020) (Sec.4,5) and the authors proposed to increase the learning rate in DiagonalCMA. We tested the effectiveness of this modification.

Empirically, we validated this for the original, Cholesky, and Diagonal CMAES algorithms. We quantified this by measuring the condition number of the covariance matrix $\kappa(C)$ and its relative distance to identity matrix $\Delta(C)$.

We found that the final generation condition number $\kappa(C)$ of the CholeskyCMA was $1.000175\pm 0.000004$ (mean $\pm$ std, $N=175$, same below), while the relative distance to identity $\Delta(C)$ was $8.05\pm 0.31\times 10^{-12}$. In comparison, condition number $\kappa(C)$ for the original CMAES was $1.002739\pm 0.000046$, and for the DiagonalCMA, $1.124728\pm 0.034774$. The relative distances $\Delta(C)$ to the identity matrix were $7.13\pm 0.01\times 10^{-8}$ and $1.63\pm 0.21\times 10^{-4}$ for original CMA and DiagonalCMA algorithms. Though as designed, DiagonalCMA updated the covariance matrix more effectively than the other two CMA algorithms, its final covariance was still quite isotropic. On the other hand, for the original and CholeskyCMA algorithm, we could safely approximate the exploration distribution $p(z^{(i)}_{t+1}|m_{t})$ by an isotropic Gaussian scaled by the step size $\sigma_{t}$

This isotropic exploration throughout the Evolution experiments simplified subsequent analyses of the algorithm.

How is this related to the performance of the algorithm? This relates to the interaction between the unit layer and optimizer, noted above (Sec. 2.2,Tab. 1,Tab. A.1): the DiagonalCMA outperformed original CMA in earlier layers of CNN but not in deeper layers. It seems the faster update of diagonal covariance matrix was only beneficial for units in earlier layers. The diagonal covariance was designed for a separable, ill-conditioned functional landscape. We noticed that although all units had highly ill-conditioned landscapes, the units in shallower layers had tuning for fewer dimensions than units in deeper layers (Sec. C). In other words, the units in earlier layers had a larger invariant space, and a diagonal covariance might suit these landscape better than the more complex ones for deeper units.

In short, we conclude that the effectiveness of CMA-type algorithm was not in its adaptation of the exploratory distribution shape, and we postulate that it could work with fixed covariance.

Next, we investigated the geometric structure of the Evolution trajectories, through the lens of Principal Component Analysis (PCA), a linear dimension-reduction algorithm. For each experiment, we computed the mean latent vector for each generation $\{\bar{z}_{t}\},t=1...T$, and applied PCA to this $T$-by-$d$ matrix of mean vectors ($T=75$).

We found a pervasive sinusoidal structure to the typical trajectory. When a given trajectory was projected onto the top principal components (PC), the projection resembled cosine waves (Fig. 3A). On the $k$th component, the projected trajectories were well-fit by cosine functions of frequency $k\pi$: $A\cos(k\pi t/T)$ (Fig. D.1A). If we allowed the phase and frequency to be fit, $A\cos(2\pi\omega(t/T+\phi))$, the fit $R^{2}$ was above 0.80 for the top 16 PCs for 95% of trajectories (Fig. D.1B). As a result, projecting the mean trajectory onto the top PCs will result in Lissajous curves (Fig. D.3). Further, we analyzed the PCs of the in vivo Evolution experiments ($N=264$), and found the same sinusoidal structure (Fig. 3C). Even for control evolutions driven by noise, this sinusoidal structure persisted (Fig. D.2).

This intriguing structure was reminiscent of the PCA of another type of high-dimensional optimization trajectory: that of neural-network parameters during training (Lorch, 2016). This structure was later observed and analytically described for high-dimensional random walks or discrete Ornstein-Uhlenbeck (OU) processes guided by potential (Antognini and Sohl-Dickstein, 2018). To test this connection, we examined the explained variance of each principal components (Fig. 3B). We found that the explained variance $\rho(k)$ of the $k$th PC was well approximated by Eq. 6, which was derived as the theoretical limit for PCA of random walk with $T$ steps in high-dimensional space⁵⁵5Our Eq.6 was adapted from Eq.12 in (Antognini and Sohl-Dickstein, 2018). We found the original equation had the incorrect normalization factor, so we corrected the denominator as it is now. (Antognini and Sohl-Dickstein, 2018). If we take the step number $T\to\infty$ to infinity in Eq. 6, the explained variance of $k$th PC scales with $k^{-2}$ (Eq. 6), showing a fast decay.

Indeed, both in theory and in our data, the first principal component explained more than $60\%$ variance of mean evolution trajectory, and the top five PC explained close to $90\%$ of the variance (Fig. 3 B). In this view, the evolution trajectory of mean latent code $\bar{z}$ of CMAES could be regarded as a (guided) random walk in high dimensional space, with its main variance residing in a low-dimensional space.

How are evolutionary trajectories related to random walks? In the extreme, when the response is pure noise, the randomly weighted average of the candidates $\{z_{t}^{(i)}\}$ will be isotropic; thus each step $p(m_{t+1}|m_{t})$ is taken isotropically. Intuitively, this is close to a random walk. In our evolutions, the selection of candidates was biased by the visual-attribute tuning of units or neurons, thus their associated trajectories were random walks guided by given potentials. As derived in (Antognini and Sohl-Dickstein, 2018), for random walks guided by quadratic potentials, the structure of trajectory should be dominated by the dimensions with the smallest driving force (curvature). As we will see below (Sec. 3.3, Sec. B,C), the latent space had a large portion of dimensions to which the units were not tuned. This may explain why our evolution trajectories looked like random walks via PCA. In summary, we interpreted this as an intriguing and robust geometric property of high-dimensional curves (Antognini and Sohl-Dickstein, 2018), but its full significance remains to be defined.

Given that individual trajectories had a low-dimensional structure, we asked if there was a common subspace that these Evolution trajectories collectively shared. We reasoned that each optimization trajectory should lead to a local optimum of the tuning function of each given CNNs unit or cortical neuron, encoding relevant visual features. Though different neurons or units prefer diverse visual attributes of different complexity(Rose et al., 2021), it was possible that the preferred visual features populated certain subspaces in the latent space of our generator.

We first collected the in silico Evolution trajectories of $N=1050$ runs, across all conditions, and represented each trajectory by the mean code from the final generation, i.e. the evolution direction $\zeta_{j}$. We shuffled the entries of each vector $\{\zeta_{j}\},j=1...N$ to form a control collection of evolution directions $\{\zeta_{shfl}\}$, preserving their vector norms. We found that the collection of observed Evolution directions was lower-dimensional than its shuffled counterpart: by PCA, the explained variance of top 7 PCs was larger than the shuffled counterpart ($p<0.005$ comparing to 500 shuffles, Fig. A.7).

Next, we examined the relationship between the collection of trajectories and the global geometry of the space (Sec. B). Previously, we reported that the latent space of the generator exhibits a highly anisotropic geometry, as quantified by the Riemannian metric tensor $H$ of this image manifold (Wang and Ponce, 2021). The bilinear form $v^{T}Hv$ defined by this Riemannian metric tensor represented the degree of image change when moving in direction $v$. Consequently, moving in top eigen-directions changes the generated image much more than other directions, while moving in the bottom eigenspace scarcely affects the image (Fig. B.1B,C). Moreover, this structure is homogeneous in the space, thus similar directions will cause rapid image change at different positions in the latent space. Thus, there exist global dimensions that affect the image a lot, and dimensions that do not. We hypothesized that the optimization trajectories might preferentially travel in some part of this eigenspace.

We used the averaged metric tensor $H$ from Wang and Ponce (2021, (Wang and Ponce, 2021)). Its eigenvectors $U=[u_{1},u_{2}...u_{d}]$ with eigenvalues $\lambda_{k},k=1..,d$ became our reference frame for analyzing trajectories. We projected the collection of evolution directions $\{\zeta_{j}\}$ onto this basis, and examined the mean projection amplitude on each eigenvector $\frac{1}{P}\sum_{i}^{P}|u_{k}^{T}\zeta_{i}|$ as a function of eigenvalue $\lambda_{k}$ (Fig. 4A). Since for in silico experiments, the initial generation was the origin, this quantity was the average distance travelled along the eigenvector across runs. We expected the trajectories to travel further for eigen-dimensions where more CNN units exhibit ”gradient”.

We applied the Kolmogorov–Smirnov test to determine if the distribution of projection coefficient in the top and bottom eigen-dimensions were different. We found that the projection coefficients onto the top $800$ eigen-dimensions as a distribution were significantly different from those onto bottom dimensions (KS statistics $0.223,P=9.3\times 10^{-18}$, Fig. 4A). Further, within the top $800$ PC dimensions, the mean traveled distance strongly correlated with the log eigenvalues (Pearson correlation $0.738,p=2.8\times 10^{-138}$, Fig. 4A), i.e. the trajectories tended to travel further for dimensions with larger eigenvalues.

As for in vivo evolutions, we examined the collection of evolution trajectories ($N=264$ from 2 monkeys, (Rose et al., 2021)) and projected them onto the Hessian eigenframe as above. We found that they preferentially aligned with the top eigenspace than the lower eigenspace (Fig. 4B): the distance traveled in the top $800$ eigenspace correlated with the eigenvalues (Pearson correlation $0.319,p=1.2\times 10^{-15}$). Since in vivo evolutions used a set of initial codes $\{z^{(i)}_{0}\}$ instead of the origin, we used noise-driven evolutions starting from these initial codes as control. As a baseline, the noise-driven evolutions had a lower correlation between the distance traveled and the eigenvalues ($r=0.198$ for the noise evolution with matching generations; $r=0.223$ for that with matching code norm). Our results were robust to the choice of cutoff dimension (600-1000).

In conclusion, this result showed us that the evolution trajectories traveled further in the the top eigenspace, and the average distance traveled was positively correlated with the eigenvalue. How do we interpret this effect? Since the top eigen-dimensions change the images more perceptibly, the tuning functions of neurons and CNN units were more likely to exhibit a ”gradient” in such subspace. As the lower eigenspace did not induce perceptible changes, those dimensions barely affected the activations of units or neurons. Thus, no signal could guide the search in lower eigenspace, inducing dynamics similar to pure diffusion. As the visually tuned units could exhibit gradient in the top eigenspace, the search would be a diffusion with a driving force, allowing the optimizers to travel farther in the top eigenspace.

Finally, we noticed one geometric structure that facilitated the convergence of the CMA algorithm. For the DeePSim generator(Dosovitskiy and Brox, 2016) with DCGAN architecture(Radford et al., 2015), the mapping $G$ was relatively linear. Namely, changing the scale of the input $z$ mainly affected the contrast of the generated image $G(z)$ (Fig. B.1). Thus, when the base vector $z_{0}$ has a larger norm, travelling the same euclidean distance $\Delta z$ will induce a smaller perceptual change; in contrast, traveling the same angular deviation $\Delta\theta$ will result in a similar pattern change regardless of the norm of the base vector $z_{0}$ (Fig. A.8). We reasoned that the angular distance in latent space would be a better proxy to the perceptual distance across generated images.

We validated this using the perceptual dissimilarity metric Learned Perceptual Image Patch Similarity (LPIPS, (Zhang et al., 2018)). We measured image variability in each generation using the mean LPIPS distance between all pairs. We measured the code variability by the standard deviation (std) of latent codes $\tilde{\sigma}_{t}$ averaged across $d$ dimensions, which was an estimate of the step size $\sigma_{t}$ and was strongly correlated with the mean L2 distance among codes. We also measured the mean angular distance between all pairs of codes in each generation. We found that the std of latent codes $\tilde{\sigma}_{t}$ decreased from $\tilde{\sigma}_{0}=3.0$ (first generation) to $\tilde{\sigma}_{T}=2.31\pm 0.04$ ($N=1050$), and hit the floor at around 25 generations (Fig. 6). In contrast, the perceptual variability kept decreasing till the end, which was better approximated by the trend of angular variability within each generation (from 1.57 rad to 0.50 rad, Fig. 6). This discrepancy of Euclidean and angular distance was mediated by the increasing vector norm during evolution: we found that the squared norm of the mean latent codes $\bar{z}_{t}$ scaled linearly with generation $\|\bar{z}_{t}\|^{2}\propto t$ ($r^{2}=0.9985\pm 0.008$, $N=1050$, Fig. D.2C), which is a classic property of a random walk. As a combined effect of the increasing vector norm and the decreasing step size $\sigma_{t}$, the angular variability within a generation decreased towards the end. Similarly, we found that across each given trajectory, angular distance was also a better proxy for the perceptual distance between images than L2 distance (Fig. 5). The mean-code images from late generations were more similar to each other than the ones from initial generations, as predicted by the angular distance between the mean vectors — but not L2 distance.

For a visual neuron or a CNN unit, the perceptual distances across images matter more. Thus, when designing the new optimizer, we reasoned it should control the angular step size as a proxy to perceptual distance, instead of the Euclidean distance ($\sigma$) (Sec. 4).