Question:
represent the scientific article to find a short tldr

E.g.  New types of compute hardware in development and entering the market hold the promise of revolutionizing deep learning in a manner as profound as GPUs. However, existing software frameworks and training algorithms for deep learning have yet to evolve to fully leverage the capability of the new wave of silicon. In particular, models that exploit structured input via complex and instance-dependent control flow are difficult to accelerate using existing algorithms and hardware that typically rely on minibatching. We present an asynchronous model-parallel (AMP) training algorithm that is specifically motivated by training on networks of interconnected devices. Through an implementation on multi-core CPUs, we show that AMP training converges to the same accuracy as conventional synchronous training algorithms in a similar number of epochs, but utilizes the available hardware more efficiently, even for small minibatch sizes, ing in shorter overall training times. Our framework opens the door for scaling up a new class of deep learning models that cannot be efficiently trained today. An emerging category of neural networks show the common trait of reacting in dynamic and unique ways to properties of their input. Networks like tree-structured recursive neural networks BID35 BID36 and graph neural networks (GNNs) BID31 BID20 BID12 take structured data types as input and and execute a computation that depends on the input structure. This defies the moden GPU-driven paradigm of minibatch-based processing, and we refer to this new class of models with dynamic control flow as dynamic neural networks. The development of dynamic neural network frameworks -Chainer BID37, DyNet BID24, and PyTorch (PyTorch core team) -speaks to the importance of this class of models and highlights the challenge of how to make it easy for users to describe them. Yet there is another big challenge: how can we train these models efficiently?Managing minibatches to keep GPUs fully utilized is typically considered a user's responsibility in these dynamic frameworks (with the exception of DyNet's autobatching feature; see Sec. 7). This means that users have to think about how to change their data feeding pipeline or even the model itself to run efficiently on GPUs, rather spending time innovating to improve the model accuracy. What if we had a hypothetical device with low memory overhead that allows perfect scaling without batching; i.e., processing 1 item is simply 100x faster than processing 100 items? Recent work on FPGAs and other specialized hardware BID10 BID4 BID17 for deep learning encourages us to investigate this question. Our premises are 1. No batching is required for efficient processing.2. Each device may not have enough memory to hold the entire model (this is a realistic constraint for current memory systems that approach the perfect scaling we require)Based on these premises, we propose an asynchronous model-parallel (AMP) training algorithm. Our idea is illustrated in FIG0. We need model parallelism because each device may be too small to hold the entire model (premise 2). However, if we perform synchronous parameter updates following the full forward and backward propagations, the only way to increase device utilization is by pipelining multiple instances into the system (see e.g., . Pipeline parallelism with synchronous updates is at odds with convergence speed due to a decreased parameter update frequency; compare FIG0 (a) and (b).To overcome this problem, we propose asynchronous parameter updates that occur without global synchronization whenever a pre-specified number of gradients have been accumulated; see Fig. 1 (c). With this design we aim for both high device utilization and update frequency. In this setting, however, model parameters may be updated between the forward and the backward computation of an instance, introducing gradient "staleness". Despite staleness, we show that AMP training can converge fast with good hardware utilization. Specifically, our contributions are:• We present the asynchronous model parallel training algorithm for efficient distributed training of dynamic networks.• We present an intermediate representation (IR) with explicit constructs for branching and joining control flow that supports AMP training. Unlike previous work that considers static computation graphs for static control flow (e.g., Caffe), and dynamic computation graphs for dynamic control flow (e.g., Chainer), our IR encodes a static computation graph to execute dynamic control flow 1. This makes training easy to distribute and parallelize.• We show that our IR can readily encode replicas, a form of data parallelism (see Sec. 5). In addition, our IR includes operators for data aggregation, which recover a form of batching, enabling our methods to be applied even on hardware where batching is beneficial.• We implement AMP training on a multi-core CPU and empirically demonstrate that AMP training converges to similar accuracies as synchronous algorithms on a variety of dynamic neural network models including Tree RNN and gated graph neural networks (GGNN).In summary, our work demonstrates the benefits of AMP training and gives a novel way to design and deploy neural network libraries with dynamic control flow. In addition, we use our implementation to estimate the performance on a hypothetical device satisfying premises 1 & 2, with 1TFLOPS compute capability (see Appendix C). Together, these contributions open up new ways to scale up dynamic networks on interconnected compute devices. Below we highlight three models with dynamic control flow that will be studied in depth in this paper:Variable-length RNNs iterate over the tokens of variable-length sequences. Pseudo-code for a simple vanilla RNN is given in Figure 2. The linear (fully connected) layer and rectified linear unit (ReLU) can be substituted with a more sophisticated unit such as a gated recurrent unit BID8. Though each instance has a different length, it is possible to add padding to enable batching. However this may lead to significant redundant compute due to variability in sequence lengths. Isu inc dec inc(s) = s{t=s.t+1} dec(s) = s{t=s. Tree-structured neural networks are powerful models used for parsing of natural language and images, semantic representation, and sentiment analysis BID34 BID3 BID35 BID36 . They require evaluation of (potentially multiple) trees with shared parameters but different topology for each instance. Each tree structure is instance-specific and batching requires nontrivial planning BID21. A simple form of tree neural network performs a bottom up traversal of the instance, starting from an embedding of the leaves. At each level the values from the child nodes are concatenated and sent through a specialized unit (e.g. LSTM). The is then propagated further up the tree. Backpropagation over the tree structure is known as backpropagation through structure BID13. BID31 BID20 BID12 combine both the temporal recurrence and recurrence over the structure. GNNs can be seen as performing aggregation/distribution operations over a general graph structure with shared parameters. Apart from the models above, there exist many recently proposed models with flexible control flow (e.g. hierarchical memory networks BID5, neural programmer interpreters BID29, adaptive computation networks BID14 BID11, and sparsely-gated mixture of experts BID33, to which our framework can be applied. In AMP training, each node of a computation graph (including control flow nodes -see next section) is associated with a worker, which is an abstraction of a compute device. Neural network training and inference is carried out by message passing among workers following algorithm 1. All workers run in parallel without synchronization. Each message contains a payload (float or int tensor), as well as a state that includes the IDs of the source and sink nodes and a label indicating the type of message (forward, backward, or update). The state is used to keep track of algorithm and control flow information. For example, in a variable-length RNN the state also contains the instance identifier, the current position in the sequence, and the total sequence length for the instance. More generally, if the neural model use (possibly nested) loops, then the state for the messages that arrive to and are produced from nodes that logically belong in loop bodies will contain sets of loop counters that together with the instance id uniquely identify the messages throughout the course of the computation. When a worker receives a message labeled as forward (or backward) then it performs the operation of the node indicated by the sink node ID on the supplied payload. This produces one or more outgoing messages that are then enqueued into the queues of the workers hosting the next sink nodes in the computation graph. The final loss layer initiates backward propagation. If the message type is update, the worker will carry out weight updates on the sink node using gradients accumulated in the appropriate slot in the worker's local memory. Since both the operation weights and weight gradients can be stored locally on the worker then workers only need to communicate activations and activation gradients, which are typically an order of magnitude smaller than the weights. The update message is typically sent from the sink node to itself as part of the backward process but it can also be sent from a controller node to simulate synchronous pipelined training. There are two important details that are not fully spelled out in Algorithm 1. First, since the messages arrive asynchronously (and possibly out of order), any operation that has more than one parent nodes need to store the payload into its local cache until all the parents send the corresponding payloads. Thus output message(s) can only be produced when all the payloads become available. The cache needs to be able to distinguish payloads received from different parents and payloads with different instance ID, and different counters (all encoded in the message states). The same is true in the backward pass for a node with multiple child nodes. Second, op * denotes the adjoint operation of op and takes the backward message msg and potentially the forward message fwd msg stored in the cache. For a nonlinear activation node (e.g., ReLU), the node will not change the state of the message in the forward pass. Thus in the backward pass, the adjoint operation will just multiply the partial derivative of the activation function to the payload of the received backward message keeping the state unchanged. By contrast, an operation that only changes the state of the message in the forward pass (e.g., increment the loop counter) will reverse the change in the backward pass leaving the payload unchanged. In the experiments we vary two hyper parameters to control the effect of asynchrony: min update interval: determines the minimum number of gradients that a parameterized operation needs to accumulate before it can update its parameters (using the update message). The staleness of a gradient can be measured by the number of updates between the forward and backward computation that produces the gradient. Small min update interval may increase gradient staleness. On the other hand, large min update interval can reduce the variance of the gradient but can in very infrequent updates and also slow down convergence. max active keys: controls the maximum number of active instances that are in-flight at any point in time. By setting max active keys = 1 we restrict to single-instance processing, typically equivalent to synchronous training. More in-flight messages generally increase hardware utilization, but may also increase gradient staleness. Section 6 demonstrates the effects of these parameters in a multi-core CPU runtime. In our asynchronous execution model, the optimal assignment of N neural network computation graph nodes to W workers (referred to as affinitization) is in general a non-trivial scheduling problem. We investigated several heuristics for assigning affinities (for example k-hop coloring BID15 to ensure subsequent heavy operations were assigned different workers). However, we find that the following procedure achieves high performance in practice in our multi-core CPU implementation, and is adopted throughout our experiments for simplicity. We first partition the nodes into H'heavy' operations (namely matrix multiplies) and (N − H)'light' operations, and then balance the heavy nodes across the workers by affinitizing the h th with the (h mod W) th worker. Finally, the light operations are affinitized randomly among the rest of the workers. Note that in scenarios where communication is over a physical network affinitization will become more critical for high performance. Overview Computation graphs are expressed using a static intermediate representation (IR) that can be a compilation target for high-level libraries (e.g. TensorFlow or our own Python and C++ frontends), and can itself admit multiple backends (e.g. the multi-core CPU runtime in this paper, or a network of accelerators). Static means that the IR graph is instance-independent. Nevertheless, it can execute dynamic and instance-dependent control flow decisions, in a forward and backward manner, by storing instance-and iteration-dependent information as the computation evolves. Each IR node comes with a forward and a backward semantics. A model is specified by (i) an IR graph, and (ii) a specialized controller loop that pumps instances and other data (e.g. initial hidden states or labels), and is responsible for throttling asynchrony. In the rest of this section we discuss the most important IR nodes along with their operational semantics, and show how they are used in the example models. Payload transformations Parameterized payload transform (PPT) nodes can be used to encode, for instance, fully connected layers. They apply a transform in the forward pass, but also record the activation in order to use it to compute gradients in the backward pass. An activation is recorded by keying on the state of the message, allowing thus to process forward and backwards messages completely asynchronously, and -in the extreme case -out of order. A PPT node requires specification of the forward and the backward transformation. It may decide to independently apply accumulated gradients to update its parameters. For transformations that do not involve parameters (e.g. ReLUs) our IR offers a simpler non-parameterized payload transform. Loops, state, and control flow A condition node (Cond f) is parameterized by a function f that queries the state (but not the payload) of the incoming message and, based on the response, routes the input to one of the successor nodes. A join node (Phi) propagates the messages it receives from each of its ancestor nodes but records the origin (using the state of the message as the key) so that in the backward pass it can backpropagate them to the correct origin. An invertible state update node (Isu f f −1) is parameterized by two functions f and f −1 that operate on the state of a message, and satisfy f −1 (f (x)) = x. Typically these are loop counter update functions. Figure 2 shows how to encode an RNN. The loop at the heart of the RNN is implemented with Cond, Phi and Isu nodes. The controller pumps sequences in a lookup table (just another PPT layer), and our Ungroup node (to be described in the next section) generates a stream of tensors each corresponding to a single token, tagged with the current time-step (loop counter). For each forward message, the Isu node increments the time-step, and the conditional node tests whether the end of the sequence has been reached. Depending on the answer it either propagates the hidden state back to Phi, or pushes the hidden state to the final linear and loss layers. In backward mode, messages pass through the Isu (which decrements the time-step), and reach the Phi node. The Phi node will (based on information from the forward phase) either back-propagate to the Cond node, or to the controller to terminate. Hence the loop is executed in both the forward and backward direction. Aggregation and disaggregation Our IR offers several constructs for aggregation and disagreggation; for example RNN requires us to concatenate (Concat) hidden states and embeddings with matching timesteps and instance ids (states). We offer a construct for broadcasting (Bcast) a message to multiple successors. We offer an ungrouping construct (Ungroup) that ungroups a matrix and emits all ing messages tagged with an extra user-provided increasing loop counter. This allows us to insert the stream of token embeddings in the middle of the RNN loop in Figure 2. In backward mode Ungroup groups back all the incoming gradients. A simpler variant of Ungroup is Replicate, which replicates a message with an extra loop counter in the state. In backwards mode Replicate sums up all incoming gradients that correspond to the state without the extra loop counter. Figure 4(a) describes a GNN that combines aggregation on the structure of a graph instance with an outer loop. The controller pumps data that contain the feature embeddings of all nodes of an input graph. In addition it pumps in a map specifiying graph topology. The Distribute node uses that information along with the graph itself to create sub-matrices (here each corresponding to edge types) and pass them through linear (fully connected) layers. The Collect node collects and re-groups the based on graph topology. These nodes correspond to a form of dynamic partition and merging. Schematically the Distribute behaviour is given in FIG1. In backward mode, based on the control information received during the forward pass (gray line) re-groups the gradients and sums together those that correspond to the same index. The Collect operator is essentially symmetric. Pipelined model parallelism can often be augmented with forms of data parallelism. Consider the RNN in Fig. 2. The only heavy operation (Linear-1) in the body of the loop will act as a bottleneck for computation. One solution is to split the linear layer into smaller tiles and compute them in parallel. This is expressible in our IR but the linear operation needs to be large enough to benefit from tiling in this way. Another approach is to replicate the linear layer in full. This requires only minimal new machinery -we can replicate the linear layer and place the replicas inside Cond and Phi nodes as in FIG2 (b). Different instances or messages from the same instance but with different position in the sequence can be processed in an (pipeline-)parallel fashion by being sent to one of the replicas chosen by a random or deterministic function of the message state. To enable parameters to be shared among the replicas, we use infrequent end-of-epoch replica synchronization (averaging) that incurs negligible communication cost. We also tried more elaborate message-passing protocols for group synchronization, but found that infrequent global synchronization was sufficient for fast convergence. We evaluate AMPNet using the dynamic models introduced in Section 2. For completeness, we additionally consider a multi-layer perceptron (MLP) as an example of a static, batched network that AMPNet is not specifically designed to tackle. Brief details of the models and data sets in the experiments are presented below, and further details are given in Appendix B. Results Our asynchronous runtime is motivated by the promise of emerging hardware (e.g. FPGA accelerators) that fulfill the premises in section 1 and are well suited to dynamic neural network execution. Here we are primarily interested in how the performance of our runtime improves as we increase the degree of asynchrony (by varying max active keys) while keeping other factors fixed. The aim is to answer the question of whether AMP training is a promising direction for novel distributed hardware that deviates from the CPU/GPU batching paradigm. To answer this question using resources available today we run the AMPNet training using a multi-core CPU runtime where each worker is a hardware thread (see Appendix A). Additionally we forecast the performance on a hypothetical 1TFLOPS device satisfying our premises by replacing all computation nodes by configurable sleep nodes. This allows us to estimate the performance on a new hardware keeping the control flow decisions dynamic. See Appendix C.It is also interesting to compare how the raw CPU performance of the AMP runtime compares with existing frameworks (TensorFlow, TensorFlow Fold and DyNet) to see that our implementation is already competitive with state of the art methods even on CPUs that do not conform to our target hardware model. We provide additional analysis for each experiment below. On MNIST, TAB3 shows 3x speedup from synchrony (max active keys = 1) to asynchrony (max active keys = 4) in terms of throughput. This is almost ideal as the first three linear layers are the heaviest operations. The number of epochs to reach the target validation accuracy increases from 3 to 4 but the overall speedup in terms of the wall clock time is 2.2x. We have also compared AMP training against pipeline parallel training (FIG0). FIG3 (a) shows that while AMP training achieves 3x throughput gain already with max active keys = 4, pipeline parallelism can only achieve 2x (in fact from FIG0, 3m/(m + 2) is the best case for max active keys = m) and higher throughput at max active keys = 8 is achieved at the cost of the sharp increase in the number of epochs to convergence. The list reduction dataset demonstrates the power of replicas. As there is only one heavy operation (Linear-1, Figure 2), the speedup from asynchrony is mild (1.3x). However we get 2.5x and 3.5x speedup for 2 and 4 replicas, respectively, which is nearly ideal. Again, the # of epochs to convergence is not affected by increasing max active keys. The slowdown in convergence for 4 replicas is due to the increased effective minibatch size -also commonly observed in data parallel training. Next the sentiment tree-RNN dataset shows that our runtime is competitive without batching to TensorFlow Fold BID21 using dynamic batching of batch size 100. It is worth mentioning that our runtime allows us to specify different min update interval parameter for each parameterized operation. We set this parameter to 1000 for the embedding layer, which is initialized by Glove vectors, and 50 for all other layers. This reduced gradient staleness in the embedding layer. The QM9 dataset demonstrates that increasing asynchrony helps on real-world tasks with complex control flow, and our method outperforms an efficient TensorFlow implementation on CPUs. Finally, we have implemented the BiLSTM w/ char model in BID25 on Wikiner dataset BID26. Our preliminary implementation without any batching achieves 130 sentences/s at max active keys = 32 without any noticeable loss in accuracy (around 94 % after one epoch). This is competitive with DyNet's performance on the same machine (23 sentences/s without and 220 with autobatching, respectively); See Sec. B.5 for details. Asynchrony Finally we provide additional analysis on the effect of asynchrony. The degree of asynchrony is controlled by hyperparameters min update interval and max active keys. In FIG3 we use an 8-replica RNN model on the list reduction dataset to investigate how these parameters affect the data and time required to converge to 96% validation accuracy. We find, in analogy with minibatch size in traditional systems, that min update interval must neither be too large nor too small. Increasing max active keys (increasing asynchrony) monotonically increases performance when the number of keys is similar to the number of individually affinitized heavy operations in the model 8 in this case). Increasing max active keys significantly beyond this point produces diminishing returns. One approach to the task of training networks with instance dependent control flow is to define the computation graph dynamically per-instance. This approach is taken in Chainer BID37, DyNet BID24, and PyTorch (PyTorch core team). There are key challenges in accelerating this approach: Model parallelism would require the dynamically generated computation graph to be scheduled on-the-fly, and BLAS level parallelism would require operations to be batched on-the-fly. Automatic dynamic batching has been implemented in DyNet BID24, and is an interesting alternative to our asynchronous execution. Similar methods are used in TensorFlow Fold BID21. The basic idea is to inspect and merge together (by depth) the unrolled computation graphs of several instances to create batched BLAS operations. The effectiveness of automatic batching greatly depends on the model -for example, it would not perform well on random permutations of a sequence of operations. By contrast, our IR would very succinctly express and achieve pipeline parallelism using a static computation graph that is easy to distribute and optimize. Theano BID1 and TensorFlow (TF) can syntactically handle instance dependent control flow with abstractions for conditional execution (ifelse in Theano and cond in TF) and loops (scan and while loop, respectively); TF also provides higher-order functions, such as map, foldl, foldr, and scan. The main difference between AMPNet and the above frameworks is that AMPNet is streaming and asynchronous whereas Theano is non-streaming and synchronous. Although not designed for streaming, TF can support streaming programmatically as it exposes first-class queues, as well as data prefetching with so called input pipelines. In our IR, all the queuing is implicit and stream-based execution is the default. TF additionally does support static description of dynamic control flow and state update, but we depart from the classic dataflow architecture that TF follows BID2: First, instead of having nodes that represent mutable reference cells, we encapsulate the state with which a message should be processed through the graph in the message itself. Second, because we encapsulate algorithmic state in the messages, we do not introduce the notion of control dependencies (which can be used to impose a specific execution order on TF operations). Our choices complicate algorithmic state management from a programming point of view and make the task of designing a high-level compiler non-trivial, but allow every node to run asynchronously and independently without a scheduler and without the need for control messages: For example, nodes that dynamically take a control flow path or split the data simply consult the state of the incoming message, instead of having to accept additional control inputs. For "small" states (e.g. nested loop counters or edge and node ids) this might be preferable than out-of-band signaling. Our IR can implement loops by simply using state-update, conditional, and phi nodes, because the state accompanies the payload throughout its lifetime, whereas TF introduces specialized operators from timely dataflow BID23 ) to achieve the same effect. Asynchronous data parallel training BID28 BID9 BID7 ) is another popular approach to scale out optimization by removing synchronization (orthogonal to and combinable with model-parallel training). For example, convolutional layers are more amenable to data-parallel training than fully connected layers, because the weights are smaller than the activations. Moreover, when control flow differs per data instance, data parallelism is one way to get an effective minibatch size > 1, which may improve convergence by reducing variance. The impact of staleness on convergence BID28 and optimization dynamics BID22 have been studied for data parallelism. It would be interesting to extend those to our setting. BID16, like us, aim to to train different parts of a model in a decoupled or asynchronous manner. More precisely, their goal is to approximate a gradient with a synthetic gradient computed by a small neural network that is locally attached to each layer. Hence, the local gradient calculation becomes independent of other layers (except for the training of the gradient predictor network) and allows asynchronous parameter updates. This would be especially useful if the evaluation of the local network is cheaper than the computation of the real gradient; for example, if the computation of the real gradient required significant communication of forward/backward messages between devices. We have presented an asynchronous model-parallel SGD algorithm for distributed neural network training. We have described an IR and multi-core CPU runtime for models with irregular and/or instance-dependent control flow. Looking forward, we aim to deploy our system on specialized hardware. Equally importantly, we plan to build a compiler that automatically deduces the information to be placed in the states and generates state keying functions from a higher-level description of the models. By unlocking scalable distributed training of dynamic models, we hope to enable exploration of this class of models that are currently only on the horizon but may become more mainstream in the future. A AMPNET RUNTIME IMPLEMENTATIONWe have implemented an AMPNet runtime for multi-core CPUs. Our runtime spawns multiple workers each associated with a hardware thread and hosting one or more IR nodes -in a more general setting each worker corresponds to a compute device. To remain faithful to a distributed environment communication is only through message passing. Each worker is equipped with a multiple-producer single-consumer queue that can accept messages for any IR node hosted on that worker. The main worker loop periodically offloads messages from the concurrent queue to a worker-local priority queue that assigns higher priority to backward messages. Backward prioritization is designed for situations when multiple IR nodes with a dependency on the IR graph end up hosted on the same worker. As a consequence, backpropagation can complete faster and new instances can be pumped in by the controller. We dequeue the top message and invoke the forward or backward method of the target IR node. These methods may update internal IR node state (such as cache the state of the incoming message and wait for more messages) or post new forward or backward messages. How to update the parameters using the gradients is a configuration option that selects amongst a range of optimization algorithms. We have implemented runtime configuration options for selecting several well-known schemes such as (momentum-)SGD and Adam BID18, and for controlling the training hyper-parameters. We provide more details of the experiment and analysis in this section. All experiments were carried out on machines with 16 cores and 112 GB of RAM. The validation curves were averaged over at least 20 independent runs. The time/epoch to reach a target accuracy was calculated as median of the time an algorithm takes to reach the target accuracy over the repetitions. We found this approach to be more reliable than reporting the time/epoch when the averaged accuracy reaches the target. TAB4 shows both the training and validation throughputs we obtained with AMPNet and our TensorFlow baselines. We train a 4-layer perceptron with ReLUs and 784-dimensional hidden units on MNIST BID19. Both the AMP runtime and a TensorFlow baseline use SGD with learning rate 0.1 and batch size of 100.Figure 6(a) shows the validation accuracy vs. time, validation accuracy vs. epochs, and throughputs of synchronous and asynchronous versions of AMPNet as well as TensorFlow. The throughput greatly increases from synchronous (max active keys = 1) to asynchronous (max active keys = 4) while the speed of convergence (middle panel) is hardly affected for mild amount of asynchrony. Taking higher max active keys = 8 increase throughput only very little (because there is no more work) and seems to rather make the convergence more unstable. This is due to the fact that our current scheduler is greedy and pumps in a forward message whenever the first layer is unoccupied, which leads to large gradient staleness. Clearly a better scheduling will remove this sensitivity. We train a vanilla RNN to perform reduction operations on variable length lists of digits. Each training instance is a sequence of at most 10 tokens: The first token indicates which of 4 reduction operations 2 is to be performed, and the remaining tokens represent the list of digits. The output is the of the calculation rounded modulo 10. The dataset consists of 10 5 training and 10 4 validation instances. We present this task as a classification problem to a vanilla RNN with ReLU activation and a hidden dimension of 128. All parameterized operations are affinitized on individual workers. We bucket training instances into batches of 100 sequences (in the baseline and in AMPNet). Figure 6(b) shows the validation accuracy vs. time and the number of epochs, and throughputs of the methods we discussed in the main text on the list reduction dataset. We first notice that increasing the asynchrony from synchronous (max active keys=1) to max active keys = 4 and max active keys = 16 affects the convergence very little at least in average. However, there (d) Sentiment Tree Bank (max active keys = 16) (e) QM9 dataset is also very little speedup without introducing replicas as we discussed in the main text. Increasing the number of replicas increases the throughput almost linearly from 13k sequences/s (synchronous) to 35k sequences/s (2 replicas) and over 70k sequences/s (4 replicas). Convergence is almost unaffected for 2 replicas. This was rather surprising because the parameters of the replicas are only synchronized after each epoch as we described in Sec. 5. A slight slow-down in convergence can be noticed for 4 replicas. Since even max active keys = 16 has almost no effect on the convergence without replicas, this is not due to asynchrony. We also tried to synchronize more frequently but this did not help. Thus we believe that the slow-down is due to the increase in the effective minibatch size ing in reduced number of updates per epoch, which is commonly observed in data parallel training. We consider the sentiment classification dataset from BID35 consisting of binarized constituency parse trees of English sentences with sentiment labels at each node. Following Tai et al. BID36, we use 8,544 trees for training, 1,101 trees for validation, and 2,210 trees for testing. We use a Tree LSTM for this classification task based on the TensorFlow Fold BID21 benchmark model. Both the AMP and Fold models are trained following BID36 with the additional architectural modifications proposed by BID21; BID32. Furthermore, we split our Tree-LSTM cell into Leaf LSTM and Branch LSTM cells. This does not affect the expressiveness of the model because the LSTM cell receives either zero input (on branch) or zero hidden states (on leaves); i.e., the two cells do not share weights except for the bias parameters, which are learned independently in our implementation. We compare the time to reach 82 % fine grained (5 classes) accuracy (averaged over all the nodes) on the validation set. Figure 6(c) shows the averaged fine grained validation accuracy for the tree RNN model with different max active keyson the Stanford Sentiment Tree Bank dataset. Interestingly although TensorFlow Fold achieves higher throughput, AMPNet converges faster (in terms of the number of epochs). This speedup is mainly due to the fact that we are not batching and updating whenever we have accumulated 50 gradients (except for the lookup table node that updates every 1000 gradients); 50 gradients correspond to roughly 2 trees. The reason for the lower throughput compared to TensorFlow Fold is that we are only grouping the leaf operations and not the branch operations. Grouping the branch operations is possible by extending our IR nodes and we are actively working on it. Figure 6(d) shows the same information for fixed max active keys = 16 and different min update interval. We can see that as we increase min update interval from the originally used 50 to larger values, the peak of the validation accuracy shifts later and lower becoming closer to the curve obtained by TensorFlow Fold. This is consistent with the parallels between min update interval and minibatch size we drew in Section 6. The min update interval parameter has marginal influence on the training throughput. We study a real-world application for GNNs: prediction of organic molecule properties from structural formulae in the QM9 dataset BID30; BID27. GNNs have previously been applied to this task in BID12.We concentrate on prediction of the norm of a molecule's dipole moment using a regression layer build on the propagation model from BID20 (corresponding to the simplest setting in BID12). We use a hidden dimension of 100 and 4 propagation steps, initializing the graph nodes (atoms) following BID12. The molecules contain up to 29 atoms and in a TensorFlow baseline we bucket molecules into batches of 20 with atom counts differing by at most 1 within a batch. Following BID12, we report regression accuracies in multiples of a target accuracy from the chemistry community. Figure 6(e) shows that GGNN can tolerate relatively large max active keys = 16, and increased the throughput significantly from 300 graphs/s (synchronous) to 1797 graphs/s (see TAB4). We compare the performance of AMP training with DyNet with and without autobatching using the BiLSTM tagger w/ char model from the DyNet benchmark suite 3. This model consists of both character-level and word-level bidirectional LSTMs. The model uses a learnable word embedding for frequent words (more than 5 times in the corpus) and character-level bidirectional LSTMs for infrequent words. We use our Distribute and Collect nodes to dynamically route the messages depending on the word frequencies. Wikiner dataset 4 is a named entity recognition dataset extracted from Wikipedia by BID26. We use the training/validation split provided by the DyNet benchmark suite with 142,153 and 1,696 sentences, respectively. The is shown in FIG5. We achieve more than 3x in throughput from max active keys = 1 to max active keys = 32 without any noticeable loss in the validation accuracy after 2 or 3 epochs. The slight decrease in the validation accuracy after the third epoch is due to overfitting and it is not related to asynchrony. In order to estimate the performance of AMPNet on a hypothetical device with 1 TFLOPS compute capability, we replaced all fully connected layers in the network with a dummy operation that simply waits for a specified time. The dummy operation waits for 2 · k · d in · d out · 10 −12 seconds when the input is k × d in and the weight matrix is d in × d out for forward, backward, and gradient accumulation operation. It waits for d in · d out · 10 −12 seconds for weight update. In this way we can maintain all the data-dependent control decisions (e.g., sequence length) identical to the original network and also measure the real time spent for all the other operations. In order to calculate the time to reach target accuracy, we take the median number of epochs the original network required to reach the target accuracy and calculate the time as time = epochs · n train throughput train + n val throughput val, where n train and n val are the number of training and validation instances, respectively, and throughput train and throughput val are the throughput for training and validation, respectively. The are shown in TAB5. For the 4-way replicated RNN, we estimate roughly 260k instances/s, which is a 3.7x speedup compared to our CPU runtime. For tree RNN and GGSNN, we estimate milder 30 -70 % speedups mainly due to the fact that they have more complicated operations like Distribute and Collect, which we did not attempt to extrapolate the computation time because the implementation on a new hardware may be drastically different from the current CPU runtime. == Using asynchronous gradient updates to accelerate dynamic neural network training != We propose accelerating Batch Normalization (BN) through sampling less correlated data for reduction operations  with regular execution pattern, which achieves up to 2x and 20% speedup for BN itself and the overall training, respectively.

