Learning and Aligning Structured Random Feature Networks

Authors: vivian  white
muawiz  chaudhary
guy  wolf
guillaume  lajoie
kameron  decker
Venue: ICLR_cc_2024_Workshop_Re-Align
Published: 2024-02-08 07:15:53

Abstract: Artificial neural networks (ANNs) are considered "black boxes'' due to the difficulty of interpreting their learned weights. 
While choosing the best features is not well understood, random feature networks (RFNs) and wavelet scattering ground some ANN learning mechanisms in function space with tractable mathematics. Meanwhile, the genetic code has evolved over millions of years, shaping the brain to develop variable neural circuits with reliable structure that resemble RFNs. We explore a similar approach, embedding neuro-inspired, wavelet-like weights into multilayer RFNs. These can outperform scattering and have kernels that describe their function space at large width. We build learnable and deeper versions of these models where we can optimize separate spatial and channel covariances of the convolutional weight distributions. We find that these networks can perform comparatively with conventional ANNs while dramatically reducing the number of trainable parameters. Channel covariances are most influential, and both weight and activation alignment are needed for classification performance. Our work outlines how neuro-inspired configurations may lead to better performance in key cases and offers a potentially tractable reduced model for ANN learning.

Main paper: 
 introduction: 
 Deep neural networks (DNNs) achieve state-of-the-art results, at the cost of large numbers of learned parameters and complex input-output mappings that are difficult to understand. These interpretability issues impact the trustworthiness of deep models. A key goal of neural network theory is to understand the roles of architecture and weight structure in DNNs to help bridge the gap between performance and interpretability. Among the extant theory for deep networks are kernels (Neal, 1995; Rahimi & Recht, 2007; Jacot et al., 2018) and wavelet scattering (Mallat, 2012; 2016) . Kernels define a reproducing kernel Hilbert space (RKHS) for functions, which can be used to characterize the properties of neural network typically at or near their random initialization. Functions that have small norms in the RKHS are easier to learn with fewer samples and have stronger generalization abilities. Wavelet scattering is a type of deep convolutional network (CNN) that cascades fixed wavelet filters with complex modulus nonlinearities and low-pass filters to separate image variation Figure 1 : Learnable random feature convolutional networks: For a given layer representation, weights in our networks are modeled as drawn from correlated Gaussians where the covariance may be learned. We separably factor the weight covariance in a convolutional layer (red box) into two parts: the covariance across n channels and the spatial covariance across d = d 1 • d 2 pixels in an input patch. The spatial and channel covariances may be fixed at initialization or learned. across scales. These networks build invariance to deformations, are mathematically tractable, and rely on fewer parameters to achieve impressive performance, particularly in low-data regimes.Whether or not the representations in the brain are "aligned" to those in DNNs is an interesting and complicated question (Sucholutsky et al., 2023) . Despite the name, neural networks on computers typically share an only superficial resemblance to the brain.In biological neurons, receptive fields are localized areas where stimuli evoke inhibitory or excitatory neuronal responses. Visual receptive fields have been optimized over millions of years to detect small visual changes including edges, orientations, and spatial frequencies (Hubel & Wiesel, 1959) . They are classically understood as comparable to wavelets, which are selective to scale and orientation (Olshausen & Field, 1996) . We connect these ideas to random feature networks (RFNs), where hidden-layer weights are drawn from random distributions and only the output weights are learned. Such structured RFNs have kernels that vary depending on weight distribution and describe their function space at large width (Pandey et al., 2022) .We build randomized scattering networks (RSNs) with weight covariance modeled from biological receptive field data and identical architecture to the wavelet scattering transform. The frozen filters in these models are more interpretable than trained weights due to their wavelet eigenbasis and neural inspiration. We compare V1-like RSNs against scattering networks and uncorrelated RSNs on image classification. We find that the RSNs outperform scattering in most cases and biologically-inspired weights may improve performance.We further investigate the importance of covariance structure by building deeper CNNs where the weight covariances are factored into a product of spatial and channel covariances (Fig. 1 ). By allowing networks to learn these covariance matrices, we gain improved performance and find that learning the channel part is more important than the spatial covariance in a deep architecture.Finally, we try aligning resampled and trained factored networks. Guth et al. (2023) introduced "rainbow networks" that cascade random feature maps with colored weight covariance matched to the covariance structure of traditionally trained CNNs. These networks assume that the weight dependencies across layers are reduced to rotations of independent random feature matrices which align the input activations. Rainbow networks thus model the joint probability distribution of the weights of trained networks. The authors sampled new weights from the learned weight distributions and found that the sampled networks retrieve similar performance to the trained models after feature alignment. Applying those techniques to our model, we show a collapse of this approach in later layers of the sampled networks. Alignment to the reference network's features is not sufficient for good performance in deep networks.Our main contributions are as follows: 1) We compare the performance of RSNs with different weight distributions and traditional wavelet scattering transforms, finding similar performance across weight initializations and improvements over scattering in most cases. 2) We build general randomized CNNs with learnable covariances, separating the importance of spatial and channel covariances. These layers can be easily swapped in for convolutional layers in any architecture, and they reach comparable performance with fewer learnable parameters, e.g. 18x fewer learnable parameters leads to only 3% reduction in accuracy on our widest ResNet-18. 3) We test two major modifications to the rainbow alignment procedure which better stabilize calculations over layers and avoid the aforementioned representation collapse problem. Our central finding is that learning covariance matrices within RFNs offers a path to achieve similar accuracy to traditional ANNs while simultaneously bridging connections to theory and neuroscience. 
