summaries: This work provides the novel approach for searching flow architectures. Compared to the standard approaches used to find the best deep architecture in standard network networks this problem is more challenging due to the need of invertible transform and requirement for determinant of the Jacobian to be easy to calculate. To solve the problem the authors propose to apply the weighting for candidate transformations. In order to enforce invertible properties of the model the authors suggest to use the mixing distribution approach instead of mixing the base transformations. They formulate the problem of learning best weights and show that optimal solution for soft weights is not optimal for binarised versions. Therefore they propose to optimize the upper bound of the proposed loss function instead. Further they show how to deal with the problem with larger number of layers. Some experiments are also performed to show the quality of the approach with respect to the baseline that is expert-based selection.

summaries: In this paper, the authors proposed to adapt a differentiable architecture search formulation (Liu et al., 2019) based on learned weighting of an ensemble of modules to automated search for Normalizing Flow architectures. The authors made several adaptations to the original approach for the normalizing flow problem, due to the invertibility constraints that prevent direct linear summation of different transform operations. Furthermore, the authors proposed to optimize the full network using an approximated upper bound of the KL divergence, instead of directly optimization. The authors proposed two methods to decompose the optimization problem: grow method, which is more straightforward and greedy, and block method, that alternatively adjusts each block. The authors experimentally compared their proposed method with manually specified architectures across various datasets, including POWER, GAS, HERMASS, MINIBOONE and BSDS300. The results seem mixed, as the searched model outperforms the manual model in some contexts but not others.

summaries: This paper proposes a DARTS-like method for searching automated normalization flow models. Instead of directly using the output ensembles, which leads to infeasible flow models, this work proposed distribution mixture to guarantee that the supernet is always a valid flow model. The upper bound of the loss function is optimized jointly with resource constraints. Experiments on small-to-medium scale datasets valid the effectiveness of the proposed method.

summaries: In this paper, the authors present an automated normalizing flow(NF) architecture search method. The method employs a mixture distribution formulation that can construct an optimal flow model with n layers of transformations from the transformation set. Besides, the authors introduce a block-wise optimization method to deal with exponentially growing optimization complexity. In the experiment, the authors proved the effectiveness of the optimization method which via approximate upper bound. And AutoNF has a better performance-cost trade-off than hand-tuned SOTA flow models.

abstractText: Although various flow models based on different transformations have been proposed, there still lacks a quantitative analysis of performance-cost trade-offs between different flows as well as a systematic way of constructing the best flow architecture. To tackle this challenge, we present an automated normalizing flow (NF) architecture search method. Our method aims to find the optimal sequence of transformation layers from a given set of unique transformations with three folds. First, a mixed distribution is formulated to enable efficient architecture optimization originally on the discrete space without violating the invertibility of the resulting NF architecture. Second, the mixture NF is optimized with an approximate upper bound which has a more preferable global minimum. Third, a block-wise alternating optimization algorithm is proposed to ensure efficient architecture optimization of deep flow models.

authors:
Jens Behrmann
Will Grathwohl
Ricky T.Q. Chen
David Duvenaud
Joern-Henrik Jacobsen

title: Invertible residual networks

venue: Proceedings of the 36th International Conference on Machine Learning,

authors:
Laurent Dinh
David Krueger
Yoshua Bengio

title: Nice: Non-linear independent components estimation

authors:
Laurent Dinh
Jascha Sohl-Dickstein
Samy Bengio

title: Density estimation using real NVP

venue: In 5th International Conference on Learning Representations,

authors:
Dheeru Dua
Casey Graff

title: UCI machine learning repository, 2017

venue: URL http://archive. ics.uci.edu/ml

authors:
Conor Durkan
Artur Bekasov
Iain Murray
George Papamakarios

title: Neural spline flows

venue: Advances in Neural Information Processing Systems,

authors:
Will Grathwohl
Ricky T.Q. Chen
Jesse Bettencourt
Ilya Sutskever
David Duvenaud

title: FFJORD: free-form continuous dynamics for scalable reversible generative models

venue: In 7th International Conference on Learning Representations,

authors:
Chin-Wei Huang
David Krueger
Alexandre Lacoste
Aaron Courville

title: Neural autoregressive flows

venue: Proceedings of the 35th International Conference on Machine Learning,

authors:
Kirthevasan Kandasamy
Willie Neiswanger
Jeff Schneider
Barnabas Poczos
Eric P Xing

title: Neural architecture search with bayesian optimisation and optimal transport

authors:
Durk P Kingma
Prafulla Dhariwal

title: Glow: Generative flow with invertible 1x1 convolutions

authors:
Hanxiao Liu
Karen Simonyan
Oriol Vinyals
Chrisantha Fernando
Koray Kavukcuoglu

title: Hierarchical representations for efficient architecture search

venue: In International Conference on Learning Representations,

authors:
Hanxiao Liu
Karen Simonyan
Yiming Yang

title: DARTS: Differentiable architecture search

authors:
D. Martin
C. Fowlkes
D. Tal
J. Malik

title: A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics

venue: In Proceedings Eighth IEEE International Conference on Computer Vision. ICCV 2001,

authors:
Bogdan Mazoure
Thang Doan
Audrey Durand
Joelle Pineau
R Devon Hjelm

title: Leveraging exploration in off-policy algorithms via normalizing flows

venue: Proceedings of the Conference on Robot Learning,

authors:
Muhammad Ferjad Naeem
Seong Joon Oh
Youngjung Uh
Yunjey Choi
Jaejun Yoo

title: Reliable fidelity and diversity metrics for generative models

authors:
Renato Negrinho
Geoff Gordon

title: Deeparchitect: Automatically designing and training deep architectures, 2017

authors:
Junier Oliva
Avinava Dubey
Manzil Zaheer
Barnabas Poczos
Ruslan Salakhutdinov
Eric Xing
Jeff Schneider

title: Transformation autoregressive networks

authors:
George Papamakarios
Theo Pavlakou
Iain Murray

title: Masked autoregressive flow for density estimation

authors:
George Papamakarios
Eric Nalisnick
Danilo Jimenez Rezende
Shakir Mohamed
Balaji Lakshminarayanan

title: Normalizing flows for probabilistic modeling and inference, 2021

authors:
Esteban Real
Sherry Moore
Andrew Selle
Saurabh Saxena
Yutaka Leon Suematsu
Jie Tan
Quoc Le
Alex Kurakin

title: Large-scale evolution of image classifiers, 2017

authors:
Danilo Rezende
Shakir Mohamed

title: Variational inference with normalizing flows

venue: In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning,

authors:
Oren Rippel
Ryan Prescott Adams

title: High-dimensional probability estimation with deep density

authors:
Masanori Suganuma
Mete Ozay
Takayuki Okatani

title: Exploiting the potential of standard convolutional autoencoders for image restoration by evolutionary search, 2018

authors:
E.G. Tabak
Cristina V. Turner

title: A family of nonparametric density estimation algorithms

venue: Communications on Pure and Applied Mathematics,

authors:
Barret Zoph
Quoc V. Le

title: Neural architecture search with reinforcement learning

heading: 1 INTRODUCTION

text:
Normalizing flow (NF) is a probabilistic modeling tool that has been widely used in density estimation, generative models, and random sampling. Various flow models have been proposed in recent years to improve their expressive power. Discrete flow models are either built based on elementalwise monotonical functions, named autoregressive flow or coupling layers (Papamakarios et al., 2017), or built with transformations where the determinant of the flow can be easily calculated with matrix determinant lemma (Rezende & Mohamed, 2015). In the continuous flow family, the models are constructed by neural ODE (Grathwohl et al., 2019).
Despite the variety of flow models, there’s yet no perfect flow concerning the expressive power and the computation cost. The flow models with higher expressive power usually have higher computational costs in either forward and inverse pass. In contrast, flows which are fast to compute are not able to model rich distributions and are limited to simple applications. For instance, autoregressive flows (Papamakarios et al., 2017) are universal probability approximators but are D times slower to invert than forward calculation, where D is the dimension of the modeled random variable x (Papamakarios et al., 2021). Flows based on coupling layers (Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018) have an analytic one-pass inverse but are less expressive than their autoregressive counterparts. Other highly expressive NF models (Rezende & Mohamed, 2015; Behrmann et al., 2019) cannot provide an analytic inverses and relies on numerical optimizations.
For different applications, the optimal flow model can be drastically different, especially if the computation cost is taken into consideration. For generative models (Dinh et al., 2015; Kingma & Dhariwal, 2018), flows with the fast forward pass are preferable since the forward transformations need to be applied to every sample from the base distribution. For density estimation (Papamakarios et al., 2017; Rippel & Adams, 2013), flows with cheap inverse will prevail. For applications where flow is utilized as a co-trained kernel (Mazoure et al., 2020), the computation cost and performance trade-off are more important, i.e., having a fast model with relatively good performance. However, in the current body of work, the architecture designs of the flow models are all based on manual configuration and tuning. To this date, there is a lack of a systematic way that could automatically construct an optimal flow architecture with a preferred cost.
In this paper, we propose AutoNF, an automated method for normalizing flow architecture optimization. AutoNF has a better performance-cost trade-off than hand-tuned SOTA flow models based on a given set of transformations. Our approach employs a mixture distribution formulation that can search a large design space of different transformations while still satisfying the invertibility requirement of normalizing flow. The proposed mixture NF is optimized via approximate upper bound which provides a better optimization landscape for finding the desired flow architecture. Besides, to deal with exponentially growing optimization complexity, we introduce a block-wise optimization method to enable efficient optimization of deep flow models.