Question: This paper concerns dictionary learning, i.e., sparse coding, a fundamental representation learning problem. We show that a subgradient descent algorithm, with random initialization, can recover orthogonal dictionaries on a natural nonsmooth, nonconvex L1 minimization formulation of the problem, under mild statistical assumption on the data. This is in contrast to previous provable methods that require either expensive computation or delicate initialization schemes. Our analysis develops several tools for characterizing landscapes of nonsmooth functions, which might be of independent interest for provable training of deep networks with nonsmooth activations (e.g., ReLU), among other applications. Preliminary synthetic and real experiments corroborate our analysis and show that our algorithm works well empirically in recovering orthogonal dictionaries. Dictionary learning (DL), i.e., sparse coding, concerns the problem of learning compact representations, i.e., given data Y, one tries to find a representation basis A and coefficients X, so that Y ≈ AX where X is most sparse. DL has numerous applications especially in image processing and computer vision . When posed in analytical form, DL seeks a transformation Q such that QY is sparse; in this sense DL can be considered as an (extremely!) primitive "deep" network .Many heuristic algorithms have been proposed to solve DL since the seminal work of , most of them surprisingly effective in practice . However, understandings on when and how DL is solvable have only recently started to emerge. Under appropriate generating models on A and X, showed that complete (i.e., square, invertible) A can be recovered from Y, provided that X is ultra-sparse. Subsequent works BID0 BID1; BID4 provided similar guarantees for overcomplete (i.e. fat) A, again in the ultra-sparse regime. The latter methods are invariably based on nonconvex optimization with model-dependent initialization, rendering their practicality on real data questionable. The ensuing developments have focused on breaking the sparsity barrier and addressing the practicality issue. Convex relaxations based on the sum-of-squares (SOS) SDP hierarchy can recover overcomplete A when X has linear sparsity BID6; ), while incurring expensive computation (solving large-scale SDP's or large-scale tensor decomposition). By contrast, showed that complete A can be recovered in the linear sparsity regime by solving a certain nonconvex problem with arbitrary initialization. However, the second-order optimization method proposed there is still expensive. This problem is partially addressed by which proved that the first-order gradient descent with random initialization enjoys a similar performance guarantee. A standing barrier toward practicality is dealing with nonsmooth functions. To promote sparsity in the coefficients, the 1 norm is the function of choice in practical DL, as is common in modern signal processing and machine learning BID10: despite its nonsmoothness, this choice often admits highly scalable numerical methods, such as proximal gradient method and alternating directionThe reader is welcome to refer to our arXiv version for future updates.method . The analyses in; , however, focused on characterizing the algorithm-independent function landscape of a certain nonconvex formulation of DL, which takes a smooth surrogate to 1 to get around the nonsmoothness. The tactic smoothing there introduced substantial analysis difficulty, and broke the practical advantage of computing with the simple 1 function. In this paper, we show that working directly with a natural 1 norm formulation in neat analysis and a practical algorithm. We focus on the problem of learning orthogonal dictionaries: given data {y i} i∈ [m] generated as y i = Ax i, where A ∈ R n×n is a fixed unknown orthogonal matrix and each x i ∈ R n is an iid Bernoulli-Gaussian random vector with parameter θ ∈, recover A. This statistical model is the same as in previous works .Write Y. = [y 1, . . ., y m] and similarly X. = [x 1, . . ., x m]. We propose to recover A by solving the following nonconvex (due to the constraint), nonsmooth (due to the objective) optimization problem: DISPLAYFORM0 |q y i | subject to q 2 = 1.(1.1)Based on the statistical model, q Y = q AX has the highest sparsity when q is a column of A (up to sign) so that q A is 1-sparse. formalized this intuition and optimized the same objective as Eq. (1.1) with a q ∞ = 1 constraint, which only works when θ ∼ O(1/ √ n). worked with the sphere constraint but replaced the 1 objective with a smooth surrogate, introducing substantial analytical and computational deficiencies as alluded above. In constrast, we show that with sufficiently many samples, the optimization landscape of formulation (1.1) is benign with high probability (over the randomness of X), and a simple Riemannian subgradient descent algorithm can provably recover A in polynomial time. Theorem 1.1 (Main , informal version of Theorem 3.1). Assume θ ∈ [1/n, 1/2]. For m ≥ Ω(θ −2 n 4 log 4 n), the following holds with high probability: there exists a poly(m, −1)-time algorithm, which runs Riemannian subgradient descent on formulation (1.1) from at most O(n log n) independent, uniformly random initial points, and outputs a set of vectors {a 1, . . ., a n} such that up to permutation and sign change, a i − a i 2 ≤ for all i ∈ [n].In words, our algorithm works also in the linear sparsity regime, the same as established in; , at a lower sample complexity O(n 4) in contrast to the existing O(n 5.5) in. 1 As for the landscape, we show that (Theorems 3.4 and 3.6) each of the desired solutions {±a i} i∈ [n] is a local minimizer of formulation (1.1) with a sufficiently large basin of attraction so that a random initialization will land into one of the basins with at least constant probability. To obtain the , we integrate and develop elements from nonsmooth analysis (on Riemannian manifolds), set-valued analysis, and random set theory, which might be valuable to studying other nonconvex, nonsmooth optimization problems. Dictionary learning Besides the many sampled above, we highlight similarities of our to. Both propose first-order optimization methods with random initialization, and several quantities we work with in the proofs are the same. A defining difference is we work with the nonsmooth 1 objective directly, while built on the smoothed objective from. We put considerable emphasis on practicality: the subgradient of the nonsmooth objective is considerably cheaper to evaluate than that of the smooth objective in , and in the algorithm we use Euclidean projection rather than exponential mapping to remain feasible-again, the former is much lighter for computation. General nonsmooth analysis While nonsmooth analytic tools such as subdifferential for convex functions are now well received in machine learning and relevant communities, that for general functions are much less so. The Clarke subdifferential and relevant calculus developed for the family of locally Lipschitz functions seem to be particularly relevant, and cover several families of functions of interest, such as convex functions, differentiable functions, and many forms of composition (; BID3 BID5 . Remarkably, majority of the tools and can be generalized to locally Lipschitz functions on Riemannnian manifolds . Our formulation (1.1) is exactly optimization of a locally Lipschitz function (as it is convex) on a Riemannian manifold (the sphere). For simplicity, we try to avoid the full manifold language, nonetheless. Nonsmooth optimization on Riemannian manifolds or with constraints Equally remarkable is many of the smooth optimization techniques and convergence can be naturally adapted to optimization of locally Lipschitz functions on Riemannian manifolds (; ; ;). New optimization methods such as gradient sampling and variants have been invented to solve general nonsmooth problems BID8 BID5; ). Almost all available convergence pertain to only global convergence, which is too weak for our purpose. Our specific convergence analysis gives us a local convergence (Theorem 3.8).Nonsmooth landscape characterization Nonsmoothness is not a big optimization barrier if the problem is convex; here we review some recent work on analyzing nonconvex nonsmooth problems. study the regularized empirical risk minimization problem with nonsmooth regularizers and show of the type "all stationary points are within statistical error of ground truth" under certain restricted strong convexity of the smooth risk.; study the phase retrieval problem with 1 loss, characterizing its nonconvex nonsmooth landscape and providing efficient algorithms. There is a recent surge of work on analyzing one-hidden-layer ReLU networks, which are nonconvex and nonsmooth. Algorithm-independent characterizations of the landscape are mostly local and require strong initialization procedures , whereas stronger global can be established via designing new loss functions , relating to PDEs , or problem-dependent analysis of the SGD . Our provides an algorithm-independent chacaterization of the landscape of non-smooth dictionary learning, and is "almost global" in the sense that the initialization condition is satisifed by random initialization with high probability. Other nonsmooth problems in application Prevalence of nonsmooth problems in optimal control and economics is evident from all monographs on nonsmooth analysis (; BID3 BID5 . In modern machine learning and data analysis, nonsmooth functions are often taken to encode structural information (e.g., sparsity, low-rankness, quantization), or whenever robust estimation is desired. In deep learning, the optimization problem is nonsmooth when nonsmooth activations are in use, e.g., the popular ReLU. The technical ideas around nonsmooth analysis, set-valued analysis, and random set theory that we gather and develop here are particularly relevant to these applications. Problem setup Given an unknown orthogonal dictionary A = [a 1, . . ., a n] ∈ R n×n, we wish to recover A through m observations of the form The coefficient vectors x i are sampled from the Bernoulli-Gaussian distribution with parameter θ ∈, denoted as BG(θ): each entry x ij is independently drawn from a standard Gaussian with probability θ and zero otherwise. The Bernoulli-Gaussian is a good prototype distribution for sparse vectors, as x i will be on average θ-sparse. For any z ∼ iid Ber(θ), we let Ω denote the set of non-zero indices, which is a random set itself. DISPLAYFORM0 We assume that n ≥ 3 and θ ∈ [1/n, 1/2]. In particular, θ ≥ 1/n is to require that each x i has at least one non-zero entry on average. First-order geometry We will focus on the first-order geometry of the non-smooth objective Eq. (1.1): DISPLAYFORM1 In the whole Euclidean space R n, f is convex with sub-differential set DISPLAYFORM2 where sign(·) is the set-valued sign function (i.e. sign = [−1, 1]). As we minimize f subject to the constraint q 2 = 1, our problem is no longer convex. The Riemannian sub-differential of f on S n−1 is defined as : DISPLAYFORM3 A point q is stationary for problem Eq. (1.1) if 0 ∈ ∂ R f (q). We will not distinguish between local maxima and saddle points-we call a stationary point q a saddle point if there is a descent direction (i.e. direction along which the function is locally maximized at q).Set-valued analysis As the subdifferential is a set-valued mapping, analyzing it requires some setvalued analysis, which we briefly present here. The addition of two sets is defined as the Minkowski summation: X + Y = {x + y : x ∈ X, y ∈ Y}. The expectation of random sets is a straightforward extension of the Minkowski sum allowing any measurable "selection" procedure; for the concrete definition see . The Hausdorff distance between two sets is defined as DISPLAYFORM4 Basic properties about the Hausdorff distance are provided in Appendix A.1.Notations Bold small letters (e.g., x) are vectors and bold capitals are matrices (e.g., X). The dotted equality. = is for definition. For any positive integer k, [k]. = {1, . . ., k}. By default, · is the 2 norm if applied to a vector, and the operator norm if applied to a matrix. C and c or any indexed versions are reserved for universal constants that may change from place to place. We now state our main , the recovery guarantee for learning orthogonal dictionary by solving formulation (1.1). Theorem 3.1 (Recovering orthogonal dictionary via subgradient descent). Suppose we observe DISPLAYFORM0 samples in the dictionary learning problem and we desire an accuracy ∈ for recovering the dictionary. With probability at least 1 − exp −cmθ 3 n −3 log −3 m − exp (−c R/n), an algorithm which runs Riemannian subgradient descent R = C n log n times with independent random initializations on S n−1 outputs a set of vectors {a 1, . . ., a n} such that up to permutation and sign change, a i − a i 2 ≤ for all i ∈ [n]. The total number of subgradient descent iterations is bounded by DISPLAYFORM1 Here C, C, C, c, c > 0 are universal constants. At a high level, the proof of Theorem 3.1 consists of the following steps, which we elaborate throughout the rest of this section.1. Partition the sphere into 2n symmetric "good sets" and show certain directional gradient is strong on population objective E[f] inside the good sets (Section 3.1).2. Show that the same geometric properties carry over to the empirical objective f with high probability. This involves proving the uniform convergence of the subdifferential set ∂f to E [∂f] (Section 3.2). 3. Under the benign geometry, establish the convergence of Riemannian subgradient descent to one of {±a i : i ∈ [n]} when initialized in the corresponding "good set" (Section 3.3).4. Calling the randomly initialized optimization procedure O(n log n) times will recover all of {a 1, . . ., a n} with high probability, by a coupon collector's argument (Section 3.4).Scaling and rotating to identity Throughout the rest of this paper, we are going to assume WLOG that the dictionary is the identity matrix, i.e. A = I n, so that Y = X, f (q) = q X 1, and the goal is to find the standard basis vectors {±e 1, . . ., ±e n}. The case of a general orthogonal A can be reduced to this special case via rotating by A: q Y = q AX = (q) X where q = A q and applying the on q. We also scale the objective by π/2 for convenience of later analysis. We begin by characterizing the geometry of the expected objective E [f]. Recall that we have rotated A to be identity, so that we have DISPLAYFORM0 Minimizers and saddles of the population objective We begin by computing the function value and subdifferential set of the population objective and giving a complete characterization of its stationary points, i.e. local minimizers and saddles. Proposition 3.2 (Population objective value and gradient). We have DISPLAYFORM1 (3.5) DISPLAYFORM2 The case k = 1 corresponds to the 2n global minimizers q = ±e i, and all other values of k correspond to saddle points. A consequence of Proposition 3.3 is that the population objective has no "spurious local minima": each stationary point is either a global minimizer or a saddle point, though the problem itself is non-convex due to the constraint. Identifying 2n "good" subsets We now define 2n subsets on the sphere, each containing one of the global minimizers {±e i} and possessing benign geometry for both the population and empirical objective, following . For any ζ ∈ [0, ∞) and i ∈ [n] define DISPLAYFORM3 For points in S DISPLAYFORM4, the i-th index is larger than all other indices (in absolute value) by a multiplicative factor of ζ. In particular, for any point in these subsets, the largest index is unique, so by Proposition 3.3 all population saddle points are excluded from these 2n subsets. Intuitively, this partition can serve as a "tiebreaker": points in S (i+) ζ0 is closer to e i than all the other 2n − 1 signed basis vectors. Therefore, we hope that optimization algorithms initialized in this region could favor e i over the other standard basis vectors, which we are going to show is indeed the case. For simplicity, we are going to state our geometry in S (n+) ζ; by symmetry the will automatically carry over to all the other 2n − 1 subsets. Theorem 3.4 (Lower bound on directional subgradients). Fix any ζ 0 ∈. We have and all indices j = n such that q j = 0, DISPLAYFORM5, we have that DISPLAYFORM6 These lower bounds verify our intuition: points inside S (n+) ζ0have subgradients pointing towards e n, both in a coordinate-wise sense and a combined sense: the direction e n − q n q is exactly the tangent direction of the sphere at q that points towards e n. We now show that the benign geometry in Theorem 3.4 is carried onto the empirical objective f given sufficiently many samples, using a concentration argument. The key behind is the concentration of the empirical subdifferential set to the population subdifferential, where concentration is measured in the Hausdorff distance between sets. Proposition 3.5 (Uniform convergence of subdifferential). For any t ∈, when 10) with probability at least 1 − exp −cmθt 2 /log m, we have DISPLAYFORM0 DISPLAYFORM1 Here C, c ≥ 0 are universal constants. The concentration guarantees that the sub-differential set is close to its expectation given sufficiently many samples with high probability. Choosing an appropriate concentration level t, the lower bounds on the directional subgradients carry over to the empirical objective f, which we state in the following theorem. Theorem 3.6 (Directional subgradient lower bound, empirical objective). There exist universal constants C, c ≥ 0 so that the following holds: for all ζ 0 ∈, when m ≥ Cn 4 θ −2 ζ −2 0 log 2 (n/ζ 0), with probability at least 1 − exp −cmθ 3 ζ 2 0 n −3 log −1 m, the following properties hold simultaneously for all the 2n subsets S DISPLAYFORM2 and all j ∈ [n] with q j = 0 and q DISPLAYFORM3 The consequence of Theorem 3.6 is two-fold. First, it guarantees that the only possible stationary point of f in S (n+) ζ0 is e n: for every other point q = e n, property (b) guarantees that 0 / ∈ ∂ R f (q), therefore q is non-stationary. Second, the directional subgradient lower bounds allow us to establish convergence of the Riemannian subgradient descent algorithm, in a way similar to showing convergence of unconstrained gradient descent on star strongly convex functions. We now present an upper bound on the norm of the subdifferential sets, which is needed for the convergence analysis. Proposition 3.7. There exist universal constants C, c ≥ 0 such that sup ∂f (q) ≤ 2 ∀ q ∈ S n−1 (3.14) with probability at least 1 − exp −cmθ log −1 m, provided that m ≥ Cn log n. This particularly implies that DISPLAYFORM4 DISPLAYFORM5 Each iteration moves in an arbitrary Riemannian subgradient direction followed by a projection back onto the sphere. We show that the algorithm is guaranteed to find one basis as long as the initialization is in the "right" region. To give a concrete , we set ζ 0 = 1/(5 log n). Theorem 3.8 (One run of subgradient descent recovers one basis). Let m ≥ Cθ −2 n 4 log 4 n and ∈ (0, 2θ/25]. With probability at least 1 − exp −cmθ 3 n −3 log −3 m the following happens. If DISPLAYFORM0 1/(5 log n), and we run the projected Riemannian subgradient descent with step size η (k) = k −α /(100 √ n) with α ∈ (0, 1/2), and keep track of the best function value so far until after iterate K is performed, producing q best. Then, q best obeys 17) provided that DISPLAYFORM1 DISPLAYFORM2, 64 DISPLAYFORM3 In particular, choosing α = 3/8 < 1/2, it suffices to let DISPLAYFORM4 Here C, C, c ≥ 0 are universal constants. The above optimization (Theorem 3.8) shows that Riemannian subgradient descent is able to find the basis vector e n when initialized in the associated region S (n+) 1/(5 log n). We now show that a simple uniformly random initialization on the sphere is guaranteed to be in one of these 2n regions with at least probability 1/2. Lemma 3.9 (Random initialization falls in "good set"). Let q ∼ Uniform(S n−1), then with probability at least 1/2, q belongs to one of the 2n sets S DISPLAYFORM5 As long as the initialization belongs to S DISPLAYFORM0 1/(5 log n), our finding-one-basis in Theorem 3.8 guarantees that Riemannian subgradient descent will converge to e i or −e i respectively. Therefore if we run the algorithm with independent, uniformly random initializations on the sphere multiple times, by a coupon collector's argument, we will recover all the basis vectors. This is formalized in the following theorem. Theorem 3.10 (Recovering the identity dictionary from multiple random initializations). Let m ≥ Cn 4 θ −2 log 4 n and ∈, with probability at least 1 − exp −cmθ 3 n −3 log −3 m the following happens. Suppose we run the Riemannian subgradient descent algorithm independently for R times, each with a uniformly random initialization on S n−1, and choose the step size as DISPLAYFORM1 Then, provided that R ≥ C n log n, all standard basis vectors will be recovered up to accuracy with probability at least 1 − exp (−cR/n) in C Rθ −16/3 −8/3 n 4 log 8/3 n iterations. Here C, C, c ≥ 0 are universal constants. When the dictionary A is not the identity matrix, we can apply the rotation argument sketched in the beginning of this section to get the same , which leads to our main in Theorem 3.1. A key technical challenge is establishing the uniform convergence of subdifferential sets in Proposition 3.5, which we now elaborate. Recall that the population and empirical subdifferentials are DISPLAYFORM0 and we wish to show that the difference between ∂f (q) and E [∂f] (q) is small uniformly over q ∈ Q = S n−1. Two challenges stand out in showing such a uniform convergence:1. The subdifferential is set-valued and random, and it is unclear a-priori how one could formulate and analyze the concentration of random sets.2. The usual covering argument won't work here, as the Lipschitz gradient property does not hold: ∂f (q) and E [∂f] (q) are not Lipschitz in q. Therefore, no matter how fine we cover the sphere in Euclidean distance, points not in this covering can have radically different subdifferential sets. We state and analyze concentration of random sets in the Hausdorff distance (defined in Section 2). We now illustrate how the Hausdorff distance is the "right" distance to consider for concentration of subdifferentials-the reason is that the Hausdorff distance is closely related to the support function of sets, which for any set S ∈ R n is defined as DISPLAYFORM0 For convex compact sets, the sup difference between their support functions is exactly the Hausdorff distance. Lemma 4.1 (Section 1.3.2,). For convex compact sets X, Y ⊂ R n, we have DISPLAYFORM1 Lemma 4.1 is convenient for us in the following sense. Suppose we wish to upper bound the difference of ∂f (q) and E [∂f] (q) along some direction u ∈ S n−1 (as we need in proving the key empirical geometry Theorem 3.6). As both subdifferential sets are convex and compact, by Lemma 4.1 we immediately have DISPLAYFORM2 Therefore, as long as we are able to bound the Hausdorff distance, all directional differences between the subdifferentials are simultaneously bounded, which is exactly what we want to show to carry the benign geometry from the population to the empirical objective. We argue that the absence of gradient Lipschitzness is because the Euclidean distance is not the "right" metric in this problem. Think of the toy example f (x) = |x|, whose subdifferential set ∂f (x) = sign(x) is not Lipschitz across x = 0. However, once we partition R into R >0, R <0 and {0} (i.e. according to the sign pattern), the subdifferential set is Lipschitz on each subset. The situation with the dictionary learning objective is quie similar: we resolve the gradient nonLipschitzness by proposing a stronger metric d E on the sphere which is sign-pattern aware and averages all "subset angles" between two points. Formally, we define d E as DISPLAYFORM0 (the second equality shown in Lemma C.1.) Our plan is to perform the covering argument in d E, which requires showing gradient Lipschitzness in d E and bounding the covering number. (4.7)As long as d E (p, q) ≤, the indicator is non-zero with probability at most, and thus the above expectation should also be small -we bound it by O(log(1/)) in Lemma F.5.To show the same for the empirical subdifferential ∂f, one only needs to bound the observed proportion of sign differences for all p, q such that d E (p, q) ≤, which by a VC dimension argument is uniformly bounded by 2 with high probability (Lemma C.5).Bounding the covering number in d E Our first step is to reduce d E to the maximum length-2 angle (the d 2 metric) over any consistent support pattern. This is achieved through the following vector angle inequality (Lemma C.2): for any p, DISPLAYFORM0 Therefore, as long as sign(p) = sign(q) (coordinate-wise) and DISPLAYFORM1 By Eq. (4.5), the above implies that d E (p, q) ≤ /π, the desired . Hence the task reduces to constructing an η = /n 2 covering in d 2 over any consistent sign pattern. Our second step is a tight bound on this covering number: the η-covering number in d 2 is bounded by exp(Cn log(n/η)) (Lemma C.3). For bounding this, a first thought would be to take the covering in all size-2 angles (there are n 2 of them) and take the common refinement of all their partitions, which gives covering number (C/η) O(n 2) = exp(Cn 2 log(1/η)). We improve upon this strategy by sorting the coordinates in p and restricting attentions in the consecutive size-2 angles after the sorting (there are n − 1 of them). We show that a proper covering in these consecutive size-2 angles by η/n will yield a covering for all size-2 angles by η. The corresponding covering number in this case is thus (Cn/η) O(n) = exp(Cn log(n/η)), which modulo the log n factor is the tightest we can get. Setup We set the true dictionary A to be the identity and random orthogonal matrices, respectively. For each choice, we sweep the combinations of (m, n) with n ∈ {30, 50, 70, 100} and m = 10n {0.5,1,1.5,2,2.5}, and fix the sparsity level at θ = 0.1, 0.3, 0.5, respectively. For each (m, n) pair, we generate 10 problem instances, corresponding to re-sampling the coefficient matrix X for 10 times. Note that our theoretical guarantee applies for m = Ω(n 4), and the sample complexity we experiment with here is lower than what our theory requires. To recover the dictionary, we run the Riemannian subgradient descent algorithm Eq. (3.16) with decaying step size η (k) = 1/ √ k, corresponding to the boundary case α = 1/2 in Theorem 3.8 with a much better base size. Metric As Theorem 3.1 guarantees recovering the entire dictionary with R ≥ Cn log n independent runs, we perform R = round (5n log n) runs on each instance. For each run, a true dictionary element a i is considered to be found if a i − q best ≤ 10 −3. For each instance, we regard it a successful recovery if the R = round (5n log n) runs have found all the dictionary elements, and we report the empirical success rate over the 10 instances. Result From our simulations, Riemannian subgradient descent succeeds in recovering the dictionary as long as m ≥ Cn 2 FIG7, across different sparsity level θ. The dependency on n is consistent with our theory and suggests that the actual sample complexity requirement for guaranteed recovery might be even lower than O(n 4) we established. 3 The O(n 2) rate we observe also matches the based on the SOS method BID6; ). Moreover, the problem seems to become harder when θ grows, evident from the observation that the success transition threshold being pushed to the right. Additional experiments A faster alternative algorithm for large-scale instances is tested in Appendix H. A complementary experiment on real images is included as Appendix I. This paper presents the first theoretical guarantee for orthogonal dictionary learning using subgradient descent on a natural 1 minimization formulation. Along the way, we develop tools for analyzing the optimization landscape of nonconvex nonsmooth functions, which could be of broader interest. For futute work, there is an O(n 2) sample complexity gap between what we established in Theorem 3.1, and what we observed in the simulations alongside previous based on the SOS method BID6; ). As our main geometric Theorem 3.6 already achieved tight bounds on the directional derivatives, further sample complexity improvement could potentially come out of utilizing second-order information such as the strong negative curvature (Lemma B.2), or careful algorithm-dependent analysis. While our applies only to (complete) orthogonal dictionaries, a natural question is whether we can generalize to overcomplete dictionaries. To date the only known provable algorithms for learning overcomplete dictionaries in the linear sparsity regime are based on the SOS method BID6; ). We believe that our nonsmooth analysis has the potential of handling over-complete dictionaries, as for reasonably well-conditioned overcomplete dictionaries A, each a i (columns of A) makes a A approximately 1-sparse and so a i AX gives noisy estimate of a certain row of X. So the same formulation as Eq. (1.1) intuitively still works. We would like to leave that to future work. Nonsmooth phase retrieval and deep networks with ReLU mentioned in Section 1.1 are examples of many nonsmooth, nonconvex problems encountered in practice. Most existing theoretical on these problems tend to be technically vague about handling the nonsmooth points: they either prescribe a rule for choosing a subgradient element, which effectively disconnects theory and practice because numerical testing of nonsmooth points is often not reliable, or ignore the nonsmooth points altogether, assuming that practically numerical methods would never touch these points-this sounds intuitive but no formalism on this appears in the relevant literature yet. Besides our work, (Laurent & von ;) also warns about potential problems of ignoring nonsmooth points when studying optimization of nonsmooth functions in machine learning. We need the Hausdorff metric to measure differences between nonempty sets. For any set X and a point p in R n, the point-to-set distance is defined as DISPLAYFORM0 For any two sets X 1, X 2 ∈ R n, the Hausdorff distance is defined as DISPLAYFORM1 Moreover, for any sets DISPLAYFORM2 (A.4) On the sets of nonempty, compact subsets of R n, the Hausdorff metric is a valid metric; particularly, it obeys the triangular inequality: for nonempty, compact subsets X, Y, Z ⊂ R n, DISPLAYFORM3 (A.5) See, e.g., Sec. 7.1 of for a proof. Lemma A.1 (Restatement of Lemma A.1). For convex compact sets X, Y ⊂ R n, we have DISPLAYFORM4 where h S (u). = sup x∈S x, u is the support function associated with the set S. Proposition A.2 (Talagrand's comparison inequality, Corollary 8.6.3 and Exercise 8.6.5 of Vershynin FORMULA1). Let {X x} x∈T be a zero-mean random process on a subset T ⊂ R n. Assume that for all x, y ∈ T we have DISPLAYFORM0 Then, for any t > 0 sup DISPLAYFORM1 with probability at least 1 − 2 exp −t 2. Here w(T). = E g∼N (0,I) sup x∈T x, g is the Gaussian width of T and rad(T) = sup x∈T x is the radius of T. Proposition A.3 (Deviation inequality for sub-Gaussian matrices, Theorem 9.1.1 and Exercise 9.1.8 of Vershynin FORMULA1). Let A be an n × m matrix whose rows A i's are independent, isotropic, and sub-Gaussian random vectors in R m. Then for any subset T ⊂ R m, we have DISPLAYFORM2 We have DISPLAYFORM3 where the last equality is obtained by conditioning on Ω and the fact that DISPLAYFORM4 The subdifferential expression comes from DISPLAYFORM5 and the fact that We first show that points in the claimed set are indeed stationary points by taking the choice v Ω = 0 in Eq. (3.5), giving the subgradient choice DISPLAYFORM6 DISPLAYFORM7. Let q ∈ S and such that q 0 = k. For all j ∈ supp(q), we have DISPLAYFORM8 On the other hand, for all j / ∈ supp(q), we always have [q Ω] j = 0, so e j E [∂f] (q) = 0. Therefore, we have that E [∂f] (q) = c(θ, k)q, and so DISPLAYFORM9 (B.6) Therefore q ∈ S is stationary. To see that {±e i : i ∈ [n]} are the global minima, note that for all q ∈ S n−1, we have DISPLAYFORM10 Equality holds if and only if q Ω 2 ∈ {0, 1} almost surely, which is only satisfied at q ∈ {±e i : i ∈ [n]}.To see that the other q's are saddles, we only need to show that there exists a tangent direction along which q is local max. Indeed, for any other q, there exists at least two non-zero entries (with equal absolute value): WLOG assume that q 1 = q n > 0. Using the reparametrization in Appendix B.3 and applying Lemma B.2, we get that E [f] (q) is directionally differentiable along [−q −n ; 1−q 2 n qn], with derivative zero (necessarily, because 0 ∈ E [∂ R f] (q)) and strictly negative second derivative. Therefore E [f] (q) is locally maximized at q along this tangent direction, which shows that q is a saddle point. The other direction (all other points are not stationary) is implied by Theorem 3.4, which guarantees that 0 / ∈ E [∂ R f] (q) whenever q / ∈ S. Indeed, as long as q / ∈ S, q has a max absolute value coordinate (say n) and another non-zero coordinate with strictly smaller absolute value (say j). For this pair of indices, the proof of Theorem 3.4(a) goes through for index j (even if q ∈ S (n+) ζ0 does not necessarily hold because the max index might not be unique), which implies that 0 / ∈ E [∂ R f] (q). For analysis purposes, we introduce the reparametrization w = q 1:(n−1) in the region S (n+) 0, following . With this reparametrization, the problem becomes DISPLAYFORM0 The constraint comes from the fact that q n ≥ 1/ √ n and thus w ≤ (n − 1)/n. Lemma B.1. We have DISPLAYFORM1 Proof. Direct calculation gives DISPLAYFORM2 as claimed. Lemma B.2 (Negative-curvature region). For all unit vector v ∈ S n−1 and all s ∈, let DISPLAYFORM3 it holds that DISPLAYFORM4 In other words, for all w = 0, ±w/ w is a direction of negative curvature. Proof. By Lemma B.1, DISPLAYFORM5 For s ∈, h v (s) is twice differentiable, and we have DISPLAYFORM6 completing the proof. Proof. For any unit vector v ∈ R n−1, define h v (t). = E [g] (tv) for t ∈. We have from Lemma B.1 DISPLAYFORM0 (B.20)Moreover, DISPLAYFORM1 We are interested in the regime of t so that DISPLAYFORM2 (B.28) DISPLAYFORM3 For any w, applying the above to the unit vector w/ w and recognizing that ∇ t h w/ w (t) = D w/ w g (w) = D c w/ w g (w), we complete the proof. We first show Eq. (3.9) using the reparametrization in Appendix B.3. We have DISPLAYFORM0 where the second equality follows by differentiating g via the chain rule. Now, by Lemma B.3, DISPLAYFORM1 For each radial direction v. = w/ w, consider points of the form tv with t ≤ 1/ 1 + v 2 ∞. Obviously, the function DISPLAYFORM2 is monotonically decreasing wrt t. Thus, to derive a lower bound, it is enough to consider the largest t allowed. In S (n+) ζ0, the limit amounts to requiring q 2 n / w DISPLAYFORM3 So for any fixed v and all allowed t for points in S DISPLAYFORM4, a uniform lower bound is DISPLAYFORM5 So we conclude that for all q ∈ S DISPLAYFORM6 We now turn to showing Eq. (3.8). For e j with q j = 0, DISPLAYFORM7 (B.37)So for all j with q j = 0, we have DISPLAYFORM8 completing the proof. C PROOFS FOR SECTION 3.2 DISPLAYFORM9 We stress that this notion always depend on θ, and we will omit the subscript θ when no confusion is expected. This indeed defines a metric on subsets of S n−1.Lemma C.1. Over any subset of S n−1 with a consistent support pattern, d E is a valid metric. Proof. Recall that (x, y). = arccos x, y defines a valid metric on S n−1. 4 In particular, the triangular inequality holds. For d E and p, q ∈ S n−1 with the same support pattern, we have DISPLAYFORM10 where we have adopted the convention that (0, v). = 0 for any v. It is easy to verify that d E (p, q) = 0 ⇐⇒ p = q, and d E (p, q) = d E (q, p). To show the triangular inequality, note that for any p, q and r with the same support pattern, p Ω, q Ω, and r Ω are either identically zero, or all nonzero. For the former case, DISPLAYFORM11 holds trivially. For the latter, since (·, ·) obeys the triangular inequality uniformly over the sphere, DISPLAYFORM12 Proof. The inequality holds trivially when either of u, v is zero. Suppose they are both nonzero and wlog assume both are normalized, i.e., u = v = 1. Then, DISPLAYFORM13 2 ) (u Ω, v Ω) > π/2, the claimed inequality holds trivially, as (u, v) ≤ π/2 by our assumption. Suppose Ω∈(DISPLAYFORM14 by recursive application of the following inequality: DISPLAYFORM15 So we have that when Ω∈( DISPLAYFORM16 as claimed. Lemma C.3 (Covering in maximum length-2 angles). For any η ∈ (0, 1/3), there exists a subset Q ⊂ S n−1 of size at most (5n log(1/η)/η) 2n−1 satisfying the following: for any p ∈ S n−1, there exists some q ∈ Q such that (p Ω, q Ω) ≤ η for all Ω ⊂ [n] with |Ω|≤ 2. DISPLAYFORM17 our goal is to give an η-covering of S n−1 in the d 2 metric. Step 1 We partition S n−1 according to the support, the sign pattern, and the ordering of the non-zero elements. For each configuration, we are going to construct a covering with the same configuration of support, sign pattern, and ordering. There are no more than 3 n · n! such configurations. Note that we only need to construct one such covering for each support size, and for each support size we can ignore the zero entries -the angle (p Ω, q Ω) is always zero when p, q have matching support and Ω contains at least one zero index. Therefore, the task reduces to bounding the covering number of DISPLAYFORM18 Step 2 We bound the covering number of A n by induction. Suppose that DISPLAYFORM19 holds for all n ≤ n − 1. (The base case m = 2 clearly holds.) Let C n ⊂ S n −1 be the correpsonding covering sets. We now construct a covering for A n. Let R. = 1/η = r k for some r ≥ 1 and k to be determined. Consider the set DISPLAYFORM20 We claim that Q r,k with properly chosen (r, k) gives a covering of DISPLAYFORM21. Each consecutive ratio p i+1 /p i falls in one of these intervals, and we choose q so that q i+1 /q i is the left endpoint of this interval. Such a q satisfies q ∈ Q r,k and DISPLAYFORM22 By multiplying these bounds, we obtain that for all 1 ≤ i < j ≤ n, DISPLAYFORM23 Take r = 1 + η/2n, we have r n−1 = (1 + η/2n) n−1 ≤ exp(η/2) ≤ 1 + η. Therefore, for all i, j,we have pj /pi qj /qi ∈ [1, 1 + η), which further implies that ((p i, p j), (q i, q j)) ≤ η by Lemma F.4. Thus we have for all |Ω|≤ 2 that (p Ω, q Ω) ≤ η. (The size-1 angles are all zero as we have sign match.)For this choice of r, we have k = log R/log r and thus DISPLAYFORM24 (C.28) and we have N (A n,R) ≤ N n.Step 3 We now construct the covering of A n \ A n,R. For any p ∈ A n \ A n,R, there exists some i such that p i+1 /p i ∈ [R, ∞), which means that the angle of the ray (p i, p i+1) is in between [arctan(R), π/2) = [π/2 − η, π/2). As p is sorted, we have that DISPLAYFORM25 So if we take q such that q i+1 /q i ∈ [R, ∞), q also has the above property, which gives that DISPLAYFORM26 Therefore to obtain the cover in d 2, we only need to consider the angles for Ω ⊂ {1, . . ., i} and Ω ⊂ {i + 1, . . ., n}, which can be done by taking the product of the covers in A i and A n−i.By considering all i ∈ {1, . . ., n − 1}, we obtain the bound DISPLAYFORM27 Step 4 Putting together Step 2 and Step 3 and using the inductive assumption, we get that DISPLAYFORM28 This shows the case for m = n and completes the induction. Step 5 Considering all configurations of {support, sign pattern, ordering}, we have DISPLAYFORM29 Lemma C.4 (Covering number in the d E metric). Assume n ≥ 3. There exists a numerical constant C > 0 such that for any ∈, S n−1 admits an -net of size exp(Cn log n) w.r.t. d E defined in Eq. (C.1): for any p ∈ S n−1, there exists a q in the net with supp (q) = supp (p) and d E (p, q) ≤. We say such nets are admissible for S n−1 wrt d E.Proof. Let η = /n 2. By Lemma C.3, there exists a subset Q ⊂ S n−1 of size at most 5n log(1/η) η DISPLAYFORM30 such that for any p ∈ S n−1, there exists q ∈ S n−1 such that supp(p) = supp(q) and (p Ω, q Ω) ≤ η for all |Ω|≤ 2. In particular, the |Ω|= 1 case says that sign(p) = sign(q), which implies that DISPLAYFORM31 Thus, applying the vector angle inequality (Lemma C.2), for any p ∈ S n−1 and the corresponding q ∈ Q, we have DISPLAYFORM32 Summing up, we get DISPLAYFORM33 Below we establish the "Lipschitz" property in terms of d E distance. Lemma C.5. Fix θ ∈. For any ∈, let N be an admissible -net for S n−1 wrt d E. Let x 1,..., x m be iid copies of x ∼ iid BG(θ) in R n. When m ≥ C −2 n, the inequality DISPLAYFORM34 holds with probability at least 1 − exp −c 2 m. Here C and c are universal constants independent of. Proof. We call any pair of p, q ∈ S n−1 with q ∈ N, supp (p) = supp (q), and d E (p, q) ≤ an admissible pair. Over any admissible pair (p, q), E [R] = d E (p, q). We next bound the deviation R − E [R] uniformly over all admissible (p, q) pairs. Observe that the process R is the sample average of m indicator functions. Define the hypothesis class H = x → 1 sign p x = sign q x: (p, q) is an admissible pair.(C.42) and let d vc (H) be the VC-dimension of H. From concentration for VC-classes (see, e.g., Eq and Theorem 3.4 of BID7 ), we have DISPLAYFORM35 for any t > 0. It remains to bound the VC-dimension d vc (H). First, we have DISPLAYFORM36 Observe that each set in the latter hypothesis class can be written as DISPLAYFORM37 the union of intersections of two halfspaces. Thus, letting DISPLAYFORM38 be the class of halfspaces, we have DISPLAYFORM39 Note that H 0 has VC-dimension n + 1. Applying bounds on the VC-dimension of unions and intersections (Theorem 1.1,), we get that DISPLAYFORM40 Plugging this bound into Eq. (C.43), we can set t = /2 and make m large enough so that C 0 √ C 3 n/m ≤ /2, completing the proof. Proposition C.6 (Pointwise convergence). For any fixed q ∈ S n−1, DISPLAYFORM0 Here C a, C b ≥ 0 are universal constants. DISPLAYFORM1 and consider the zero-mean random process {X u} defined on S n−1. For any u, v ∈ S n−1, we have DISPLAYFORM2 where we write DISPLAYFORM3 where we have used Lemma F.1 to obtain the last upper bound. If DISPLAYFORM4 and we can use similar argument to conclude that DISPLAYFORM5 Thus, {X u} is a centered random process with sub-Gaussian increments with a parameter C 4 / √ m. We can apply Proposition A.2 to conclude that DISPLAYFORM6 which implies the claimed . Throughout the proof, we let c, C denote universal constants that could change from step to step. Fix an ∈ (0, 1/2) to be decided later. Let N be an admissible net for S n−1 wrt d E, with |N | ≤ exp(Cn log(n/)) (Lemma C.4). By Proposition C.6 and the union bound, DISPLAYFORM0 For any p ∈ S n−1, let q ∈ N satisfy supp (q) = supp (p) and d E (p, q) ≤. Then we have DISPLAYFORM1 by the triangular inequality for the Hausdorff metric. By the preceding union bound, term I is bounded by t/3 as long as the bad event does not happen. For term II, we have DISPLAYFORM2 where the last line follows from Lemma F.5. As long as ≤ ct/ log(1/t), the above term is upper bounded by t/3. For term III, we have DISPLAYFORM3 By Lemma C.5, with probability at least 1 − exp(−c 2 m), the number of different signs is upper bounded by 2m for all p, q such that d E (p, q) ≤. On this good event, the above quantity can be upper bounded as follows. Define a set T. = {s ∈ R m : s i ∈ {+1, −1, 0}, s 0 ≤ 2m } and consider the quantity sup s∈T Xs, where DISPLAYFORM4 uniformly (i.e., indepdent of p, q and u). We have DISPLAYFORM5 Noting that 1/ √ θ · X has independent, isotropic, and sub-Gaussian rows with a parameter C/ √ θ, we apply Proposition A.3 and obtain that DISPLAYFORM6 with probability at least 1 − 2 exp −t 2 0. So we have over all admissible (p, q) pairs, DISPLAYFORM7 Setting t 0 = ct √ m and = ct θ/log m, we have that DISPLAYFORM8 provided that m ≥ C t −2 n = Ct −1 n θ/log m, which is subsumed by the earlier requirement m ≥ Ct −2 n. Putting together the three bounds Eq. (C.62), Eq. (C.67), Eq. (C.80), we can choose DISPLAYFORM9. A sufficient condition is that m ≥ Cnt −2 log 2 (n/t) for sufficiently large C. When this is satisfied, the probability is further lower bounded by 1 − exp(−cmθt 2 /log m). Define DISPLAYFORM0 (C.85) By Proposition 3.5, with probability at least 1 − exp −cmθ 3 ζ 2 0 n −3 log −1 m we have DISPLAYFORM1 0 log (n/ζ 0). We now show the properties Eq. (3.12) and Eq. (3.13) on this good event, focusing on S (n+) ζ0 but obtaining the same for all other 2n − 1 subsets by the same arguments. For Eq. (3.12), we have ∂ R f (q), e j /q j − e n /q n = ∂f (q), I − qq (e j /q j − e n /q n) = ∂f (q), e j /q j − e n /q n.(C.87) Now sup ∂f (q), e n /q n − e j /q j = h ∂f (q) (e n /q n − e j /q j) (C.88) = Eh ∂f (q) (e n /q n − e j /q j) − Eh ∂f (q) (e n /q n − e j /q j) + h ∂f (q) (e n /q n − e j /q j) (C.89)≤ Eh ∂f (q) (e n /q n − e j /q j) + e n /q n − e j /q j sup DISPLAYFORM2 By Theorem 3.4(a), DISPLAYFORM3 Moreover, e n /q n − e j /q j = 1/q 2 n + 1/q 2 j ≤ 1/q 2 n + 3/q 2 n ≤ 2 √ n. Meanwhile, we have DISPLAYFORM4 We conclude that inf ∂f (q), e j /q j − e n /q n = − sup ∂f (q), e n /q n − e j /q j (C.94) DISPLAYFORM5 as claimed. For Eq. (3.13), we have by Theorem 3.4(b) that sup ∂f (q), e n − q n q = h ∂f (q) (e n − q n q) (C.97) = Eh ∂f (q) (e n − q n q) − Eh ∂f (q) (e n − q n q) + h ∂f (q) (e n − q n q) (C.98)≤ Eh ∂f (q) (e n − q n q) + e n − q n q sup DISPLAYFORM6 (C.100)As we are on the good event DISPLAYFORM7 we have inf ∂f (q), q n q − e n = − sup ∂f (q), e n − q n q (C.102) DISPLAYFORM8 q − e n for all q with q n ≥ 0 completes the proof. C.5 PROOF OF PROPOSITION 3.7For any q ∈ S n−1, DISPLAYFORM9 by the metric property of the Hausdorff metric. On one hand, we have DISPLAYFORM10 On the other hand, by Proposition 3.5, DISPLAYFORM11 (C.107) with probability at least 1 − exp −c 1 mθ log −1 m, provided that m ≥ C 2 n 2 log n (simplified using θ ≥ 1/n). Combining the two complete the proof. For w with w = 1/2, DISPLAYFORM12 So, back to the q space, DISPLAYFORM13 Combining the in Eq. (C.115) and Eq. (C.121), we conclude that with high probability DISPLAYFORM14, which is equivalent to w ≤ 1/2 in the w space. Under this constraint, by Lemma B.3, DISPLAYFORM15 So, emulating the proof of Eq. (3.9) in Theorem 3.4, we have that for q ∈ S (n+) ζ0with q −n ≤ 1/2, (C.125) where at the last inequality we use q n = 1 − w 2 ≥ √ 3/2 when w ≤ 1/2. Moreover, we emulate the proof of Eq. (3.13) in Theorem 3.6 to obtain that C.126) with probability at least 1 − exp −cmθ 3 log −1 m, provided that m ≥ Cθ −2 n log n. DISPLAYFORM16 DISPLAYFORM17 The last step of our proof is invoking the mean value theorem, similar to the proof of Proposition C.7. For any q, we have DISPLAYFORM18 for a certain t ∈ and a certain v ∈ ∂g (tw). We have ). Set η = t 0 /(100 √ n) for t 0 ∈. For any ζ 0 ∈, on the good events stated in Proposition 3.7 and Theorem 3.6, we have for all q ∈ S (n+) ζ0 \ S (n+) 1 DISPLAYFORM19 and q + being the next step of Riemannian subgradient descent that DISPLAYFORM20 In particular, we have q + ∈ S (n+) ζ0.Proof. We divide the index set [n − 1] into three sets DISPLAYFORM21 We perform different arguments on different sets. We let g (q) ∈ ∂ R f (q) be the subgradient taken at q and note by Proposition 3.7 that g ≤ 2, and so |g i | ≤ 2 for all i ∈ [n]. We have DISPLAYFORM22 Provided that η ≤ 1/(4 √ n), 1 − 2η √ n ≥ 1/2, and so DISPLAYFORM23 where the last inequality holds when η ≤ 1/ √ 40n. For any j ∈ I 1, DISPLAYFORM24 where the very last inequality holds when η ≤ 1/(26 √ n). DISPLAYFORM25, I 2 is nonempty. For any j ∈ I 2, q 2 +,n q 2 DISPLAYFORM26 Since g j /q j ≤ 2 √ 3n, 1 − ηg j /q j ≥ 1/2 when η ≤ 1/ 4 √ 3n. Conditioned on this and due to that g j /q j − g n /q n ≥ 0, it follows DISPLAYFORM27, we have q 2 n / q −n 2 ∞ ≤ 2, so there must be a certain j ∈ I 2 satisfying q 2 n /q 2 j ≤ 2. We conclude that when DISPLAYFORM28 the index of largest entries of q +,−n remains in I 2.On the other hand, when η ≤ 1/(100 √ n), for all j ∈ I 2, DISPLAYFORM29 So when η = t/(100 √ n) for any t ∈, DISPLAYFORM30 completing the proof. Proposition D.2. For any ζ 0 ∈, on the good events stated in Proposition 3.7 and Theorem 3.6, if the step sizes satisfy DISPLAYFORM31 the iteration sequence will stay in S DISPLAYFORM32 where the last inequality holds provided that η ≤ (1 − ζ 0) /(9 √ n). Combining the two cases finishes the proof. D.2 PROOF OF THEOREM 3.8As we have DISPLAYFORM33, the entire sequence q DISPLAYFORM34 will stay in S For any q and any v ∈ ∂ R f (q), we have v, q = 0 and therefore DISPLAYFORM35 So q − ηv is not inside B n. Since projection onto B n is a contraction, we have DISPLAYFORM36 where we have used the bounds in Proposition 3.7 and Theorem 3.6 to obtain the last inequality. Further applying Proposition C.7, we have DISPLAYFORM37 Summing up the inequalities until step K (assumed ≥ 5), we have DISPLAYFORM38 Substituting the following estimates DISPLAYFORM39 and noting 16 q (K) − e n 2 ≤ 32, we have DISPLAYFORM40 and when K ≥ 1, DISPLAYFORM41 (D.27) So we conclude that when DISPLAYFORM42 When this happens, by Proposition C.8, DISPLAYFORM43 Plugging in the choice ζ 0 = 1/(5 log n) in Eq. (D.28) gives the desired bound on the number of iterations. E PROOFS FOR SECTION 3.4E.1 PROOF OF LEMMA 3.9Lemma E.1. For all n ≥ 3 and ζ ≥ 0, it holds that DISPLAYFORM44 We note that a similar appears in but our definitions of the region S ζ are slightly different. For completeness we provide a proof in Lemma F.3.We now prove Lemma 3.9. Taking ζ = 1/(5 log n) in Lemma E.1, we obtain DISPLAYFORM45 By symmetry, all the 2n sets S1/(5 log n), S DISPLAYFORM46 have the same volume which is at least 1/(4n). As q ∼ Uniform(S n−1), it falls into their union with probability at least 2n · 1/(4n) = 1/2, on which it belongs to a uniformly random one of these 2n sets. Assume that the good event in Proposition 3.7 happens and that in Theorem 3.6 happens to all the 2n sets S (i+) 1/(5 log n), S DISPLAYFORM0, which by setting ζ 0 = 1/(5 log n) has probability at least DISPLAYFORM1 By Lemma 3.9, random initialization will fall these 2n sets with probability at least 1/2. When it falls in one of these 2n sets, by Theorem 3.8, one run of the algorithm will find a signed standard basis vector up to accuracy. With R independent runs, at least S. = 1 4 R of them are effective with probability at least 1 − exp −(R/4) 2 /(R/4 · 2) = 1 − exp (−R/8), due to Bernstein's inequality. After these effective runs, the probability any standard basis vector is missed (up to sign) is bounded by DISPLAYFORM2 where the second inequality holds whenever S ≥ 2n log n. Lemma F.1. For x ∼ BG(θ), x ψ2 ≤ C a. For any vector u ∈ R n and x ∼ iid BG(θ), DISPLAYFORM0 Proof. For any λ ∈ R, DISPLAYFORM1 So x ψ2 is bounded by a universal constant. Moreover, DISPLAYFORM2 for any t ≥ 0. Here C a, C b ≥ 0 are universal constants. Consider the zero-centered random process defined on S n−1: DISPLAYFORM0 where we use the estimate in Lemma F.1 to obtain the last inequality. Note that X q is a mean-zero random process, and we can invoke Proposition A.2 with w(S n−1) = C 4 √ n and rad S n−1 = 2 to get the claimed . Lemma F.3. For all n ≥ 3 and ζ ≥ 0, it holds that DISPLAYFORM1 Proof. We have DISPLAYFORM2 where we write ψ(t). = 1 √ 2π t −t exp −s 2 /2 ds. Now we derive a lower bound of the volume ratio by considering a first-order Taylor expansion of the last equation around ζ = 0 (as we are mostly interested in small ζ). By symmetry,h = 1/(2n). Moreover, we have DISPLAYFORM3 2 /2 x 2 ψ n−1 (x) dx (F.15) and combining the above integral , we conclude thath (ζ) ≥ 0 and complete the proof. Lemma F.4. Let (x 1, y 1), (x 2, y 2) ∈ R 2 >0 be two points in the first quadrant satisfying y 1 ≥ x 1 and y 2 ≥ x 2, and y2/x2 y1/x1 ∈ [1, 1 + η] for some η ≤ 1, then we have ((x 1, y 1), (x 2, y 2)) ≤ η. DISPLAYFORM4 Proof. For i = 1, 2, let θ i be the angle between the ray (x i, y i) and the x-axis. Our assumption implies that θ i ∈ [π/4, π/2) and θ 2 ≥ θ 1, thus ((x 1, y 1), (x 2, y 2)) = θ 2 − θ 1, so we have (F.30) tan ((x 1, y 1), (x 2, y 2)) = tan θ 2 − tan θ 1 1 + tan θ 2 tan θ 1 = y 2 /x 2 − y 1 /x 1 1 + y 2 y 1 /(x 2 x 1) = y2/x2 y1/x1 − 1 y 2 /x 2 + x 1 /y 1 ≤ y 2 /x 2 y 1 /x 1 − 1 ≤ η. Therefore ((x 1, y 1), (x 2, y 2)) ≤ arctan(η) ≤ η. where we have used θ ≤ 1/2 and ≤ 1/2. The Riemannian subgradient descent is cheap per iteration but slow in overall convergence, similar to many other first-order methods. We also test a faster quasi-Newton type method, GRANSO, 6 that employs BFGS for solving constrained nonsmooth problems based on sequential quadratic optimization . For a large dictionary of dimension n = 400 and sample complexity m = 10n 2 (i.e., 1.6 × 10 6), GRANSO successfully identifies a basis after 1500 iterations with CPU time 4 hours on a two-socket Intel Xeon E5-2640v4 processor (10-core Broadwell, 2.40 GHz)-this is approximately 10× faster than the Riemannian subgradient descent method, showing the potential of quasi-Newton type methods for solving large-scale problems. To experiment with images, we follow a typical setup for dictionary learning as used in image processing . We focus on testing if complete (i.e., square and invertible) dictionaries are reasonable sparsification bases for real images, instead on any particular image processing or vision tasks. so that nonvanishing singular values of Y are identically one. We then solve formulation (1.1) round (5n log n) times with n = 64 using the BFGS solver based on GRANSO, obtaining round (5n log n) vectors. Negative equivalent copies are pruned and vectors with large correlations with other remaining vectors are sequentially removed until only 64 vectors are left. This forms the final complete dictionary. Results The learned complete dictionaries for the two test images are displayed in the second row of FIG1. Visually, the dictionaries seem reasonably adaptive to the image contents: for the left image with prevalent sharp edges, the learned dictionary consists of almost exclusively oriented sharp corners and edges, while for the right image with blurred textures and occasional sharp features, the learned dictionary does seem to be composed of the two kinds of elements. Let the learned dictionary be A. We estimate the representation coefficients as A −1 Y. The third row of FIG1 contains the histograms of the coefficients. For both images, the coefficients are sharply concentrated around zero (see also the fourth row for zoomed versions of the portions around zero), and the distribution resembles a typical zero-centered Laplace distribution-which is a good indication of sparsity. Quantitatively, we calculate the mean sparsity level of the coefficient vectors (i.e., columns of A −1 Y) by the metric · 1 / · 2: for a vector v ∈ R n, v 1 / v 2 ranges from 1 (when v is one-sparse) to √ n (when v is fully dense with elements of equal magnitudes), which serves as a good measure of sparsity level for v. For our two images, the sparsity levels by the norm-ratio metric are 5.9135 and 6.4339, respectively, while the fully dense extreme would have a value √ 64 = 8, suggesting the complete dictionaries we learned are reasonable sparsification bases for the two natural images, respectively.

Efficient dictionary learning by L1 minimization via a novel analysis of the non-convex non-smooth geometry.

Represent the short summary for finding the corresponding scientific article (long):

General analysis of sign-based methods (e.g. signSGD) for non-convex optimization, built on intuitive bounds on success probabilities.