related work: 
 A number of previous works have used correlated Gaussian structure in multilayer networks. Jacobsen et al. (2016) showed that learnable weights structured from a fixed Gaussian derivative basis improve performance over unstructured weights in limited dataset sizes and show improvements against wavelet scattering on large datasets. Feinman & Lake (2019) used a correlated multivariate Gaussian distribution to incorporate smooth and spatially correlated structure seen in V1 receptive fields into random convolutional weights, which improves generalization performance. Garriga-Alonso & van der Wilk (2021) found that spatial weight correlations are useful in deep networks, and Fortuin et al. (2022) show that using spatially correlated Gaussian priors in Bayesian neural networks improve performance compared to isotropic Gaussian priors. Pandey et al. (2022) explored the use of a correlated Gaussian model to generate random receptive fields inspired by neurons in the mouse primary visual cortex (V1). Structured random feature networks have fixed hidden weights sampled from a Gaussian process with a covariance modeled from experimental receptive field data space. The weight covariances strongly impact the performance of these networks by rotating and filtering the input before random projection. In single hidden layer networks, weights with V1-inspired covariance achieve faster training times and higher accuracies than uncorrelated weights.We also draw upon prior work related to wavelet scattering. Angles & Mallat (2018) find that embedding images using a wavelet scattering transform achieves strong performance in image generation without the use of a discriminator or encoder. Gauthier et al. (2022) show that learning the wavelet filters via parameterization can improve performance over traditional wavelet scattering filters. Li & Bonner (2022) find that the wavelet scattering transform is a strong and interpretable method for deep learning models related to V1. 