heading: 2 RELATED WORK

text:
Normalizing Flows: Various normalizing flow models have been proposed since the first concept in (Tabak & Turner, 2013). Current flow models can be classified into two categories: finite flows based on layer structure, and continuous flow based on neural ODE (Grathwohl et al., 2019). The finite flow family includes flows based on elemental-wise transformation (Papamakarios et al., 2017; Kingma & Dhariwal, 2018) and flows whose transformations are restricted to be contractive (Behrmann et al., 2019). In elemental-wise transformation flows, autoregressive flow and coupling layers are two major flavors and extensive work has been proposed to improve the expressive power of both flow models. In Huang et al. (2018), the dimension-wise scalar transformation is implemented by a sigmoid neural network, which increases the expressive power at the cost of being not analytically invertible. In Durkan et al. (2019), piecewise splines are used as drop-in replacement of affine or additive transformations (Dinh et al., 2015; 2017) and is the current SOTA flow model. Consequently many recent research efforts have been devoted to closing the gap of expressive power, albeit at the cost of more complex and expensive transformations. Moreover, there has been no quantitative trade-off analysis between the performance and cost among different flows.
Neural Architecture Search: Many algorithms have been proposed or applied for neural architecture search. For instance, reinforcement learning (Zoph & Le, 2017), genetic algorithm (Real et al., 2017; Suganuma et al., 2018; Liu et al., 2018), Monte Carlo tree search (Negrinho & Gordon, 2017) or Bayesian optimization (Kandasamy et al., 2018). However, these methods all face the challenge of optimizing on a large discrete space and can take thousand of GPU days to find a good architecture. To address this issue, DARTS (Liu et al., 2019) proposes to relax the search space from discrete to continuous and allows efficient differentiable architecture search with gradient method which could reduce the search time to a single GPU day while still producing the SOTA architecture. However, all current NAS methods focus on optimizing traditional neural network structures (CNN, RNN) and there has yet been any implementation on normalizing flow.
Necessity for the Trade-off Between Performance and Cost: Despite various transformations proposed in the literature, there is no perfect transformation with strong expressive power and low computational cost. Autoregressive flows have better expressive power, but the inverse computation cost grows linearly with data dimension. Coupling layers’ inverse calculation is as fast as the forward pass, but their expressive power is generally worse than autoregressive flow with the same element-wise transformation. Even in the same autoregressive flow or coupling layer family, flows with different element-wise transformations have different performance and computation costs. For instance, additive or affine coupling layers (Dinh et al., 2017; 2015) have very fast forward and inverse calculation with limited expressive power while the flow in (Durkan et al., 2019) are highly expressive but are more demanding on computation. In most applications, it is necessary to find the best performance while minimizing at least one specific component of the cost. Unfortunately, the current design of flow models is empirical and therefore cannot ensure the optimal trade-offs.

heading: 3 METHOD

text:
In this work, we aim to tackle the challenge of finding an optimal flow model for a given task via an automated architecture search algorithm.
Assumptions: In the remaining part of this paper, without losing generality, we assume that the transformation is properly modeled such that during the training process, only forward computation is needed. Under this assumption, when the flow model is used for density modeling (Durkan
et al., 2019), the forward calculation is the dominant computation. When the flow model is used for random sampling (Kingma & Dhariwal, 2018), the inverse calculation is computationally intensive. When the flow model is utilized as a module and trained together with other components, e.g., policy network in maximum entropy learning (Mazoure et al., 2020), the training cost of the flow model is an important consideration.
Problem Definition: Given a transformation set with m options {T 1, T 2, ...Tm}, the goal is to construct an optimal flow model with n layers of transformations from the set. The flow model pNF (x;θ) = pT1T2...Tn(x;θ) should minimize the KL divergence between the target distribution p∗(x) and itself while minimizing its computational cost CNF . Here, θ are the parameters of the transformation in the flow model. In this paper, we use the forward KL divergence as our target loss function (Papamakarios et al., 2021):
θ∗ =argmin θ {DKL[p∗(x) || pT1T2...Tn(x;θ)] + λ · CNF }
s.t. Ti ∈ {T 1, T 2, ...Tm} (1)
While λ is a tuning factor capturing the relative importance of the performance-cost trade-off. Finding this optimal flow model is a discrete optimization problem with exponential complexity. To enable efficient architecture optimization, we use proposed method of relaxing the discrete search space to continuous space as suggested in Liu et al. (2019).

heading: 3.1 MIXED FLOW ENSEMBLE

text:
For the ith transformation layer with m options, we introduce a corresponding weight w j i for each option T j which reflects how likely the transformation will be selected. The weight is parameterized by a vector α and made continuous via softmax:
wji = exp(αji )∑m j=1 exp(α j i )
(2)
By applying this parameterization for each transformation layer, we can construct a mixed flow ensemble pMix(x;θ,α), where each layer in this mixed model reflects a weighted combination of the effect of all possible transformations. In this case, the architecture optimization problem is reduced to learning the weight vector for each layer and, at the end of the optimization process, weights will be binarized and the transformation with the highest weight in one layer will be selected as the final transformation. The mixed flow ensemble thus degrades to a normal flow model. The whole procedure is illustrated in Fig. 1 (left).
As adopted in (Liu et al., 2019), training of the flow ensemble becomes joint optimization of the architecture parameterα and the model parameter θ over the training and validation datasets, which could be written as the following bi-level optimization problem:
α∗ =argmin α DvalKL[p ∗(x) || pMix(x;θ∗,α)] + λ · CMix(α)
s.t. θ∗ = argmin θ DtrainKL [p ∗(x) || pMix(x;θ,α)],
∀ T ∈ pMix, T ∈ {T 1, T 2, ...Tm},
(3)
While the optimization problem is well defined, the key challenge is to construct the flow ensemble within the normalizing flow framework. This is different from traditional neural architecture search, which can mix various operations with no additional issue. Normalizing flow has its unique requirement for the invertibility of transformations and a preferred simple Jacobian calculation, which requires careful handling.
The mixed flow ensemble pMix(x;θ∗,α) must satisfy two requirements. First, it must be a legal density function such that it can be optimized by the KL divergence formulation. Second, each transformation layer in pMix(x;θ∗,α) should represent a weighted combination of all possible transformations. Consider the ith layer in the mixed flow ensemble with input random variable xin and output random variable xout, and pxin(xin) and pxout(xout) are their corresponding density functions. This layer has m transformation options in {T 1i , T 2i , ...Tmi } and w j i is the corresponding
weight for each transformation. As discussed in Assumption, we assume all transformations directly model the inverse transformation, i.e. xin = T j i (xout). Two approaches can be used to construct the mixed flow ensemble.
Construction by Mixed Transformations: The straight forward way of building the ith mix flow ensemble layer is to mix all transformations by weighted summation, as shown in Fig. 1 (right-top). The final weighted transformation for this layer can be thus represented as:
Ti(xin) = m∑ j=1 wji · T j i (xout) (4)
There are two drawbacks of this formulation despite its simplicity. First, definition of normalizing flow requires the mixed transformation Ti be invertible and differentiable in order to ensure pxout(xout) legal density function. However, this invertibility is not guaranteed even if all candidate transformations are invertible. Second, even if the mixed transformation is invertible, there is no easy way to calculate the Jacobian determinant of this weighted summation of transformations. Meeting the requirement of invertibility and ease of calculating Jacobian determinant brings strict restrictions on the candidate transformations and prevents the optimization of flow architectures on a wider search space. As a result, the construction of the mixed flow ensemble by weighted summation of transformations is not adopted in this paper.
Construction by Mixed Distributions: An alternating way is to build the mixed flow ensemble by mixing distributions. For a given transformation T ji in this ith layer, applying the transformation to the input random variable will result in a new distribution:
pT ji (xout) = pxin(T j i (xout)) · | detJT ji (xout)| (5)
By applying this to every transformation option in {T 1i , T 2i , ...Tmi }, we can obtain k different distributions, and it is possible to mix all the density functions together by their weighted summation, to get a mixture model as shown in eq.(6).
pTi(xout) = m∑ j=1 wji · pT ji (xout) (6)
An illustration of this process is shown in Fig. 1 (right-bottom). Different from the previous approach, the mixture model has a legal density function as: pTi(xout). By the definition of normalizing flow, we can assume that there exists an invertible and differentiable transformation Ti, which transforms xin to xout, although the transformation itself can not be explicitly written out.
For the next (i + 1)th layer, the density of the mixture model will be used as the input density function pxin(xin) as in the previous layer. By applying this formulation for n layers, the final mixed flow ensemble can be written as:
pMix(x;θ,a) = mn∑ k=1 Wk · pT1T2...Tn(x,θ) = mn∑ k=1 Wk · pi(x;θi)
where each Wk = n∏
i=1 wi and mn∑ k Wk = 1
(7)
Each wi is defined in eq.(2) and we use pk(x;θk) to represent a “normal flow architecture” with n transformation layers. Clearly, the final mixed flow ensemble is a legal density function which is in fact, a weighted summation of all possible flow models built with n layers of transformations.

