# Memory-Efficient Adaptive Optimization

Source: https://arxiv.org/abs/1901.11150

Adaptive gradient-based optimizers such as [AdaGrad](https://paperswithcode.com/method/adagrad) and [Adam](https://paperswithcode.com/method/adam) are among the
defacto methods of choice in modern machine learning.These methods tune the learning rate for each parameter during the optimization process using cumulative second-order statistics. These methods provide superior convergence properties and are very attractive in large scale applications due to their moderate time and space requirements which are linear in the number of parameters.


However, the recent advances in natural language processing such as [BERT](https://paperswithcode.com/method/bert) and GPT2 show that models with 108 to 1010 parameters, trained with adaptive optimization methods, achieve state-of-the-art results. In such cases, the memory overhead of the optimizer can restrict the size of the model that can be used as well as the batch size, both of which can have a dramatic effect on the quality of the final model.


Here we construct a new adaptive optimization method that retains most of the benefits of standard per-parameter adaptivity while significantly reducing memory overhead.


We observe that in standard neural networks that certain entries of the stochastic gradients have (on average) similar values, and exhibit what we refer to as an activation pattern. For example, in gradients of embedding layers of deep networks, an entire row (or column) is either zero or non-zero. Similarly, in intermediate layers we often observe that gradients associated with the same unit are of similar order of magnitude. In these cases, a similar phenomenon is observed in the second-order statistics maintained by adaptive methods. With this key observation, to reduce the memory overhead of the optimizer our method takes in a cover set of the parameters. Cover sets are typically selected in practice such that parameters in each of the sets have second order statistics of similar magnitude. Our method is general enough that it can easily be extended to arbitrary cover sets. For parameters of deep networks that are organized as a collection of tensors, we form a cover consisting of slices of codimension one for each tensor. Thus, for an m x n parameter matrix, the cover consists of rows and columns of the matrix. The memory requirements therefore drop from mxn to merely m+n. For a parameter tensor of rank p, with dimensions n1 ... np, the reduction in memory consumption is even more pronounced, dropping from product of all the dimensions to the sum of all dimensions. This virtually eliminates the memory overhead associated with maintaining the adaptive learning rates!

Another practical aspect worthy of note is that our method does not require an external hand engineered learning rate decay schedule but instead relies on the per parameter adaptivity that is natural to its update rule which makes it easier to tune. We provide details in the supplementary section of the paper.

## Advice on using SM3 on your model

### Learning rate warm-up:

```python
learning_rate = lr_constant * tf.minimum(1.0, (warm_up_step / global_step) ** p)
```

* p = 1, linear ramp up of learning rate.
* p = 2, quadratic ramp up of learning rate [preferred].

We typically set `warm_up_step` as 5% of overall steps. Initially, the norm of the preconditioned gradient is much larger than norm of the weights. Learning rate warmup allows us to heuristically fix this scale mismatch.

### Learning rate decay:

We make use accumulated gradient squares for the decay. This means that each coordinate gets its own natural decay based on the scales of the gradients over time. Hence, users need not put in an external learning rate decay schedule. Moreover, we found in our experiments with translation and language models that this approach is superior to a hand-tuned learning rate decay schedules which is typically combined with exponential moving averages of the gradient squares.

Having said that if users want to add exponential moving averages instead of the standard accumulated gradient squares - It's easy to modify the optimizer implementation to switch to exponential moving averages.

For rank > 1:

| from | to |
|-------------------------------------|-------------------------------------|
| current_accumulator += grad * grad | current_accumulator = beta * current_accumulator + (1-beta) * grad * grad |


For rank <= 1:


| from | to |
|-------------------------------------|-------------------------------------|
| current_accumulator = tf.assign_add(accumulator, grad * grad) | current_accumulator = tf.assign(accumulator, beta * accumulator + (1-beta) * (grad * grad)) |


### [Polyak averaging](https://paperswithcode.com/method/polyak-averaging) of parameters: 
It's useful to run [polyak averaging](https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage) of the parameters. These parameters are then used in inference / serving. Using the averaged parameters instead of the last iterate typically improves the overall performance of the model.

An **alternative** to polyak averaging which does not make use of extra memory is to decay the learning rate from the constant to zero for the last 10% of the steps of your entire training run, we term the phase a **cool-down** phase of the model. As training makes smaller and smaller steps the final iterate can be thought of as an average iterate.

Memory Efficient Adaptive Optimization

A **Dueling Network** is a type of Q-Network that has two streams to separately estimate (scalar) state-value and the advantages for each action. Both streams share a common convolutional feature learning module. The two streams are combined via a special aggregating layer to produce an
estimate of the state-action value function Q as shown in the figure to the right.

The last module uses the following mapping:

$$ Q\left(s, a, \theta, \alpha, \beta\right) =V\left(s, \theta, \beta\right) + \left(A\left(s, a, \theta, \alpha\right) - \frac{1}{|\mathcal{A}|}\sum\_{a'}A\left(s, a'; \theta, \alpha\right)\right) $$

This formulation is chosen for identifiability so that the advantage function has zero advantage for the chosen action, but instead of a maximum we use an average operator to increase the stability of the optimization.

Dueling Network

Dueling Network Architectures for Deep Reinforcement Learning

Chinchilla is a 70B parameters model trained as a compute-optimal model with 1.4 trillion tokens. Findings suggest that these types of models are trained optimally by equally scaling both model size and training tokens. It uses the same compute budget as Gopher but with 4x more training data. Chinchilla and Gopher are trained for the same number of FLOPs. It is trained using [MassiveText](/dataset/massivetext) using a slightly modified SentencePiece tokenizer. More architectural details in the paper.

Chinchilla

Training Compute-Optimal Large Language Models

An attention mechanism for content-based filtering of multi-level features. For example, recurrent features obtained by forward and backward passes of a bidirectional RNN block can be combined using attention feature filters, with unprocessed input features/embeddings as queries and recurrent features as keys/values.

Attention Feature Filters

NeuriCam: Key-Frame Video Super-Resolution and Colorization for IoT Cameras

**Ternary Weight Splitting** is a ternarization approach used in [BinaryBERT](https://www.paperswithcode.com/method/binarybert) that exploits the flatness of ternary loss landscape as the optimization proxy of the binary model. We first train the half-sized ternary BERT to convergence, and then split both the latent full-precision weight $\mathbf{w}^{t}$ and quantized $\hat{\mathbf{w}}^{t}$ to their binary counterparts $\mathbf{w}\_{1}^{b}, \mathbf{w}\_{2}^{b}$ and $\hat{\mathbf{w}}\_{1}^{b}, \hat{\mathbf{w}}\_{2}^{b}$ via the TWS operator. To inherit the performance of the ternary model after splitting, the TWS operator requires the splitting equivalency (i.e., the same output given the same input):

$$
\mathbf{w}^{t}=\mathbf{w}\_{1}^{b}+\mathbf{w}\_{2}^{b}, \quad \hat{\mathbf{w}}^{t}=\hat{\mathbf{w}}\_{1}^{b}+\hat{\mathbf{w}}\_{2}^{b}
$$

While solution to the above equation is not unique, we constrain the latent full-precision weights after splitting $\mathbf{w}\_{1}^{b}, \mathbf{w}\_{2}^{b}$ to satisfy $\mathbf{w}^{t}=\mathbf{w}\_{1}^{b}+\mathbf{w}\_{2}^{b}$. See the paper for more details.

Ternary Weight Splitting

BinaryBERT: Pushing the Limit of BERT Quantization

**FCOS** is an anchor-box free, proposal free, single-stage object detection model. By eliminating the predefined set of anchor boxes, FCOS avoids computation related to anchor boxes such as calculating overlapping during training. It also avoids all hyper-parameters related to anchor boxes, which are often very sensitive to the final detection performance.

FCOS

FCOS: Fully Convolutional One-Stage Object Detection

Rana Mostafa, Hoda Baraka and AbdelMoniem Bayoumi

**LMOT**, i.e., Light-weight Multi-Object Tracker,  performs joint pedestrian detection and tracking. LMOT introduces a simplified DLA-34 encoder network to extract detection features for the current image that are computationally efficient. Furthermore, we generate efficient tracking features using a linear transformer for the prior image frame and its corresponding detection heatmap. After that, LMOT fuses both detection and tracking feature maps in a multi-layer scheme and performs a two-stage online data association relying on the Kalman filter to generate tracklets. We evaluated our model on the challenging real-world MOT16/17/20 datasets, showing LMOT significantly outperforms the state-of-the-art trackers concerning runtime while maintaining high robustness. LMOT is approximately ten times faster than state-of-the-art trackers while being only 3.8% behind in performance accuracy on average leading to a much computationally lighter model.

Code: https://github.com/RanaMostafaAbdElMohsen/LMOT
Paper: https://doi.org/10.1109/ACCESS.2022.3197157

LMOT

**Siren**, or **Sinusoidal Representation Network**, is a periodic activation function for implicit neural representations. Specifically it uses the sine as a periodic activation function:

$$ \Phi\left(x\right) = \textbf{W}\_{n}\left(\phi\_{n-1} \circ \phi\_{n-2} \circ \dots \circ \phi\_{0} \right) $$

Siren

Implicit Neural Representations with Periodic Activation Functions

**SimpleNet** is a convolutional neural network with 13 layers. The network employs a homogeneous design utilizing 3 × 3 kernels for convolutional layer and 2 × 2 kernels for pooling operations. The only layers which do not use 3 × 3 kernels are 11th and 12th layers, these layers, utilize 1 × 1 convolutional kernels. Feature-map down-sampling is carried out using nonoverlaping 2 × 2 max-pooling. In order to cope with the problem of vanishing gradient and also over-fitting, SimpleNet also uses batch-normalization with moving average fraction of 0.95 before any [ReLU](https://paperswithcode.com/method/relu) non-linearity.

SimpleNet

Lets keep it simple, Using simple architectures to outperform deeper and more complex architectures

The **MelGAN Residual Block** is a convolutional [residual block](https://paperswithcode.com/method/residual-block) used in the [MelGAN](https://paperswithcode.com/method/melgan) generative audio architecture. It employs residual connections with dilated convolutions. Dilations are used so that temporally far output activations of each subsequent layer has significant overlapping inputs. Receptive field of a stack of [dilated convolution](https://paperswithcode.com/method/dilated-convolution) layers increases exponentially with the number of layers. Incorporating these into the MelGAN generator allows us to efficiently increase the induced receptive fields of each output time-step. This effectively implies larger overlap in the induced receptive field of far apart time-steps, leading to better long range correlation.

MelGAN Residual Block

MelGAN: Generative Adversarial Networks for Conditional Waveform Synthesis

NeRF represents a scene with learned, continuous volumetric radiance field $F_\theta$ defined over a bounded 3D volume. In a NeRF, $F_\theta$ is a multilayer perceptron (MLP) that takes as input a 3D position $x = (x, y, z)$ and unit-norm viewing direction $d = (dx, dy, dz)$, and produces as output a density $\sigma$ and color $c = (r, g, b)$. The weights of the multilayer perceptron that parameterize $F_\theta$ are optimized so as to encode the radiance field of the scene. Volume rendering is used to compute the color of a single pixel.

NeRF

NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis

**Hierarchical Feature Fusion (HFF)** is a feature fusion method employed in [ESP](https://paperswithcode.com/method/esp) and [EESP](https://paperswithcode.com/method/eesp) image model blocks for degridding. In the ESP module, concatenating the outputs of dilated convolutions gives the ESP module a large effective receptive field, but it introduces unwanted checkerboard or gridding artifacts. To address the gridding artifact in ESP, the feature maps obtained using kernels of different dilation rates are hierarchically added before concatenating them (HFF). This solution is simple and effective and does not increase the complexity of the ESP module.

Hierarchical Feature Fusion

ESPNet: Efficient Spatial Pyramid of Dilated Convolutions for Semantic Segmentation

**Cascade R-CNN** is an object detection architecture that seeks to address problems with degrading performance with increased IoU thresholds (due to overfitting during training and inference-time mismatch between IoUs for which detector is optimal and the inputs). It is a multi-stage extension of the [R-CNN](https://paperswithcode.com/method/r-cnn), where detector stages deeper into the cascade are sequentially more selective against close false positives. The cascade of R-CNN stages are trained sequentially, using the output of one stage to train the next. This is motivated by the observation that the output IoU of a regressor is almost invariably better than the input IoU. 

Cascade R-CNN does not aim to mine hard negatives. Instead, by adjusting bounding boxes, each stage aims to find a good set of close false positives for training the next stage. When operating in this manner, a sequence of detectors adapted to increasingly higher IoUs can beat the overfitting problem, and thus be effectively trained. At inference, the same cascade procedure is applied. The progressively improved hypotheses are better matched to the increasing detector quality at each stage.

Cascade R-CNN

Cascade R-CNN: Delving into High Quality Object Detection

**BasicVSR** is a video super-resolution pipeline including optical flow and [residual blocks](https://paperswithcode.com/method/residual-connection). It adopts a typical bidirectional recurrent network. The upsampling module $U$ contains multiple [pixel-shuffle](https://paperswithcode.com/method/pixelshuffle) and convolutions. In the Figure, red and blue colors represent the backward and forward propagations, respectively.  The propagation branches contain only generic components. $S, W$, and $R$ refer to the flow estimation module, spatial warping module, and residual blocks, respectively.

BasicVSR

BasicVSR: The Search for Essential Components in Video Super-Resolution and Beyond

An **Inception Module** is an image model block that aims to approximate an optimal local sparse structure in a CNN. Put simply, it allows for us to use multiple types of filter size, instead of being restricted to a single filter size, in a single image block, which we then concatenate and pass onto the next layer.

Inception Module

Going Deeper with Convolutions

**Virtual Batch Normalization** is a normalization method used for training generative adversarial networks that extends batch normalization. Regular [batch normalization](https://paperswithcode.com/method/batch-normalization) causes the output of a neural network for an input example $\mathbf{x}$ to be highly dependent on several other inputs $\mathbf{x}'$ in the same minibatch. To avoid this problem in virtual batch normalization (VBN), each example $\mathbf{x}$ is normalized based on the statistics collected on a reference batch of examples that are chosen once and fixed at the start of training, and on $\mathbf{x}$ itself. The reference batch is normalized using only its own statistics. VBN is computationally expensive because it requires running forward propagation on two minibatches of data, so the authors use it only in the generator network.

Virtual Batch Normalization

Improved Techniques for Training GANs

**Prioritized Sweeping** is a reinforcement learning technique for model-based algorithms that prioritizes updates according to a measure of urgency, and performs these updates first. A queue is maintained of every state-action pair whose estimated value would change nontrivially if updated, prioritized by the size of the change. When the top pair in the queue is updated, the effect on each of its predecessor pairs is computed. If the effect is greater than some threshold, then the pair is inserted in the queue with the new priority.

Source: Sutton and Barto, Reinforcement Learning, 2nd Edition

Prioritized Sweeping

**G3D** is a unified spatial-temporal graph convolutional operator that directly models cross-spacetime joint dependencies. It leverages dense cross-spacetime edges as skip connections for direct information propagation across the 3D spatial-temporal graph.

Viet-Anh on Software

AI Glossary

Tech Terminology

SM3