methods: 
 Our networks are multilayer CNNs, although nothing prevents our models from applying to fullyconnected or other layers. Layers are indexed as l = 0, . . . , L, where l = 0 and l = L refer to the input and output layers of the network, respectively. Convolution is assumed to operate on patches of size d 1 × d 2 containing d = d 1 • d 2 pixels per patch, n l-1 input channels and n l output channels. Each n l × n l-1 × d 1 × d 2 weight tensor is thus equivalent to an n l × (n l-1 d) matrix W l . Layer L is always a standard linear readout layer trained with standard methods.We build networks with different kinds of uncorrelated and correlated random weights. Uncorrelated weights are drawn from either Gaussian or uniform distributions. The correlated weight model draws the weights onto a neuron at layer l as w ∼ N (0, C l ), where C l is a positive semi-definite (n l-1 d) × (n l-1 d) covariance matrix (Fig. 1 ). For V1-like weights, we use a model of receptive field data from mouse V1 (Pandey et al., 2022) . In this setting C V1 depends on two parameters, the receptive field size s and spatial bandwidth f . We ran a hyperparameter sweep over these parameters for an RSN trained on CIFAR10, finding the optimal values s = 2 and f = 0.1. We use these values across all V1-like initializations. The eigenbasis of V1-like random weights consists of Hermite wavelets.We consider just multilayer feedforward networks. If C l = R T l R l is the Cholesky square root of the covariance of the layer l weights W l , then the length n l feature vector z l satisfies the recursionEQUATIONwhere ρ(•) = max(•, 0). The Gaussian part G l is an n l × (n l-1 d) matrix of i.i.d. entries drawn from a univariate Gaussian that is not learned with backpropagation. We discuss the structured convolutional RFN kernel which can be applied to our networks in Appendix A.6.In an RFN, the weights at all layers except the readout are frozen. To soften this constraint, we allow the covariance C l but not the initialization G l to be learned. For CNNs, we factor R l into the separable product of two lower-dimensional matrices, ann l-1 × n l-1 channel factor R 1 and a d × d spatial factor R 2 1 , as shown in Fig. 1. This yields R l = R 1 ⊗ R 2 and C l = (R 1 ⊗ R 2 ) T (R 1 ⊗ R 2 ).We build factored learnable covariance layers to replace traditional Conv2d layers in PyTorch (Paszke et al., 2019) . These have R 1 and R 2 as parameters. In our parameterization, the diagonals of R 1 and R 2 are exponentiated and the lower triangular part is zero; these are initialized as the identity. The matrix G is initialized with kaiming normal default scaling. We test a V1-like initialization with R 2 from the Cholesky decomposition of C V1 and test networks where channel and spatial covariances (R 1 and R 2 ) are learned (LC and LS) or unlearnable (UC and US).From a multilayer network with fixed covariance structure, we create a new network where the random matrices G ′ are resampled. Our goal is to see whether the new untrained network can recover the performance of the trained network purely through alignment, defined as follows: Given two random vectors ϕ and ϕ ′ , we align by solving the orthogonal Procrustes problem minA:A T A=In E∥Aϕ ′ -ϕ∥ 2 F ,which leads to A = EF T := A(ϕ, ϕ ′ ) where E and F are the orthogonal matrices of left and right singular vectors of the cross-correlation matrix E[ϕ(ϕ ′ ) T ]. If the features are deterministic matrices Φ and Φ ′ , then using the SVD of Φ(Φ ′ ) T minimizes ∥AΦ ′ -Φ∥ 2 F . We consider three different kinds of alignment: activation alignment (AA, ϕ = z), and weight alignment on the input (IWA, Φ = G T ) and output (OWA, Φ = G) dimensions. Guth et al. (2023) studied AA but for reference networks without enforced Gaussian weight structure. In all cases, we align earlier before later layers and adapt batchnorm statistics in-between. The final layer weights are copied from the reference model:W ′ L = W L .In the resampled network, the weightsW ′ l = G ′ l R l and features z ′ l = ρ(W ′ l z ′ l-1). To align activations, we average over both batches and spatial locations when computing the cross-correlation Guth et al., 2023) . With some abuse of notation, we will still write A = A(z l , z ′ l ), but the n l × n l alignment acts on only the channel dimension asE[z l (z ′ l ) T ] (as in(A ⊗ I d )z ′ l .In our code this is implemented as a forward pre hook to layer l + 1. Equivalently, we may absorb the alignment into R l+1 of the next layer:EQUATIONwhereR l+1 (A ⊗ I d ) = R 1,l+1 A ⊗ R 2,l+1. By taking the Cholesky decomposition of (R 1,l+1 A) T R 1,l+1 A, a new upper-triangular matrix R1,l+1 can be stored to make an equivalent factored covariance model.Given a resampled matrix G ′ and reference G, we also consider aligning these matrices directly. IWA does this on the input side of the weights: Ĝ = G ′ A, with A = A(G T , G ′T ), acting on the (n l-1 d)-dimensional row-space of the initialization matrices. On the other hand, OWA treats the output side of the weights: Ĝ = AG ′ , with A = A(G, G ′ ) acting on n l -dimensional space. In either case, we directly replace G ′ with Ĝ in the corresponding layer.RSNs are CNNs with random feature layers in the same architecture as a wavelet scattering network. We test an RSN composed of two convolutional layers and compare these networks against 2-scale wavelet scattering networks from Kymatio (Andreux et al., 2020) . We build versions of these networks (RSN-Fact) where the convolutional layers are replaced by factored covariance layers. Additional details of the architecture and parameter counts are given in Appendix A.1. To explore the effect of our structured weights on deeper networks, we also test ResNet-18 models (He et al., 2015) . We replace the convolutional layers with factored covariance layers (ResNet-Fact), disabling the learning of R 2 in the 1 × 1 skip connection layers.We run image classification experiments on the CIFAR10 and CIFAR100 datasets (Krizhevsky, 2009) . In RFNs, all weights except the readout are kept at their initialization. In all other models except traditional networks, the weight initializations G l are fixed while the R l are learned. To test performance in the limited data regime, we evaluate our models on CIFAR10 with small sample size (CIFAR10-S) using only 50 training samples per class. To understand variance due to initialization and training process, we train ten different networks on each dataset and report the average and standard error accuracy across seeds. We run experiments with RSN models initialized with V1-like weights, Gaussian weights, and uniform weights, and a two-scale wavelet scattering network. We include results of our RSN-Fact and ResNet-Fact models with V1-like spatial covariance initialization.Additionally, we run unsupervised learning experiments comparing RSNs to the wavelet scattering embedding described in Angles & Mallat (2018) and see similar results across image generation tasks (Fig. 9 in Appendix A.5).We also run extensive resampling and alignment experiments for our various models and alignment methods (AA, IWA, OWA). After alignment, we then fine-tune the readout layer. We estimate a sampled accuracy and deviation using 5 different samples of models generated from this process, before finetuning the readout layer over ten epochs of the training set.The code for our experiments and to implement these networks and learning rules is available at https://github.com/glomerulus-lab/fact-conv/. 
learning covariance structure for image classification: 
 We train RSN, RSN-Trad, RSN-Fact, ResNet-Trad, and ResNet-Fact models on CIFAR10, CIFAR10-S, and CIFAR100. Visualizations of the eigenspectra for the learned spatial and channel covariance matrices in an RSN-Fact network are shown in Fig. 2 (see also Fig. 7 in Appendix A.2).The leading spatial eigenvectors are relatively smooth and display characteristics of edge and blob detectors. The spectra of our spatial and channel matrices are effectively low-dimensional, with a few leading components accounting for most of the variance. These results are consistent across layers and initializations. Although the exact structure of the eigenvectors varies, similar patterns emerge across initializations (not shown).Image classification results are displayed in Fig. 3 for models with fixed features (left) and learnable features (right), including RSN, wavelet scattering, and ResNet architectures. Among the fixedfeature RSN models with different weight distributions, we see moderate improvement with V1-like initialization on CIFAR10. All of the RSN models outperform scattering except on CIFAR10-S, the small data setting where scattering is known to be powerful (Oyallon et al., 2019; Gauthier et al., 2022) .When features are learnable and we use lots of data (all settings except CIFAR10-S), all variations of the factored covariance models outperform the fixed-feature networks. More learning tends to improve performance for both RSN and ResNet. On the other hand, when in the small sample setting, bigger models and more learning do not lead to good performance due to overfitting. ResNet-Fact models performs worse than RSN-Fact on CIFAR10-S. RSN-USLC outperforms the other learnable models, but it is still a 3% worse than wavelet scattering. This hints that the other models may be overfitting spatial features. In the ResNet-Fact models, we see similar accuracies between the USLC and LSLC models among all three datasets, showing that the majority of the ResNet performance comes from learning the channel covariances.Furthermore, in CIFAR10 and CIFAR100, these models approach (but fall slightly short of) the performance of a traditional model with 17x fewer learned parameters. Full parameter counts are shown in beat the performance of IWA alone. Probe accuracies of AA and IWA degrade past layers 9 and 14, respectively.Finally, Fig. 5 shows the cosine similarity of the aligned and reference weights for a width-1 ResNet across layers. The different alignment methods lead to different aligned weights Ŵl = Ĝl Rl that may be compared to the reference model weights W l = G l R l ; studying the composite weights captures both weight and activation alignment effects. AA fails to align the weight matrices any more than a randomly resampled network without any alignment step: aligned and reference weights are nearly orthogonal for all layers. This is expected since the first layer weights are different. IWA, on the other hand, leads to high similarity with the reference model. Interestingly, IWA + AA has lower similarity than IWA alone but performs better in accuracy. OWA produces an intermediate value of similarity (not shown). 
discussion: 
 We use correlated Gaussian weights as a flexible model for multilayer networks where the weights are frozen or learned. This formulation is tested in Randomized Scattering Networks (RSNs) and ResNets. RSNs follow a similar architecture to the wavelet scattering transform and use fixed neuroinspired filters as feature extractors. We find that RSNs can outperform wavelet scattering models and that biologically-inspired weights can lead to an advantage in supervised learning. To highlight the similarity between RSNs and wavelet scattering, we build image generation embeddings described in Angles & Mallat (2018) , seeing similar performance with both models (Appendix A.5).Our main finding is that introducing covariance learnability into RSN and ResNet models leads to improved performance. In particular, learning the channel covariance reaches within 5% of the performance of a traditional CNN of the same architecture. This is achieved with 17x fewer parameters than fully trained networks. In limited experiments, we find effectively low rank spatial and diagonal channel covariances, suggesting further parameter reductions to try.We experiment with activation alignment (AA), based on pre-activations to a layer, and input weight alignment (IWA), where we directly align the weight matrices. With both of these methods, our linear probe study shows a degradation of accuracy in later layers of ResNet-Fact. Combining these two methods of alignment avoids the degradation and comes close to approximating original performance at the cost of requiring direct access to the pretrained weights and destroying their randomness. With this alignment technique, sampled weight performance converges to finetuned performance at higher width, suggesting that finetuning the classifier head isn't necessary. Guth et al. (2023) studied AA applied to wide, moderately deep networks, where colored Gaussian models were fit to each layer of a traditionally trained network. Among other results, they found that AA alone can lead to good performance in 7-layer scattering networks. An important conclusion was that this would not affect the network kernel, since the weights at a given layer were independent once conditioned on the previous layer weights. Weight alignment breaks this kernel equivalence. There are many possible reasons for our contradictory results: narrower layers, different architecture, or the covariance factorization. However, supplemental experiments (Fig. 8 ) with rainbow sampling reveal similar discrepancies in ResNets. Probing readout accuracies from intermediate layers leads us to hypothesize that in deeper networks, more than AA is needed to achieve strong performance.Limitations to our work include the few architectures that were considered and emphasis on thumbnail image datasets. We only considered the ReLU nonlinearity; other choices could lead to other results. In particular, it is unclear how the signal processing ideas underlying our work apply to very different domains such as text processing.In future work, we plan to investigate further ways to compress networks using weight factorization. Further speedups may be possible using methods such as fast random projections (Le et al., 2013) .We also want to study the randomized scattering kernel, to strengthen our theoretical understanding, and more thoroughly analyze the optimal representations and dynamics of covariance learning.A.1 NETWORK ARCHITECTURESThe RSN models are composed of an input layer followed by two convolutional layers with (7 × 7) kernels, a hidden dimension of 100, a stride of 1, and padding set to 3. Each of these layers go through batch normalization and average pooling before being flattened, concatenated, and passed through a linear classifier. Fig. 6 models the architecture, which is similar to the scattering transform with scale J = 2 and input passed through convolution, nonlinearity, and averaging. We train all RSNs for 90 epochs using a stochastic gradient descent optimizer and update the learning rate by 0.2 every 20 epochs. Table 1 displays the number of learnable parameters in the feature and classifier layers of the RSN along with a two-scale wavelet scattering model, RSN-Fact models with learnable and unlearnable spatial and channel covariances, and a traditional version of RSN with learnable weights in the convolution layers.The ResNet models follow the architecture in He et al. (2015) . We train all ResNets for 200 epochs using stochastic gradient descent and update the learning rate using a cosine annealing scheduler. We explore the effect of width by applying a scale factor to the number of channels in the convolutional models. Table 2 displays the number of learnable parameters in variations of ResNet-Fact across widths ranging from 2 -3 to 2 2 . Notably, we see a 14-18x decrease in learnable parameters between LSLC and a fully trained ResNet, with the compression factor increasing with width.A.2 LEARNED SPATIAL EIGENVECTORS Fig. 7 displays the top six learned spatial eigenvectors in the second RSN convolutional layer (left) and the first ResNet convolutional layer (right). We observe less smoothness than in the first RSN layer (Fig. 2 ). The ResNet spatial filters are 3 × 3 so these eigenvectors are of a smaller shape, but the first could be interpreted as a centered blob detector. We compared eigenvectors across different random initializations, noting similar patterns across layers (not shown). The number of learnable parameters in ResNet-Fact models trained on CIFAR10 at different widths compared to traditional models. The compression ratio of LSLC learnable parameters relative to traditional increases with width; this ratio is higher for USUC, LSUC, and USLC models. We see less smoothness than the first-layer eigenvectors (Fig. 2 ). Right: The first layer of a ResNet-Fact. The ResNet is dominated by a 1x1 centered bump eigenfeature with significantly outlying eigenvalue.Table 3 shows the result of default versus V1-like initialization with ResNet-Fact and RSN-Fact models for the spatial covariance matrix R 2 . We see that the V1-like initialization may lead to slightly improved performance. In the ResNet models across both initializations, we see close performance between USLC and LSLC, indicating that learning channels is beneficial. However, in most of the RSN models (CIFAR10 and CIFAR100), we see that the LSUC models reach the closest performance to LSLC, indicating that spatial covariance may be more important in this architecture.A.4 NETWORK RESAMPLING Fig. 8 shows a supplemental experiment applying the method of Guth et al. (2023) to ResNet-18.In this scenario, we trained a ResNet-Trad model then at each layer fit a Gaussian covariance to the trained weights and resample new weights accordingly. After weight sampling, we apply the same alignment methods as studied in the main text. Our code was adapted from that provided at https://bonnerlab.github.io/ccn-tutorial/pages/analyzing_neural_ networks.html.Each linear probe (Alain & Bengio, 2018; Richter et al., 2021) 2023) applied to ResNet-18. Left: The accuracy of traditional ResNets and their resampled and aligned (dashed) and resampled, aligned, and fine-tuned (solid) counter parts over width. At larger widths we are unable to approximate the performance of larger width networks. Right: We inspect the linear probe performance on various methods of sampling a network from a pretrained convolutional network. AA accuracy declines after layer 11. Aligning a sampled weight with a corresponding pretrained network's layer weight, IWA, prevents this decline. IWA + AA gives closest performance to that of the originator pretrained network, but a large difference remains in later layers.To test our RSNs against wavelet scattering in the unsupervised learning domain, we build a scattering-based image generator described in Angles & Mallat (2018) . Their network provides a strong mathematical framework with similar properties to generative adversarial networks (GANs) and variational auto-encoders (VAEs). In their model, images are embedded with a wavelet scattering transform, then passed through an affine whitening layer. The whitened scattering representation is then inverted by the generator model, which is composed of 5 convolutional layers with upsampling, reflection padding, batch normalization, and ReLU nonlinearity.We investigate whether embedding the input images with fixed V1-like features achieves similar performance to the wavelet scattering embedding in three image generation areas: linear interpolation between two images, generating new images from Gaussian white noise, and reconstructing images from the training and test sets. We use the wavelet scattering package Kymatio (Andreux et al., 2020) to regenerate results from Angles & Mallat (2018) . In our RSN version of the wavelet scat-this can be evaluated as the kernel for unstructured features acting on inputs z l-1 (x) transformed as R l z l-1 (x) (Pandey et al., 2022) .For a network with convolutional structure, z l [p] = ρ(G l ΛU z l-1 [p]) where p is a patch of the representation at each layer. This implies a convolutional kernel representation (Mairal et al., 2014; Guth et al., 2023) : (4)k l (x, x ′ ) =Note that the structure of this function is additive (Harris, 2019) over patches. Concatenating features, as in our RSN, yields a kernel of the form:k(x, x ′ ) = x T x ′ + L l=1k l (x, x ′ ).(5)An interesting direction for future work is to study how the transformation induced by the covariance of the correlated weights (R l or equivalently the spectral decomposition of the covariance C l ) affects the RKHS associated to the network. Perhaps the techniques of Xiao (2021) Here, we do not explicitly show the dependence of R1 and R2 on the layer, although we allow this to vary. 
 
 Appendix:  



Reviews: review: title: title: Does not seem like a good fit to the goals of this workshop
main_review: main_review: This paper investigates how neuron-inspired deep network architectures might lead to better performance. I have two main concerns. First, this does not really seem relevant to the goals of the workshop (as I understand them). Second, the empirical performance of these methods (see Fig. 3) seems very underwhelming to me. Conventional ANNs with learned features perform considerably better than the proposed method. Moreover, the authors only use CIFAR-10 and CIFAR-100 benchmarks. Comparing performance on a more extensive benchmark (e.g. ImageNet) is necessary. It is also not clear to me that the V1-like random feature model performs better than other random feature models. Indeed if I understand it correctly, the V1-like random feature model was optimized based on an experiment in mice---a species with much lower visual acuity than, for example, nonhuman primates.

I am marking this review with a low confidence score since my primary interest is not in machine learning benchmarks.
strength_weakness: strength_weakness: N/A
title: title: Review of "Learning and Aligning Structured Random Feature Networks"
main_review: main_review: This paper proposes a class of wavelet scattering networks with structured covariances. This paper is clearly relevant and interesting to the audience of this workshop, as it combines features of models commonly used in neuroscience and in machine learning (I have in mind shallower LN models and rainbow networks, respectively). The experiments are clearly described and appear well-motivated. 

I have only a few, relatively minor comments: 

- In discussing related works, it would be useful to mention the line of work in theoretical machine learning on computing sharp asymptotics for the generalization performance of networks with random weights. I have in mind Schröder et al. ICML 2023, Bosch et al. 2023 (on arXiv), and Zavatone-Veth & Pehlevan NeurIPS 2023. The former two works consider nonlinear models with i.i.d. weights, while the later considers correlated Gaussian weights. Given appropriate Gaussian equivalence results, one could compute the asymptotic generalization error for one of the models considered here given knowledge of the adapted covariance spectrum. 

- In Figure 3, it might be clearer to use a shaded background to highlight which models have fixed feature maps. 

- Do you understand why the alignment for IWA+AA in Figure 5 fluctuates from layer to layer in such a pronounced way? Also, the label hear reads AA+IWA rather than IWA+AA. 

- How sensitive are the results presented to the choice of ReLU nonlinearity? 

- For a broad audience, it would be useful to discuss rainbow networks in greater depth.
strength_weakness: strength_weakness: The paper is reasonably clearly written, and is relevant to the workshop audience.

Comments: title: Revised version
comment: Thank you to both reviewers for your time looking at our work and for your helpful comments. We have revised the paper to hopefully address the main issues.

* In discussing related works, it would be useful to mention the line of work in theoretical machine learning on computing sharp asymptotics for the generalization performance of networks with random weights. I have in mind Schröder et al. ICML 2023, Bosch et al. 2023 (on arXiv), and Zavatone-Veth & Pehlevan NeurIPS 2023. The former two works consider nonlinear models with i.i.d. weights, while the later considers correlated Gaussian weights. Given appropriate Gaussian equivalence results, one could compute the asymptotic generalization error for one of the models considered here given knowledge of the adapted covariance spectrum.

We have added to end (appendix): “Additionally, we could apply methods from Bosch et al., Schr ̈oder et al., and Zavatone-Veth & Pehlevan, who computed sharp asymptotics for the generalization performance of networks with random weights. The former two works consider nonlinear models with i.i.d. weights, while the latter considers correlated Gaussian weights.”
* In Figure 3, it might be clearer to use a shaded background to highlight which models have fixed feature maps. 

Done
* Do you understand why the alignment for IWA+AA in Figure 5 fluctuates from layer to layer in such a pronounced way? Also, the label hear reads AA+IWA rather than IWA+AA.

We do not have a good explanation for the fluctuations from layer-to-layer. The label has been fixed, thanks.

* How sensitive are the results presented to the choice of ReLU nonlinearity?

We added this caveat to the limitations: “We only considered the ReLU nonlinearity; other choices could lead to other results.”
However, we would expect that the learnable covariance model would apply to other types of nonlinearities.

* For a broad audience, it would be useful to discuss rainbow networks in greater depth.

We added these sentences to the discussion of rainbow networks in the introduction: “These networks assume that the weight dependencies across layers are reduced to rotations of independent random feature matrices which align the input activations. Rainbow networks thus model the joint probability distribution of the weights of trained networks.”

* ImageNet?

Running on more diverse image datasets is left for future work.

-----------------
A minor error was caught in revisions: In our original paper, the ResNet-trad results in Fig 3 were a few percent lower than in the final version due to a mistake in the initialization. This has been fixed for the final version. We have also added a link to our code repository, although this remains in-progress.