heading: 3.2 OPTIMIZATION WITH APPROXIMATED UPPER BOUND

text:
Optimizing the forward KL divergence between the target distribution and the mixed flow ensemble can be written as:
LOpMix = DKL [p ∗(x) || pMix(x;θ,α)]
= −Ep∗(x)[log( mn∑ k=1 Wk · pk(x;θk))] (8)
We will demonstrate that direct optimization of this original loss can lead to underside mixture models. In the whole search space of the flow ensemble, we are interested only in ”normal flow architectures” points, i.e. the points where the weight of one architecture is 1 and others are all 0. However, it can be easily proven that the global optimum of LOpMix may not be the desired normal flow architecture (the red points in Fig. 2). Instead, optimization is very likely to end up in a mixture model that is globally optimal with similar weight for each possible flow architecture (the green point in Fig. 2). In this case, we will encounter difficulty when extracting a normal flow architecture with the search result. A common practice in differentiable architecture search (Liu et al., 2019) is to binarize the weights and select corresponding transformations. However, there is no guarantee that the binarized architecture will have a lower loss, and finding this nearest binarization may lead to performance drop. As a result, optimization with the original loss function is not suitable, and could be risky.
In this paper, we propose to optimize an upper bound of the original loss function to provide a better global optimum for the search of best normal flow architectures. Our method utilizes Jensen’s inequality log( ∑ W · x) ≥ ∑ W · log(x) as follows, since we have ∑ W = 1 and the log function is concave, we can obtain an upper bound of the KL divergence given as:
LOpMix = −Ep∗(x)[log( mn∑ k Wk · pk(x;θk)] ≤ LUpMix = −Ep∗(x)[ mn∑ k Wk · log(pk(x;θk))] (9)
The benefit of optimizing the upper bound can be summarized as follows:
Proposition 1: The global minimum point of LUpMix is defined by a normal flow architecture.
Proof: Suppose each flow model pk(x;θk) has an optimal parameter θ∗k that minimizes the KL divergence between p∗(x) and it: −Ep∗(x)[log(pk(x;θ∗k)] ≤ −Ep∗(x)[log(pk(x;θk)] (10) There also exists a flow architecture (pz(x;θ∗z)) that has the minimal KL divergence: −Ep∗(x)[log(pz(x;θ∗z)] ≤ −Ep∗(x)[log(pk(x;θ∗k)], ∀k ∈ mn (11) We can then prove the proposition by showing that:
LUpMix = −Ep∗(x)[ mn∑ k Wk · log(pk(x;θk))] ≥ −Ep∗(x)[ mn∑ k Wk · log(pk(x;θ∗k))]
≥ −Ep∗(x)[ mn∑ k Wk · log(pz(x;θ∗z))] = −Ep∗(x)[log(pz(x;θ∗z)]
(12)
Proposition 2: At normal architecture points (Wk = 1,W−k = 0), LUpMix = L O pMix .
The proof of proposition 2 is apparent.
With the above propositions and under the assumption that the global optimum can be reached at the end of the optimization, we can show that the solution set, i.e. all possible normal flow architectures are the same in both LOpMix and L U pMix , and we can do optimization with proposed upper bound without violating the original definition. Furthermore, since the global optimum of the upper bound will always lead to a normal flow architecture, we will not end up in finding a mixture model with the need to do heuristic and risky binarization of weights W .

heading: 3.3 EFFICIENT ARCHITECTURE OPTIMIZATION FOR DEEP FLOW MODELS

text:
While the flow ensemble by mixed density formulation could reflect the weighted effect of all possible transformation combinations, the architecture optimization complexity grows exponentially with respect to the number of considered transformation types and the number of transformation layers. In this scenario, efficient optimization of the whole flow architecture will not be possible. It is natural to decompose the original problem into sequential optimization of few different blocks, where each block could be optimized in one time with a limited number of layers. We propose two methods to decompose the problem.
Grow Method: The first approach is a straightforward greedy method which we call ”Grow”. Each time, a block is optimized until convergence, and the weights of the transformation layer are binarized. The searched transformations in this block will be directly added to the searched layer in the previous block. The architecture optimization of later blocks will be based on the existing layers and, the growth of layers stops when reaching the total number of layers constraint. Despite its simplicity, the downside of the “Grow” method is that the optimization is short-sighted. The block being optimized has no information about the architectures which could be added later, and the whole architecture is more likely to be trapped in local minimum.
Block Method: To avoid the issue of getting stuck in a local minimum, we propose another method named “Block” optimization. Blocks B in this approach are optimized alternatively to allow each block to adjust their architectures with respect to other blocks. In fact, the first “Grow” approach is a specific case of the “Block” method, where all the blocks are initialized as identity transformations and optimized only once.
Algorithm 1 Algorithm flow for AutoNF Require: Transformations: {T 1, T 2, ...Tm}, Blocks: B = {B1, B2, ...Bl}, Cost: CMix Ensure: n-layer flow model:
1: while not converged do 2: for each Bi ∈B do 3: while not convergence do 4: αBi = argminαBi D val KL[p
∗(x) || pMix(x;θ∗B,αBi)] + λ · CMix(αBi) 5: θB = argminθB D train KL [p
∗(x) || pMix(x;θB,αBi)] 6: end while 7: Fix architecture for Bi 8: end for 9: end while

heading: 3.4 COST MODEL AND ALGORITHM FLOW

text:
As discussed in section II, we are interested in modeling the training cost (forward calculation cost) and the inverse calculation cost, since each of them plays a different role based on desired applications. We use an independent experiment to model the cost of different types of flows and summarized in a table which are included in Appendix B. With the cost model, the total cost of the mixed flow ensemble could be extracted based on emphasize on different costs, e.g. if training cost is the major concern, only training cost of different flows will be calculated. This total cost CMix is then added as an regularization term into the training loss function.
In our paper, gradient based method is used for optimization which is efficient in this very high dimensional search space. The architecture parameter α and the flow model parameter θ are optimized alternatively with first order approximation in (Liu et al., 2019). The final algorithm flow of our proposed AutoNF method can be summarized in Algorithm 1.

heading: 4 EXPERIMENTS

heading: 4.1 EVALUATION OF PROPOSED UPPER BOUND

text:
Setup: We use a simple example to demonstrate the necessity of doing optimization with our proposed upper bound. We use AutoNF to build a 4 layer flow model with 2 transformation options including planar flow and radial flow from (Rezende & Mohamed, 2015). We use the POWER dataset as the target and optimize with original loss (name M1) and our proposed upper bound (named M2). We use Adam optimizer for both architecture parameter and model parameter with a learning rate of 0.002. The batch size is 512 and the training iteration is 10000.
The results are shown in Fig.3. For both M1 and M2, we present the weight for planar and radial flow for each layer as well as the training and validation loss during the search process. The final weight for each layer, searched architectures after binarization and the test score are shown in the right-bottom table.
Analysis: Optimization with our proposed upper bound (M2) shows a concrete convergence of weight to 0 or 1 for each layer, which leads to a desired normal flow architecture, while the optimization with the original loss function (M1) ends up in a mixture model instead of a normal flow architecture, as shown in Fig.3(left). This is within in our expectation as shown in Fig.2. Moreover, although the mixture model is mostly likely to be the optimal in the original loss, the normal flow architecture after binarization however, is not an optimal model. As shown in the right-bottom table, the architecture found by M2 has a significantly better test score than M1, and this clearly supports our statement of doing optimization with our proposed upper bound.

heading: 4.2 SEARCH FOR FLOW MODELS WITH BEST PERFORMANCE COST TRADE-OFF

text:
Transformation Options and Reference Designs: To evaluate our AutoNF framework, we setup our experiments with four types of non-linear flows and one linear flow. In autoregressive family, we
choose affine autoregressive flow (Papamakarios et al., 2017) and rational quadratic autoregressive flow (Durkan et al., 2019). Affine autoregressive flow has limited expressive power but the computation cost is lower, while the later has the state of art performance in autoregressive family with higher cost. Affine coupling layer (Dinh et al., 2015) and rational quadratic coupling layer (Durkan et al., 2019) are selected from coupling layer family. For linear transformation, we combine a reverse permutation and an LU linear layer together as a single layer. Random permutation (Durkan et al., 2019; Oliva et al., 2018) is not used since it is difficult to reproduce in architecture optimization. Every non-linear transformation layer is paired with a linear transformation layer suggested by Durkan et al. (2019) as a final transformation option, i.e., a layer in our experiment contains a reverse permutation, an LU-linear layer and one of the non-linear transformation layer listed above.
We use the rational quadratic flows family, including rational quadratic autoregressive flow (RQ-AF) and Rational quadratic coupling layer (RQ-C) in (Durkan et al., 2019) which have top 2 performance as the baseline. For fair comparison, we use RQ-AF as the baseline when emphasizing forward cost since it has better performance and use RQ-C as the baseline when emphasizing inverse cost since RQ-C has significantly lower inverse cost.
Evaluation Metric and Datasets: Evaluating the performance-cost trade-off is an open question in NF, we propose to use a new metric to address the difficulty of negative log-likelyhood (NLL). NLL is a common measurement for density estimation (lower, the better), however, the order of magnitude of NLL is different across different datasets and it is not suitable to use percentage difference to measure how a model is exactly better than another.
In this paper, We proposed to utilize density and coverage (Naeem et al., 2020) to evaluate the performance of NF models. Density and coverage are recently proposed method to evaluate the sample quality of generative models. The density metric reflects the fidelity of the model and is consistent with NLL metric. Across different datasets, density and coverage are at the same order of magnitude and allows evaluation of architecture across datasets. In our experiments, 10000 samples are drawn from the trained flow models and compare with 10000 samples from the test data. The results of three independent runs are averaged as the final reported results.
To evaluate the performance-cost trade-off, we define a figure of merit (FOM) as FOM = cost reduction% + density drop% compared to reference SOTA designs. In principle, the weight of the two terms can be manually adjusted to reflect the importance. For demonstration purpose, we use the equally weighted summation to report the results.
The performance of the flow models are evaluated with density estimation for UCI (Dua & Graff, 2017) and BSDS300 (Martin et al., 2001) datasets.
Analysis: The architecture search results are reported in Table.1 which includes the test NLL, density, coverage, cost and corresponding FOM. Table.1 shows that our AutoNF clearly helps to find architectures that have better performance-cost trade-off. out AutoNF can reach to up to 3.66X cost reduction and up to 75.2% improvement in FOM compared with SOTA literature results. Across all five different datasets, AutoNF demonstrates an average improvement of 58.67% on FOM with emphasis on forward cost and an average improvement of 52.57% on FOM with emphasis on inverse cost.

heading: 5 DISCUSSION

text: Normalizing flow is highly parameterized module and designing a flow model and use it for application requires a lot of hands-on experience and domain knowledge. In this paper, we show that the AutoNF framework is very effective in balancing performance-cost trade-offs when building complex flow models. Moreover, although not demonstrated in this paper, the framework could also be used to help decide hyper parameters in complex flow model, e.g. the hidden features and number of bins in the SOTA coupling layer (Durkan et al., 2019). In additional, the proposed optimization method with upper bound can be easily extended to other suitable probabilistic kernels. one example is to identify the best parameterized distribution(s) within a mixture model. We believe our framework will be very useful in many machine learning applications where normalizing flows are needed.

title: AUTONF: AUTOMATED ARCHITECTURE OPTIMIZA-


summaries: This paper tries to give a measurement method to evaluate the learned model, which is highly correlated with the final test accuracy. The measurement method depends on two key factors: Intrinsic Dimension (ID) and CLuster Learnability (CL). This paper claims that the model with higher ID and CL performs better. Using ID and CL to predict top-1 accuracy can achieve a Pearson correlation coefficient of 0.93, which is better than existing predictors (e.g., alignment and uniformity). Inspired by the above observation, this paper proposes a modified DeepCluster algorithm to increase the final performance.

summaries: This paper proposes two metrics for assessing the quality of self-supervised learning representations in terms of expressiveness and learnability. Expressiveness expects to maximize the mutual information between the representation and the original data, while learnability emphasizes the representation suppose to be simpler and more learnable. Based on these two intuitions, the author applies Intrinsic Dimension (ID) and Cluster Learnability (CL) as the measurements of expressiveness and learnability respectively to predict downstream classification performance. The authors collect 30 checkpoints of recent self-supervised methods to validate the proposed metrics on two datasets, ImageNet and STL_10, and the results show a high correlation between the proposed metrics and the downstream classification performance.

summaries: This paper proposes Intrinsic Dimension (ID) and Cluster Learnability (CL) as an alternative solution to efficiently evaluate learned representation quality of pre-trained networks, without using down-stream tasks/labels. These two frameworks are inspired by finding “expressiveness” – finding the smallest number of variables to approximate the representation; and “learnability” – number of data samples needed for KNN clustering.  Extensive experiments have been evaluated on models trained with both supervised and unsupervised methods, and results showed that ID + CL would produce a high correlation with prediction accuracy from these methods, comparing to baselines. Finally, the authors argue that CL can also be incorporated into DeepCluster (an unsupervised method), as an auxiliary loss to further improve the prediction performance.

summaries: This paper proposes to use Cluster Learnability (CL) and Intrinsic Dimension (ID) to evaluate the representations learned by self-supervised learning methods. The authors collected 30 checkpoints and show that their method is more predictable when compared to other methods, e.g., alignment and uniformity. Moreover, the authors modified the labels generated by K-means in DeepCluster to improve the learnability of the representations and the results demonstrate slight improvements.

abstractText: We argue that representations induced by self-supervised learning (SSL) methods should both be expressive and learnable. To measure expressiveness, we propose to use the Intrinsic Dimension (ID) of the dataset in representation space. Inspired by the human study of Laina et al. (2020), we introduce Cluster Learnability (CL), defined in terms of the learning speed of a KNN classifier trained to predict Kmeans cluster labels for held-out representations. By collecting 30 state-of-art checkpoints, both supervised and self-supervised, using different architectures, we show that ID and CL can be combined to predict downstream classification performance better than the existing techniques based on contrastive losses or pretext tasks, while having no requirements on data augmentation, model architecture or human labels. To further demonstrate the utility of our framework, we propose modifying DeepCluster (Caron et al., 2018) to improve the learnability of the representations. Using our modification, we are able to outperform DeepCluster on both STL10 and ImageNet benchmarks. The performance of intermediate checkpoints can also be well predicted under our framework, suggesting the possibility of developing new SSL algorithms without labels.

authors:
Alessio Ansuini
Alessandro Laio
Jakob H Macke
Davide Zoccolan

title: Intrinsic dimension of data representations in deep neural networks

authors:
Philip Bachman
R Devon Hjelm
William Buchwalter

title: Learning representations by maximizing mutual information across views

authors:
Philip Bachman
R. Devon Hjelm
William Buchwalter

venue: In NeurIPS,

authors:
Clay Beckner
Janet B Pierrehumbert
Jennifer Hay

title: The emergence of linguistic structure in an online iterated learning task

venue: Journal of Language Evolution,

authors:
Mohamed Ishmael Belghazi
Aristide Baratin
Sai Rajeshwar
Sherjil Ozair
Yoshua Bengio
Aaron Courville
Devon Hjelm

title: Mutual information neural estimation

venue: In International Conference on Machine Learning,

authors:
Léonard Blier
Yann Ollivier

title: The description length of deep learning models

venue: arXiv preprint arXiv:1802.07044,

authors:
Mathilde Caron
Piotr Bojanowski
Armand Joulin
Matthijs Douze

title: Deep clustering for unsupervised learning of visual features

venue: In Proceedings of the European Conference on Computer Vision (ECCV),

authors:
Mathilde Caron
Ishan Misra
Julien Mairal
Priya Goyal
Piotr Bojanowski
Armand Joulin

title: Unsupervised learning of visual features by contrasting cluster assignments

venue: arXiv preprint arXiv:2006.09882,

authors:
Mathilde Caron
Hugo Touvron
Ishan Misra
Hervé Jégou
Julien Mairal
Piotr Bojanowski
Armand Joulin

title: Emerging properties in self-supervised vision transformers

venue: arXiv preprint arXiv:2104.14294,

title: Thinking About Godel And Turing: Essays On Complexity, 1970–2007

venue: World scientific,

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: arXiv preprint arXiv:2002.05709,

authors:
Ting Chen
Simon Kornblith
Kevin Swersky
Mohammad Norouzi
Geoffrey E Hinton

title: Big self-supervised models are strong semi-supervised learners

authors:
Xinlei Chen
Kaiming He

title: Exploring simple siamese representation learning

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Xinlei Chen
Haoqi Fan
Ross Girshick
Kaiming He

title: Improved baselines with momentum contrastive learning

venue: arXiv preprint arXiv:2003.04297,

title: Present position and potential developments: Some personal views statistical theory the prequential approach

venue: Journal of the Royal Statistical Society: Series A (General),

authors:
Weijian Deng
Liang Zheng

title: Are labels always necessary for classifier accuracy evaluation

authors:
Elena Facco
Maria d’Errico
Alex Rodriguez
Alessandro Laio

title: Estimating the intrinsic dimension of datasets by a minimal neighborhood information

venue: Scientific reports,

authors:
Thomas L Griffiths
Brian R Christian
Michael L Kalish

title: Revealing priors on category structures through iterated learning

venue: In Proceedings of the 28th annual conference of the Cognitive Science Society,

authors:
Jean-Bastien Grill
Florian Strub
Florent Altché
Corentin Tallec
Pierre Richemond
Elena Buchatskaya
Carl Doersch
Bernardo Avila Pires
Zhaohan Guo
Mohammad Gheshlaghi Azar

title: Bootstrap your own latent-a new approach to self-supervised learning

authors:
Shangmin Guo
Yi Ren
Serhii Havrylov
Stella Frank
Ivan Titov
Kenny Smith

title: The emergence of compositional languages for numeric concepts through iterated learning in neural agents

authors:
Kaiming He
Haoqi Fan
Yuxin Wu
Saining Xie
Ross Girshick

title: Momentum contrast for unsupervised visual representation learning

authors:
Irina Higgins
Loı̈c Matthey
Arka Pal
Christopher P. Burgess
Xavier Glorot
Matthew M. Botvinick
Shakir Mohamed
Alexander Lerchner

title: beta-vae: Learning basic visual concepts with a constrained variational framework

venue: In ICLR,

authors:
R Devon Hjelm
Alex Fedorov
Samuel Lavoie-Marchildon
Karan Grewal
Phil Bachman
Adam Trischler
Yoshua Bengio

title: Learning deep representations by mutual information estimation and maximization

venue: arXiv preprint arXiv:1808.06670,

title: The learnability of syntactic categories

venue: LOT Occasional Series,

authors:
Simon Kirby
Tom Griffiths
Kenny Smith

title: Iterated learning and the evolution of language

venue: Current opinion in neurobiology,

authors:
Iro Laina
Ruth Fong
Andrea Vedaldi

title: Quantifying learnability and describability of visual concepts emerging in representation learning

authors:
Elizaveta Levina
Peter Bickel

title: Maximum likelihood estimation of intrinsic dimension

authors:
Chunyuan Li
Heerad Farkhoor
Rosanne Liu
Jason Yosinski

title: Measuring the intrinsic dimension of objective landscapes

venue: arXiv preprint arXiv:1804.08838,

authors:
Fushan Li
Michael Bowling

title: Ease-of-teaching and language structure from emergent communication

venue: In Advances in Neural Information Processing Systems,

authors:
Junnan Li
Pan Zhou
Caiming Xiong
Steven Hoi

title: Prototypical contrastive learning of unsupervised representations

title: Self-organization in a perceptual network

authors:
Chenxi Liu
Piotr Dollár
Kaiming He
Ross Girshick
Alan Yuille
Saining Xie

title: Are labels necessary for neural architecture search

venue: In European Conference on Computer Vision,

authors:
Yuchen Lu
Soumye Singhal
Florian Strub
Olivier Pietquin
Aaron Courville

title: Countering language drift with seeded iterated learning

venue: arXiv preprint arXiv:2003.12694,

authors:
Ishan Misra
Laurens van der Maaten

title: Self-supervised learning of pretext-invariant representations

authors:
Aaron van den Oord
Yazhe Li
Oriol Vinyals

title: Representation learning with contrastive predictive coding, 2018

authors:
Karl W. Pettis
Thomas A. Bailey
Anil K. Jain
Richard C. Dubes

title: An intrinsic dimensionality estimator from near-neighbor information

venue: IEEE Transactions on Pattern Analysis and Machine Intelligence,

authors:
Phil Pope
Chen Zhu
Ahmed Abdelkader
Micah Goldblum
Tom Goldstein

title: The intrinsic dimension of images and its impact on learning

authors:
Eva Portelance
Michael C Frank
Dan Jurafsky
Alessandro Sordoni
Romain Laroche

title: The emergence of the shape bias results from communicative efficiency

venue: arXiv e-prints, pages arXiv–

authors:
Anna N Rafferty
Thomas L Griffiths
Marc Ettlinger

title: Exploring the relationship between learnability and linguistic universals

venue: In Proceedings of the 2nd Workshop on Cognitive Modeling and Computational Linguistics,

authors:
Colorado J Reed
Sean Metzger
Aravind Srinivas
Trevor Darrell
Kurt Keutzer

title: Selfaugment: Automatic augmentation policies for self-supervised learning

authors:
Yi Ren
Shangmin Guo
Matthieu Labeau
Shay B Cohen
Simon Kirby

title: Compositional languages emerge in a neural iterated learning model

venue: arXiv preprint arXiv:2002.01365,

authors:
Kenny Smith
Monica Tamariz
Simon Kirby

title: Linguistic structure is an evolutionary trade-off between simplicity and expressivity

venue: In Proceedings of the annual meeting of the cognitive science society,

authors:
Samarth Swarup
Les Gasser

title: Simple, but not too simple: Learnability vs. functionality in language evolution. In The Evolution Of Language, pages 299–306

venue: World Scientific,

authors:
Yonglong Tian
Chen Sun
Ben Poole
Dilip Krishnan
Cordelia Schmid
Phillip Isola

title: What makes for good views for contrastive learning

venue: arXiv preprint arXiv:2005.10243,

authors:
Michael Tschannen
Josip Djolonga
Paul K. Rubenstein
Sylvain Gelly
Mario Lucic

title: On mutual information maximization for representation learning

authors:
Ankit Vani
Max Schwarzer
Yuchen Lu
Eeshan Dhekane
Aaron Courville

title: Iterated learning for emergent systematicity in {vqa

authors:
P. Vincent
H. Larochelle
Y. Bengio
P.-A. Manzagol

title: Extracting and composing robust features with denoising autoencoders

venue: In International Conference on Machine Learning proceedings

authors:
Tongzhou Wang
Phillip Isola

title: Understanding contrastive representation learning through alignment and uniformity on the hypersphere

authors:
William F Whitney
Min Jae Song
David Brandfonbrener
Jaan Altosaar
Kyunghyun Cho

title: Evaluating representations by the complexity of learning low-loss predictors

venue: In Neural Compression: From Information Theory to Applications–Workshop@ ICLR 2021,

authors:
Zhirong Wu
Yuanjun Xiong
Stella X Yu
Dahua Lin

title: Unsupervised feature learning via nonparametric instance discrimination

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Jing Xu
Thomas Griffiths
Mike Dowman

title: Replicating color term universals through human iterated learning

venue: In Proceedings of the Annual Meeting of the Cognitive Science Society,

authors:
Fei Xue
Hongbing Ji
Wenbo Zhang
Yi Cao

title: Self-supervised video representation learning by maximizing mutual information

venue: Signal Processing: Image Communication,

authors:
Yuchen Zhou
Dan Yurovsky

title: A common framework for quantifying the learnability of nouns and verbs

heading: 1 INTRODUCTION

text:
We consider the unsupervised learning scenario, which aims at learning distributed representations for complex natural data without human annotations. Despite great recent progress in selfsupervised learning (Chen et al., 2020a; Caron et al., 2021; Grill et al., 2020; Caron et al., 2020; 2018; He et al., 2020; Chen and He, 2021), the problem of properly evaluating the quality of the learned representations has not been fully explored. We approach this question by drawing an analogy with the evolution of human language. Recent literature (Swarup and Gasser, 2008; Smith et al., 2013) suggests that linguistic structure arises as a trade-off between expressiveness and learnability. On one hand, expressiveness is the ability to discriminate objects and concepts of the world. On the other hand, Kirby et al. (2014) states that language evolution undergoes an iterated learning process, where each individual acquires language by observing a limited amount of language outputs of others (Chomsky, 1980). This process induces a form of compression, which makes the language simpler and more learnable for the next generation. Linguistic structure, e.g., compositionality, emerges between these opposing forces.
A similar guiding principle can be applied when learning representations of natural data. Many existing methods in self-supervised learning can be considered as attempts to increase expressiveness of the representations, i.e., by maximizing or controlling the mutual information between inputs and representations, either through reconstruction (Vincent et al., 2008; Higgins et al., 2017) or contrastive learning (Oord et al., 2018; Bachman et al., 2019b). We argue that in addition to maximizing expressiveness, a good representation should also be learnable, e.g. support generalization from few examples (Laina et al., 2020).
In this paper, we propose an evaluation framework for self-supervised learning relying on expressiveness and learnability. Compared with previous work (e.g., Wang and Isola, 2020), our framework is independent of (i) the availability of downstream task labels or predefined data augmentation techniques, and (ii) architecture choices and dimensionality of the latent representations, which makes it applicable to all previous SSL methods. We propose to measure expressiveness in terms estimators of the intrinsic dimension (ID) of the data representations (Ansuini et al., 2019). Taking inspiration
from Laina et al. (2020), we quantify learnability in terms of the learning speed (in terms of number of examples) of a KNN learner trained to predict the K-Means cluster labels for held out representations. We dub this method Cluster Learnability (CL). We show that high ID and CL correlate with high downstream performance for 30 different pretrained checkpoints across different architectures.
To further demonstrate the usefulness of our framework, we present a case-study in which we improve upon DeepCluster (Caron et al., 2018), which is a self-supervised learning method that iterates between clustering the representations and fitting the cluster labels. Inspired by our CL measure, we refine the cluster labels by executing KNN at the end of each K-means clustering round. The resulting method, KNN-DeepCluster, is able to outperform DeepCluster in both ImageNet and STL10. When visualizing the training trajectories on the CL-ID plane, we find that our method indeed improves learnability at convergence. The correlation between performance and learnability holds also for checkpoints obtained during training, which suggests the possibility of using our framework to guide the development of new SSL algorithm, without labels at hand.
Our contributions are summarized below:
• We analyze self-supervised learning through the lens of expressiveness and learnability, and we propose tractable way to estimate them with Intrinsic Dimension (ID) and Cluster Learnability (CL).
• We carry out a large-scale evaluation on ImageNet and STL10 with 30 state-of-the-art selfsupervised and supervised checkpoints. We show that CL-ID is more correlated with the Top1 accuracy than previous approaches, with a Pearson’s r of 0.93.
• We propose a modification to DeepCluster to increase its cluster learnability. The resulting algorithm, KNN-DeepCluster is able to outperform DeepCluster.

heading: 2 RELATED WORK

text:
Representation Evaluation Several recent works address the question of representation evaluation in self-supervised learning. Whitney et al. (2021) propose to use the learning dynamic of the downstream classifiers to measure the representation complexity, however the method depends on extra human labels. Wang and Isola (2020) decompose the contrastive loss into an alignment loss and a uniformity loss, which can be used to analyze the representation. However, their method requires access to the projection head output, limiting their applicability to general architecture. Pretext tasks like jigsaw, rotation prediction are shown to be well correlated with the supervised evaluation (Reed et al., 2021; Deng and Zheng, 2021) or architecture search (Liu et al., 2020). However, as is shown in our experiments, they only provide limited predictive power when applying to a large pretrained family of checkpoints including both supervised and self-supervised models. In addition, these methods depends on predefined data augmentation, which might be task specific.
Manifold Intrinsic Dimension Intrinsic dimension can be thought of as the smallest number of variable needed to approximate the representation manifold. Note that it is different from the intrinsic dimension of loss landscape (Li et al., 2018), which measures minimum dimension of parameter space needed to optimize. We build upon extensive work on estimating intrinsic dimension. Applying local neighborhood information to estimate the intrinsic dimension is not a new idea, and it is shown to be more superior to the global eigenvalue approach (Pettis et al., 1979b). Ansuini et al. (2019) apply the TwoNN estimator (Facco et al., 2017) to the non-linear representation manifold of modern deep neural nets. They find that the intrinsic dimension is inversely correlated to the classification accuracy, and that the deep layers have smaller intrinsic dimensions. Their work is further extended to confirm that natural images lies in a low-dimension manifold (Pope et al., 2021), and lower ID datasets leads to better generalization. Our work presents a more nuanced view about ID. It suggests that in order to learning representation from scratch, a least amount of ID is needed to convey the information from the dataset, e.g., a constant-value representation has an ID of 0, but no capacity to generalize at all. More discussion can be found in Section 5.
Learnability, Ease-of-Transmission and Compression Learnability has been argued to be a hallmark of the human language in order to be effortlessly transmitted through generations (Kirby et al., 2014; Rafferty et al., 2011; Beckner et al., 2017; Zhou and Yurovsky, 2021; Kampen, 2004), and
it is also true for visual concepts like color (Xu et al., 2010), categories (Griffiths et al., 2006), shapes (Portelance et al., 2021) etc. In deep learning, it has been explored in the context of emergent communication (Ren et al., 2020; Guo et al., 2019; Li and Bowling, 2019), language drift (Lu et al., 2020), and neural module networks (Vani et al., 2021), but it is less explored for vision representation learning, except for a human study on just two SSL methods Laina et al. (2020). Learnability has a tight connection to compression (Chaitin, 2007) and prequential codelength (Dawid, 1984), which quantifies the compression levels with the online learning error. Existing works (Blier and Ollivier, 2018) has use it to support the generalization ability of the learner (e.g., deep neural nets) on the dataset (e.g., labeled images). However, we use it to quantify the learnability of the representation, in the sense that if the emerged Kmeans clustering is more learnable, then the same KNN learner could achieve a lower compression bound via prequential coding.

heading: 3 PROPOSED FRAMEWORK

text: In this section, we introduce our setup and define our measures of expressiveness and learnability.

heading: 3.1 NOTATION

text:
We consider the following setting: we assume we have a dataset of images {xi}Ni=1 represented in a pixel space X ⊂ Rd, where xi ∼ P are sampled i.i.d from some natural image distribution P over X . We also consider representation mapsF : X → Rm, typically pretrained neural networks, which represent any image as an m-dimensional vector. Any such map forms a representation dataset Z = {zi}Ni=1 where zi = F(xi). Given zi ∈ Z and an integer k ≥ 1, we denote by NN(zi, k) the k-th nearest neighbor of zi. While we generally mean ‘nearest’ with respect to Euclidean distance, in practice we will also use the cosine distance function,1
D(z1, z2) = 2− 2 cos(z1, z2) (1) Let rik = D(zi, NN(zi, k)) be the distance of the data point zi to its k-th nearest neighbor.

heading: 3.2 INTRINSIC DIMENSION (ID)

text:
To measure expressiveness, we propose to use the notion of intrinsic dimensionality (ID) of the data in the representation space (Pettis et al., 1979a). Intuitively, as the representation becomes more expressive, we expect it to be able to represent images coming from more fine-grained categories, which could have multiple disentangled features such as color, shape, positions in space, etc. If these disentangled features correspond to different directions in some non-linear manifold, ID should yield a good surrogate for the expressiveness of the representation.
Inferring the intrinsic dimension of a highly nonlinear manifold is a challenging problem (e.g., Levina and Bickel, 2005). In this work, we leverage the nearest neighbor-based method of Facco et al. (2017) to estimate ID, which uses only the first two nearest neighbors of each point. This estimator (TwoNN) is shown to be reliable with respect to representation dimensions and scalable to real-world datasets with deep neural networks (Ansuini et al., 2019). Assuming the data density is approximately constant around data points up to the distance of their neighbors, this method exploits the simple relationship between the distribution of the ratio of first and second nearest neighbor distances and the intrinsic dimension.
Formally, let µi = ri2/ri1 be the ratio of distances for the data point zi. We suppose the points are sampled on a manifold with intrinsic dimension d. It can be shown that, under the above local uniformity assumption, µi follows a Pareto distribution with parameter d+ 1 on [1,∞). That is
P (µ1, · · · , µN |d) = dN N∏ i=1 µ −(d+1) i (2)
While d can be computed by maximizing the likelihood, we follow a much simpler method proposed by Ansuini et al. (2019) and estimate d with a linear regression on the empirical cumulate of the distribution of µ, which is shown to have good practical performance (Facco et al., 2017). This estimator is exact for uniformly distributed data. It is also shown to remain asymptotically correct for non-uniform sampling of zi when the image dataset is large enough (Ansuini et al., 2019).
1Both distances produce similar results in our experiments.

heading: 3.3 CLUSTER LEARNABILITY (CL)

text:
To measure the representation learnability, we take inspiration from Laina et al. (2020). In this work, it is shown that the classes obtained by clustering the representations of pretrained self-supervised models can be easily learned by humans. We propose a cluster learnability (CL) metric as a way to automate their procedure and measure the learnability of the representation.
Let {zi, ỹi}Ni=1 be the labelled dataset obtained by clustering the representation (e.g., with Kmeans). We want to define the learnability of the representation as the learning speed of a KNN classifier with this labelled dataset. To formalize this, we use a prequential approach (Dawid, 1984). Let ŷi = KNN(zi|{zk<i, ỹk<i}) denote the prediction of the i-th label after seeing the pairs (zk, ỹk) for k < i. Learnability is defined as the online learning accuracy of our KNN classifier:
CL = 1
N N∑ i=1 [ŷi == ỹi] (3)
This metric can be connected to the prequential codelength of the cluster assignment dataset where the conditional model is a formed by a KNN classifier. We discusss this further in Appendix B.

heading: 4 EXPERIMENT: EMPIRICAL EVALUATION OF THE CL-ID FRAMEWORK

text:
In this section, to assess the main hypothesis of the paper and the relevance of our metrics, we perform experiments designed to answer the following questions:
Q1: Are high-performing models consistently both expressive and learnable?
Q2: Can our CL-ID framework robustly predict the model performance better than the existing SSL evaluation frameworks?

heading: 4.1 SETUP

text:
We select 11 different self-supervised checkpoints and download the official ImageNet checkpoints. We include checkpoints with different training epochs and different architecture. Combining the supervised checkpoints results in 30 checkpoints in total. A complete list can be found in Table 3 in the appendix. We mainly use the KNN evaluation on the validation data using the ground-truth labels to measure the performance of the model, which has been shown to be well correlated with the linear evaluation but computationally less expensive (Caron et al., 2021). CL is computed on the train set to have a better estimate, while the ID is computed on the validation. For STL-10, we use the unlabeled split, which has more images.
In our experiments, we use 1000 clusters for ImageNet and 320 for STL-10, which are approximately the square root of the dataset size. We report results with 1 neighbor for our KNN learner. We normalize the features and use cosine distance for the K-means clustering and KNN learner. Other configurations of cluster numbers and neighbor numbers are also explored in section 4.4.

heading: 4.2 QUALITATIVE EVALUATION

text:
We conduct a qualitative evaluation (Figure 1) to address Q1. We choose Alignment vs. Uniformity (Align-Unif) (Wang and Isola, 2020) as our baseline, since this work also does not depend on downstream labels. We compute both alignment loss and uniformity loss on the validation data.
On both dataset, the proposed CL-ID framework show that good representations tend have high cluster CL as well as high ID, forming clusters on the upper-right Pareto frontier. In contrast, the Align-Unif fails to show the results with a lower alignment and higher uniformity for good representations. We hypothesize that this is because Align-Unif framework is derived from the contrastive loss and therefore requires access to the post-projection output. However, not every checkpoint has the projection layer. Our results suggest that the Align-Unif framework is less flexible than ours and might not be applicable to the raw outputs coming from a pretrained backbone.

heading: 4.3 QUANTITATIVE EVALUATION

text:
We also perform a quantitative evaluation (Figure 2) to address Q1 and Q2 by comparing the Top1 accuracy predictors from different framework. We come up with two Top1 accuracy predictors for our method:
Vanilla CL-ID Predictor : CL+ ID Linear CL-ID Predictor : K1 · CL+K2 · ID +K3
where K1,K2,K3 are obtained by linear regression with our collected 30 checkpoints.
The first baseline is linear Align-Unif Predictor, which is the linear regression of the alignment loss and uniformity loss. It is also learnt from the collected checkpoints.
Linear Align-Unif Predictor : K1 ·Align+K2 · Unif +K3.
The second baseline is rotation prediction, which is one of the best pretext task proposed proposed by Reed et al. (2021). We randomly rotate the training images with 0, 90, 180, 270 degrees, train
a KNN classifier on the train images to predict a 4-way classification, and then report the rotation prediction accuracy on the validation images.
Our results show that, for both datasets, our linear predictor beats our vanilla one, which already beats all the other methods in terms of the correlation scores. While the linear Align-Unif predictor achieves about 0.5 correlation, the learnt parameters K1,K2 are negative, which support our observation in Figure 1 that alignment-uniformity framework does not produce qualitatively correct results. The pretext tasks like rotation prediction have certain correlation, but it is not as good as our methods when applied to a wider range of SSL methods and it is data augmentation heavy.
We also tried combining alignment loss, uniformity loss and rotation prediction into one linear regression formula, but that still does not beat our predictors.

heading: 4.4 ROBUSTNESS ANALYSIS

text:
In this experiment, we want check whether our results in the previous section is robust (Q2). In our framework, ID does not depend on any hyper-parameters. Therefore, we mainly examine the following hyper-parameters for CL: the number of clusters and the number of neighbors used in the KNN Learner. The results can be found in Figure 3. The resulting linear predictor can still produce a higher correlation than the rotation prediction for a wide range of the configurations we have tried, for both ImageNet and STL-10.
On both datasets, we observe that in general decreasing neighbor number leading to a better predictor. We hypothesize that lowering the neighbor number would results in a stronger KNN classifier that over-fits the structure-less cluster assignments more easily, resulting in better distinction among different checkpoints. For practitioner, we recommend setting number of neighbors to be 1.
On both datasets, we observe that there is an optimal cluster numbers so that the choice of neighbor numbers can be more flexible. If the cluster number is too low, the results become worse. Recall that for different checkpoints, the clusters are produced with the same K-means algorithm, so the separability among clusters might be roughly controlled. As a result, we hypothesize that here the cluster numbers here might control the underlying difficulty of the classification problem faced by the KNN learner. In the extreme cases we can have either only 1 cluster or as many clusters as data points. In the former, KNN accuracy would always be 100% and in the latter always be close to 0%. As a result, neither of them would be suitable enough to distinguish the collected checkpoints. As a
rule of thumb, we recommend using a reasonably large number of clusters, e.g., the square root of the dataset size.

heading: 5 EXPERIMENT: DEEPCLUSTER WITH IMPROVED CLUSTER LEARNABILITY

text:
In this section, we propose a simple modification for DeepCluster (Caron et al., 2018), which increase the learnability of the representations. Our method is called KNNDeepcluster. The purpose of this experiment is not to focus on a new state-of-the-art SSL method, but rather to understand whether our framework can improve an existing SSL method. The experiments would address the following questions:
Q3: Does KNN-DeepCluster out-perform DeepCluster consistently? Does it end up having a higher CL?
Q4: Can we predict the performances of the new checkpoints with the CL-ID predictors from section 4.3?
DeepCluster alternates between running K-means on the representation space and learning a representation such that it is highly predictive of the target cluster label associated with it. As a result, if we can make target cluster label to be more learnable, we can also make the representation more learnable. Following our approach of CL described above , we will enforce the K-means labels to be consistent with the KNN prediction, We proceed in the following way: at the beginning of the epoch, after each K-means clustering round, we change the K-means labels to the leave-one-out nearest neighbor predictions2. The pseudocode of the approach can be found in Algo. 1.
2In practice, instead of performing a KNN prediction by querying the whole dataset, we use a queue of representations stored in each GPU worker. While this does not result in the exact neighbors, we find that it’s more scalable to large datasets.
Algorithm 1: KNNDeepCluster. X is input images and F is the model while Not converged do
Compute zi = F(xi; θ) for i = 1, ..., N Kmeans clustering on {zi}Ni=1 and get cluster label {ỹi}Ni=1. Get ŷi = KNNClassifier(zi|{zk, ỹk}k 6=i) for i = 1, ..., N . Leave one out Update θ on the dataset {xi, ŷi} for one epoch.
end

heading: 5.1 RESULTS ON IMAGENET AND STL-10

text: In our experiments, we set the number of neighbors to be 3 in our KNNDeepCluster for both datasets, and we find that this does not affect our results much. We perform experiments on both STL10 and ImageNet. Our results are shown in Table 1. There are consistent gains on both STL10 and ImageNet, but the gain on ImageNet is smaller, and only appears with a weaker downstream classifier such as a KNN or with less training data.

heading: 5.2 ANALYSIS WITH THE CL-ID FRAMEWORK

text:
We visualize the training trajectories in Figure 4. We observe that the training trajectories for both methods follow similar general patterns across datasets: in the initial stage, the models tend to decrease CL while increasing ID. After the model is properly warmed up, there is an increase of ID
along with CL. Finally, especially for the ImageNet models as well as KNN-DeepCluster in STL-10, ID hits the plateau and starts decreasing while CL keeps increasing. Despite the slight difference in the final stage between two datasets, we find that KNN-DeepCluster converges to a solution that has a higher CL in comparison to DeepCluster, without a significant drop in ID.
We also find that the relative Top-1 classification performance of the intermediate checkpoints can still be well predicted by our linear CL-ID predictor from Section 4.3, especially after the models are property warmed up. The lack of correlation in the early stage checkpoints could be caused by the fact that our linear CL-ID predictor is learnt from the collected 30 checkpoints, which are all about late stage (with Top1 Acc greater than 50%).
Previous work (e.g., Ansuini et al., 2019) claims that ID is inversely correlated to the Top-1 accuracy, seemingly countering our expressiveness argument. Our findings are more nuanced. For ImageNet, at least for the early and middle stages, the ID is increased (before epoch 119) so that, intuitively, the model could explore as much as possible to encode the information from dataset. Then in the later stage, a form of compression occurs while CL increases for better generalization – presumably this stage is the regime of analysis done in Ansuini et al. (2019).

heading: 6 CONCLUSION & DISCUSSION

text:
Inspired by an analogy with language emergencewe propose to analyze self-supervised representation learning through the lens of expressiveness and learnability. We propose to estimate expressiveness with TwoNN intrinsic dimension (ID), and learnability with the acquisition speed of a KNN learner on the K-means clustering of the representation. Our large-scale evaluation with 30 checkpoints supports our claim. Our method better predicts the top-1 accuracy than other evaluation frameworks, while making less assumptions on data distributions, training algorithms or underlying architectures. According to our method, we modify DeepCluster and our version outperforms it by achieving better cluster learnability on both ImageNet and STL-10. Visualizations show that our framework could predict the performance of the unseen intermediate checkpoints.
We believe our work is a solid step towards understanding the current SSL algorithms and opens up interesting research directions. Future works could further explore the expressiveness-learnability in a more theoretical context, extend this framework to other field like pretrained language models, as well as devise new SSL algorithms that directly maximize intrinsic dimension or cluster learnability.

heading: A FUTHER DISCUSSION ON CURRENT SSL PRACTICE

text:
We further discuss the motivation behind our framework, and how can approach some of the SSL practice throught the lense of expressiveness and learnability.
Contrastive Learning With our new view on representation quality by expressivenesslearnability, we can try to analyze some of the practice in the contrastive learning and see why it works and fails. Contrastive learning is a typical example of such a trade-off, and it has the following objective:
Ex,x+,x− log ( eF(x)TF(x+) eF(x)TF(x+) + eF(x)TF(x−) ) It aims to output distinguishable feature vectors between positive and negative examples which turns out to increase the expressiveness. Meanwhile, it aims to make the output vectors among positive examples to be similar by which a higher value of learnability is ensured, since ideally the another learner only need to see one example to generalize to other positive examples. As a result, the degree of learnability here can be controlled by the definition of positive examples. A more diverse selection of the positive examples should make theF leaning toward the learnability side, since each example could represent a larger fraction of the inputs. It is well acknowledged in the self-supervised community that a diverse data augmentation is an essential part of the contrastive learning, and under this framework, it is a source of the learnability pressure.
Clustering-based Learning Perhaps our expressiveness-learnability framework is more obvious in the clustering-based algorithm like DeepCluster (Caron et al., 2018) and SwAV (Caron et al., 2020). In these algorithms, you learn not only F but also a bank of cluster centroids C as well as assignments ỹ = CTF(x). As a result, we can also view these algorithms as finding ỹ = G(x), where G is composed of F and C. In these methods, the expressiveness is achieved by using a large number of clusters as well as some uniformity constraint, and this can ensure each data has a distinctive cluster assignment. SwAV achieves it with the SinkHorn iterations, while the DeepCluster, at least in the early version, would balance the sampling ratio of images to make the cluster uniform as well as resolving the empty cluster issues.
The learnability requirements are implicitly enforced as well. Since the space of ỹ is usually less than x, there are usually multiple data being assigned to the same cluster. This already ensure that a new learner should able to see a few examples to generalize within the cluster. In addition, these algorithms also leverage the data augmentation, so that G(x) and G(x+) should have the same ỹ, which further improves the learnability.

heading: B CLUSTER LEARNABILITY AND PREQUENTIAL CODELENGTH

text:
Note that Epn. 3 can be connected to the compression lower bound of the emerged cluster assignments datasets {zi, ỹi}Ni=1 via prequential coding (Dawid, 1984). The prequential codelength is defined as
Lpreq = − ∑ i log p(ỹi|(zk≤i, ỹk<i)) (4)
where p(ỹi|(zk≤i, ỹk<i)) is a conditional model, which can be computed from our KNN classifier probabilities. If Eqn. 3 is the online learning accuracy, then Eqn. 4 is online learning cross-entropy loss. Higher CL lead to lower prequential compression bounds on the emerged clusters.

heading: C OTHER ATTEMPTS ON CHARACTERIZING EXPRESSIVENESS AND LEARNABILITY

text:
In this section, we discussed our other unsuccessful attempts on characterizing expressiveness and learnability.
Previous works propose maximizing mutual information between inputs and the representations as a way of self-supervised learning (Bachman et al., 2019a; Xue et al., 2020; Hjelm et al., 2018; Linsker, 1988). However, Tschannen et al. (2020) argue that there is a more holistic picture behind their success story, beyond mutual information alone, e.g., the trade-off between the amount of information that can be stored against how hard it is to extract it in the downstream tasks.
We hypothesize that mutual information is indeed part of the solution and it corresponds to the expressiveness of representation. E.g., a one-to-one map between images and representations would have a high MI, while a degenerate solution would not. As a result, we tried using MINE (Belghazi et al., 2018) to compute a lower bound on the mutual information. We use ResNet18 as our MINE network and train for 5 epochs and report the estimate for validation set.
For learnability, we also try fitting a student ResNet18 from scratch to learn the map from images xi to representations zi produced by the pretrained checkpoints. Since the representation space is continuous, we experiment with both L2 loss and cosine similarity. Furthermore, inspired by the cluster learnability, we propose to also propose to let the student network to predict emerged cluster labels ỹi from images xi. We also use ResNet18 as our student, train for 5 epochs, and report the validation metrics as the learnability.
The qualitative results can be found in Figure 5. We find all these attempts does not produce any meaningful signal. There many problems. Firstly estimating MI usng MINE is known to be sensitive to many things like learning rate and architecture of MINE network (Tschannen et al., 2020). Secondly the loss of directly mapping to the continuous vectors zi is heavily effected by dimensionality of the space. Thirdly, there is a architecture bias when we use ResNet18 as the students. ViT checkpoints in general performs poorer since it is too different from the ResNets. As a result, our proposed method CL and ID do not have these problem, mainly because we avoid training another neural network from scratch.

heading: D MORE RESULTS ON KNNDEEPCLUSTER

text:
We performed more experiments on STL-10 between KNNDeepCluster and DeepCluster under different number of prototypes. The results are in Table 2. We find that our method is able to consistently out-perform the baseline under different configurations.
To further understand the effect of replacing the K-means labels with the KNN predictions, we perform some clustering analysis on a pretrained DeepCluster checkpoint in Figure 6 on STL-10.
We find that after the initial K-means clustering there is an uneven distribution of cluster size, and larger clusters have lower Davies–Bouldin index (DB Index)are more separable from the smaller cluster. We find the reassignment with KNN prediction would punish the smaller clusters with larger DB index, and their cluster members are absorbed into large clusters. As a result, in the next epoch, the network pays more attention on fitting labels coming from large and well-separable clusters