# CSANet: Channel Spatial Attention Network for Robust 3D Face Alignment and
Reconstruction

Yilin Liu
Carnegie Mellon University
<EMAIL_ADDRESS>
& Xuezhou Guo
Carnegie Mellon University
<EMAIL_ADDRESS>
& Xinqi Wang
Mount-Sheikh University
<EMAIL_ADDRESS>
Fangzhou Du
Carnegie Mellon University
<EMAIL_ADDRESS>
###### Abstract

3D face reconstruction from monocular image has been a popular topic in the
computer vision community. Recent breakthroughs are mainly facilitated by face
alignment, which aims to locate face vertices on a registered model. Although
existing research works successfully tackle accuracy problems caused by large
head pose, they still fall short when the faces are partially occluded, or in
large light changing condition. Motivated by these challenges, this project
proposes an end-to-end 3D face alignment and reconstruction network. The
backbone of our model is built by BottleNeck structure via Depthwise Separable
Convolution. In order to improve the model’s performance in extreme
conditions, we integrate Coordinate Attention mechanism and Spatial Group-wise
Enhancement to extract more representative features in face regions. For more
stable training process and better convergence, we jointly use Wing loss and
the Weighted Parameter Distance Cost (WPDC) to learn parameters for 3D
Morphable model (3DMM) and 3D vertices. From the results, our proposed model
outperforms all baseline models both quantitatively and qualitatively.
Specifically, the face alignment accuracy is improved by 7.44% on AFLW and
10.16

_Keywords_ 3D Face Reconstruction $\cdot$ Computer Vision $\cdot$ Artificial
Intelligence

##  1 Introduction

During the past few years, 3D image processing is gaining more attention due
to the limitations of 2D image application. Many research fields benefit from
the booming of 3D images, such as robotics, security and bio-medical. 3D
images can associate with 2D images, while providing more enriched information
(e.g. space directions). Among all 3D image related tasks, 3D face
reconstruction is considered as one of the most challenging ones, since using
scanning device to obtain ground-truth 3D faces can be expensive and time-
consuming. Therefore, reconstructing 3D face models from a 2D face image has
been widely explored, as an alternative to synthesize accurate 3D faces with
limited budget. Generally, 3D face reconstruction can be decomposed into two
steps, the initial face alignment step and the face detail synthesis step [1].
The initial modeling mainly involves face alignment, which focuses on locating
facial vertices of certain face model, and face synthesis aims at recovering
3D facial geometry and details from the input image. In particular, the
accuracy of the alignment module has a significant effect on the overall
performance of the reconstruction model. Our project aims to improve the
performance of 3D face alignment frameworks. Furthermore, based on the
alignment results, we can generate more precise 3D face reconstruction
results. Initially, we re-implemented the 3D face reconstruction model 3DDFA
[2]. The authors proposed to use the Pose Adaptive Feature (PAF) and Projected
Normalized Coordinate Code (PNCC) to encode human face features. However,
these two processes turned out to involve tremendous computation, which is not
efficient during evaluation and interference, thus making it impossible to be
applied in real-world applications. Besides, when the face image is partially
occluded (e.g. by hair or sunglasses) or in large lighting change conditions,
3DDFA and some other proposed methods struggle to generate desirable alignment
and reconstruction results. To address these problems, we improve the baseline
model in three aspects. Firstly, inspired by the MobileNet [3] structure and
DAMDNet [4], we adopt Depth-wise Separable Convolution instead of normal
convolution for efficient computation. Then, we integrate the state-of-the-art
attention mechanisms along both the spatial and channel dimensions to deal
with faces with occlusions and lighting condition changes. Lastly, two
different losses are jointly used to learn parameters for the projected 3D
Morphable Model (3DMM) [5, 6]. Extensive experiments are conducted on the AFLW
and AFLW2000-3D benchmark and the performance of the proposed model is then
compared with that of other baseline models.

##  2 Related Work

###  2.1 Face Alignment

Face alignment aims to accurately locate facial vertices. It can be considered
as a model parameter optimization problem [7]. In the past, face alignment
focuses on detecting sparse 2D facial landmarks in the image plane. Many
traditional works [8, 9] are based on the Active Appearance Models (AMM) [10]
or Active Shape Models (ASM) [11], which only represent shape/texture
variations with linear subspace. To better model non-linear pose variations,
such as face rotation, some improvements are proposed to add non-linearity
into face model [12, 13], or into the model fitting formulation [14]. Despite
the introduction of non-linearity to improve the representative ability of 2D
face alignment methods, there are still limitations on detecting invisible
landmarks (e.g. occlusions, profile view), since 2D points lose geometric
correspondence with actual locations. Therefore, 3D face alignment approaches
have become the mainstream research in recent years. 3D face alignment can be
separated into 2 strategies. (1) Bypass constructing and fitting face models.
[33, 35] directly regress all 3D point coordinates by a fully convolutional
neural network. However, the resolution of resulted faces depends on the
feature map size, and it is not time-effective during inference [15]. (2)
Implement parameter regression of a certain 3D face model [16, 17, 18, 19, 20]
to obtain a full 3D structure, in order to guide the actual landmark
localization. Among all 3D face representing methods, 3D Morphable Model
(3DMM) [21, 22] is the most widely-used one. 3DMM is a vector space of linear
combinations of objects’ shapes (i.e. vertex coordinates) and textures (i.e.
RGB colors). It applies Principle Component Analysis (PCA) to 3D face scans.
The non-linear face transformations are separate from linear subspaces, which
models linear transformations for shapes or expressions.

###  2.2 Face Reconstruction

There are two main categories of 3D face reconstruction: model-based and
model-free approaches. The model-based reconstruction methods are highly
related with 3D dense face alignment if the face model is well registered.
Initially, earlier studies [23, 24, 25, 26] reconstruct 3D faces by regressing
3DMM parameters and as solving a nonlinear optimization problem. Later, works
[27, 28, 29] based on Convolution Neural Networks (CNNs) started to replace
traditional parameter-learning methods with more accurate results. Since the
reconstruction geometry is limited by linear bases of 3DMM, some works
proposed to use non-linear models. For instance, [30] learns a nonlinear 3DMM
decoder to reconstruct 3D faces from in-the-wild face images, instead of face
scans. Similarly, [31] learn a nonlinear 3DMM from diverse sources, including
RGB-D images. Some other works [32, 33, 34] directly estimate 3D face geometry
without depending on parametric models. These methods do not require accurate
alignment or dense correspondence. [35] adopts a volumetric representation of
3D face shape and directly regressed the voxels.

###  2.3 Attention

Attention mechanism was first introduced in natural language processing field.
Due to its effectiveness in feature extraction, it has also been deployed in
various visual-based research works [36, 37, 38]. The visual-based attention
mechanism has shown the potential ability in refining extracted features
obtained by convolutions. Specifically, existing attention modules mainly
refines features along the channel and spatial dimensions. The channel
attention determines which features are more important, while spatial
attention focuses on which portion of the image can be more informative. By
applying attention mechanism, the model can adaptively ignore the less
important information by enhanced feature representation. In this project, we
improve the baseline model by adding two attention mechanisms to capture the
relation between channel and spatial feature groups. One is Spatial Groupwise
Enhancement(SGE) [39] added inside each layer to enhance spatial group-wise
feature capture, another is Coordinate Attention (CA) [40] added between each
layer to aggregate features along two spatial coordinate directions
respectively.

##  3 Baseline Implementation

Among existing popular frameworks [33, 35, 34, 31, 2], we chose 3DDFA [2] as
the face alignment baseline model to deal with large pose 3D regression. While
being one of the state-of-the-art solutions for large-pose face alignment,
3DDFA is also one of the fundamental models, which is frequently cited by
other 3DDFA-based variations or novel face alignment and reconstruction
techniques for comparison and benchmark performance evaluation.

###  3.1 3D Morphable Model

3DMM [5, 6] is a way of representing face shape with the formula:

$S=\overline{S}+A_{id}\alpha_{id}+A_{exp}\alpha_{exp}$ (1)

, where $\overline{S}$ is the average face, $A_{id}$ is the principle
component of training 3D face scans, $\alpha_{id}$ is the shape parameter,
$A_{exp}$ is the principle axes trained on the offsets between expression
scans and neutral scans, and $\alpha_{exp}$ is the expression parameter. After
attaining the 3D face model, it can be projected onto 2D plane with the
transformation formula:

$V(p)=f*Pr*R*(\overline{S}+A_{id}\alpha_{id}+A_{exp}\alpha_{exp})+t_{2d}$ (2)

, where $Pr=\begin{pmatrix}1&0&0\\\ 0&1&0\end{pmatrix}$ represents the
orthographic projection matrix, $f$ if the scale factor, $R$ is the rotation
matrix and $t_{2d}$ is the translation vector.

###  3.2 Network Structure and Cost Function

In 3DDFA, the main pipeline is a two-stream CNN network designed with unified
network structure across the cascade [2], as shown in Figure 1. For each
iteration, the network first constructs the Projected Normalized Coordinate
Code (PNCC) and the Pose Adaptive Feature (PAF) [2] based on current
parameters for input images. The calculated PNCC and PAF data, combined with
the original face images, are used to train the two streams. To be more
specific, the first CNN stream is trained with input images combined with PNCC
values, while the second CNN stream is trained with input images combined with
PAF values. The outputs of the two CNN streams are concatenated and passed
through fully connected layers to generate the predicted parameter update
values $\Delta p$ that will be used to update the network [2].

Figure 1: 3DDFA Pipeline

In the original work, three cost functions are discussed: Parameter Distance
Cost (PDC), Vertex Distance Cost (VDC), and Weighted Parameter Distance Cost
(WPDC) [2]. Among the three cost functions, the Weighted Parameter Distance
Cost (WPDC) proposed in the paper showed the best performance as the loss
function. WPDC aims at shrinking the difference between the current parameters
$p_{0}+\Delta p$ and the ground-truth $p_{g}$ while assigning different
weights to different parameters to emphasize their importance [2].

###  3.3 Implementation Results

Our re-implementation of 3DDFA baseline model is based on the source code
published by the original authors [15]. The reason that we re-implemented the
original proposed method with source code published by original authors is
that the published implementation does not feature the same techniques stated
in the original paper. The authors also continuously updated their
implementation with some strategies not discussed in the paper, and optimized
their design targeting inference efficiency instead of accuracy by reducing
the two-stream CNN structure to single-stream with more advanced CNN
structure. Since we could not verify the impact of these changes on face
alignment and reconstruction accuracy, we discarded efficiency-targeting
changes not described in the original paper and re-implement the two-stream
CNN as the paper proposed. This ensure accurate comparison among methods.

Test results are obtained by training from scratch with the same data provided
by the original authors, and are compared with that obtained by model trained
with released source code, released pre-trained model, as well as that the
paper reported values as shown in Table 1. As the results show, our re-
implemented version of the two-branch structure performed worse on the
benchmark compared to the published pre-trained models and the results
reported in the paper, but slightly better than the model we obtained strictly
following the source code and training procedures release by the authors. One
possible reason for this is variance and lack of parameters tuning. We got
better performance with PDC and worse with WPDC comparing to the reported
values. Still, as the original authors stated in the documentation of their
publish source code, accuracy may vary depending on selection of training and
testing datasets [15], and our test results are within the error range.
Therefore, we consider our implementation as valid.

Table 1: Comparison of the Normalized Mean Error (NME%) of face alignment
results
on AFLW2000-3D using WPDC loss Our re-Implementation | Published Source Code | Published Pre-trained | Paper
---|---|---|---
$4.759$ | $4.823$ |  $4.252$ $\pm$ $0.996$ | $4.05$

##  4 Channel Spatial Attention Network

###  4.1 Model Structure

The network architecture for our model is shown in Figure 2. The main pipeline
is inspired by [4]. The input to our model is a single RGB face image, and the
output is a 62-dimensional 3D parameter vector, including 40-dimensional
identity parameter vector, 10-dimensional expression parameter vector and
12-dimensional pose vector. There are 7 main layers in the model backbone.
Each layer is a cascade of several units. Each unit is shown in the red box in
Figure 2. It is a bottleneck structure followed by a Spatial Group-wise
Enhancement (SGE) module. Between each layer, we add another attention block:
coordinate attention. Both attention mechanisms and the cost function will be
discussed in more details in the following sections.

Figure 2: The proposed model structure

###  4.2 Attention Modules

####  4.2.1 Spatial Group-wise Enhance Module

In face analysis tasks, CNNs are widely used to extract semantic features that
can be grouped along the channel-dimension as sub-features [39]. However,
these approaches overlooked the spatial dimension of the convolution feature
maps. In fact, it would be desired to generate the corresponding features from
the correct location in the original image for a particular semantic group.
Still, this can be made challenging by noises in the background and image
patterns. To account for the spatial dimension, we employed the Spatial Group-
wise Enhance (SGE) [39] module shown in Figure 3 to strengthen the the spatial
distribution of facial semantic features. This attention mechanism treats each
feature layer as a sub-feature group, and generates an attention factor based
on the similarity of global and local features within each group. First, the
feature layers will be grouped along the channel dimension and each group is
examined. Ideally, this group contains features representing a specific
semantic response such as eyes or mouths, but this is obscured by the noises
in the actual image. To account for this instability, it is better to use the
entire group space to improve the learning of the critical regions. The global
statistical feature can be obtained as

$g=F_{gp}(X)=\frac{1}{m}\sum^{m}_{i=1}x_{i}$ (3)

where $F_{gp}(\cdot)$ is a spatial averaging function. With the global
feature, we calculate the initial attention mask by spatially comparing each
local feature group $x_{i}$ with the global pooling feature $g$, which can be
obtained by the dot product that measure the similarity between the them.
Next, the final attention mask $a_{i}$ is obtained by a sigmoid operation on
the normalized initial attention mask, and the enhanced features $\hat{x}_{i}$
can be generated as

$\hat{x}_{i}=x_{i}\cdot\sigma(a_{i})$ (4)

Figure 3: Spatial Group-wise Enhance Module(SGE)

####  4.2.2 Coordinate Attention Module

To better represent precise positional information while also effectively
capturing inter-channel relationships, Coordinate Attention [40] blocks
illustrated in Figure 4 are introduced into the network. Specifically, it will
squeeze spatial content by two pooling kernels along the horizontal and
vertical coordinate to encode each channel information. This will yield a pair
of feature maps, which aggregate features along two spatial directions
respectively. After encoding the positional information, we generate
intermediate feature maps by applying $1\times 1$ convolution $F_{1}$ on the
concatenated feature maps as

$f=\delta(F_{1}([z^{h},z^{w}]))$ (5)

where $z^{h}$ and $z^{w}$ are the generated feature maps and $\delta$ is a
non-linear Swish activation function, which is simply $x\ast sigmoid(x)$. $f$
is then split along the spatial dimension into $f^{h}$ and $f^{w}$ and
processed again with $1\times 1$ convolutions followed by sigmoid activation,
giving

$g^{h}=\sigma(F_{h}(f^{h}))$ (6) $g^{w}=\sigma(F_{w}(f^{w}))$ (7)

Lastly, the attentions are applied to the input of the block as

$y_{c}(i,j)=x_{c}(i,j)\times g^{c}_{c}(i)\times g^{w}_{c}(j)$ (8)

Figure 4: Coordinate Attention Module(CA)

###  4.3 Loss Function

In this project, we adopt a joint loss [4], which is composed of the Wing Loss
[41] and the WPDC loss [2]. The WPDC loss, as discussed before, reduces the
differences between the generated face model parameters and the target
parameters; the Wing loss improves the accuracy of CNN-based facial landmark
localisation. It deals with both large errors that occur at the beginning of
the training process and smaller errors during later phases. They are combined
as

$loss_{merged}=loss_{wing}+0.5loss_{wpdc}$ (9)

##  5 Experiments

###  5.1 Implementation

Our model is implemented using the PyTorch library. The training and
validation processes are implemented on Amazon EC2 g4dn.2xlarge instances. The
model was trained for 50 epochs, approximately 25 hours to achieve the
obtained results. For 3DDFA baseline re-implementation, on the same device, it
costs about 50 hours to train. Therefore, our proposed model is two times
faster in terms of training speed. We used Stochastic Gradient Descent (SGD)
as the parameter optimizer with and initial learning rate at 0.02, and reduced
it by a factor of $0.6$ after no significant improvements for 4 epochs.
Details of training/validation datasets and benchmark evaluation are discussed
in the following sections.

###  5.2 Dataset and Data Transformation

We used 300W-LP as the training dataset. It is a synthesized face alignment
dataset containing 122,450 samples with large head poses, which are generated
by profiling the face images in the original 300W dataset. 300W dataset
consists of multiple re-annotated in-the-wild datasets with 68 facial
landmarks.

We evaluated the trained model on the AFLW and the AFLW2000-3D datasets. AFLW
is a large-scale collection of annotated face images (21,080 in-the-wild
faces) gathered from the web, exhibiting a large variety in appearance (e.g.,
pose, expression, gender). Each image is annotated with 21 landmarks.
AFLW2000-3D is developed from AFLW dataset. It contains the ground truth 3D
face scans and the corresponding 68 landmarks for the first 2,000 AFLW
samples. This dataset fills in the gap between mapping from 2D landmarks to
its 3D scans.

In addition, images from all datasets are pre-processed before feeding to the
network. Each image is cropped according to the region of interest, which is
the bounded face region in the original image. The standard size of cropped
images are $120\times 120$. During the benchmark evaluation, images are also
categorized by yaw angles, representing small, medium and large poses
respectively. For example, yaw angle within $[0^{\circ},30^{\circ}]$ can be
regarded as small head movement.

###  5.3 Evaluation Metric

Normalized Mean Error (NME) is used to evaluate the face alignment performance
of the network. It is defined as the normalized mean Euclidean distance
between corresponding points in the predicted result $p$ and the ground truth
$\hat{p}$.

$NME=\frac{1}{N}\sum_{i=1}^{N}\frac{\|p_{i}-\hat{p}_{i}\|_{2}}{d}$ (10)

where the normalization factor $d$ is the bounding box size of ground truth
landmarks.

For face reconstruction results, we fitted the obtained 3D parameters into
3DMM and visualize it in 3D vector space. The face reconstruction results are
consistent with face alignment landmarks and further discussed in the Results
section.

##  6 Results

The model performance is evaluated both quantitatively and qualitatively.
Numerical results are shown in Table 2. We compare a series of 3DDFA models
with different cost functions, one of the state-of-the-art face reconstruction
model named DAMDNet and our proposed model. The baseline 3DDFA model is the
published version, which is trained by Weighted Parameter Distance Cost (WPDC)
loss and fine-tuned on Vertex Distance Cost (VDC), denoted as 3DDFA-WPDC-VDC.
Other 3DDFA models are trained by Parameter Distance Cost (PDC), VDC and WPDC
separately, denoted by the name of adopted loss function in the table.
Besides, we also include the evaluation of the proposed model without SGE
blocks as a comparison, to validate the effectiveness of the incorporation of
SGE into the model.From the result, our proposed network outperforms other
baseline models, especially in large pose conditions. Compared with the chosen
baseline model 3DDFA, the face alignment accuracy is improved by 7.44% on the
AFLW dataset and 10.16% on the AFLW2000-3D dataset.

Table 2: The NME(%) of face alignment results on AFLW and AFLW2000-3D with the
first and the second best results highlighted.

AFLW Dataset AFLW2000-3D Dataset Method ${[}0^{\circ},30^{\circ}]$
${[}30^{\circ},60^{\circ}]$ ${[}60^{\circ},90^{\circ}]$ Mean Std
${[}0^{\circ},30^{\circ}]$ ${[}30^{\circ},60^{\circ}]$
${[}60^{\circ},90^{\circ}]$ Mean Std 3DDFA-PDC $6.937$ $8.441$ $9.175$ $8.184$
$0.932$ $5.254$ $6.673$ $8.123$ $6.683$ $1.172$ 3DDFA-VDC $5.756$ $7.222$
$8.072$ 7.017 $0.957$ $4.040$ $5.207$ $6.654$ $5.300$ $1.069$ 3DDFA-WPDC
$5.282$ $6.606$ $7.464$ $6.451$ $0.898$ $3.658$ $4.857$ $5.962$ $4.826$
$0.941$ 3DDFA-WPDC-VDC 4.488 $5.381$ $6.334$ 5.401 $0.754$ $3.099$ $4.172$
$5.486$ $4.252$ $0.976$ 3DDFA-ours $5.206$ $6.496$ $6.975$ $6.226$ $0.747$
$3.897$ $4.783$ $5.788$ $4.823$ $0.733$ DAMDNet $4.539$ 5.209 6.028 5.199
$0.682$ 2.907 3.830 4.953 3.897 $0.837$ Ours-CA $4.978$ $6.072$ $6.885$ 5.978
$0.781$ $3.440$ $4.300$ $5.376$ $4.372$ $0.792$ Ours 4.254 4.947 5.796 4.999
$0.630$ 2.914 3.738 4.809 3.820 $0.776$

For quantitative result, we visualize reconstructed faces on AFLW2000 dataset,
as shown in Figure 5. In face occlusion conditions (e.g. hair-covered face
region, incomplete image), our proposed model can achieve an obviously better
result, which captures more accurate face contour, eye corner and nose tips.
The most prominent example can be observed in the figure where a partially
present face is also blocked by sunglasses. Both the 3DDFA model and DAMDNet
struggled to generate a reasonable face shapes, whereas our proposed network
successfully detects the face and generates desirable reconstruction results.
Apart from these face details, our model can generate a more reasonable facial
expression and face angles in conditions with large light change Figure 6.

Figure 5: Qualitative Visual Results of 3DDFA(a), DAMDNet(b) and our proposed
model(c) on AFLW2000-3D dataset. Face alignment landmarks are shown on RGB
images, with a consistent face reconstruction result

##  7 Future Work

Although our proposed model was able to increase face alignment accuracy on
face images with partial occlusion or significant light variation, our model
can still be improved to achieve better results. From Figure6 it can be seen
that the dark-side face is not accurate enough. First, since the backbone of
our model is inspired by MobileNet [3], a relatively lightweight network, the
ability of our model to capture complex face features is limited. Therefore,
increasing the network capacity might be an appropriate step to improve
detection accuracy and capture of feature details. Another potential approach
is to experiment with loss functions with tunable parameters. Although the
weight parameter of the merged loss function can be experimented, the effects
did not prove to be very significant. Thus, in future experiments, tunable
loss functions may be explored to improve the training of the model.

Our Current method utilizes attention methods to improve robustness against
occlusions on a single image. However, many tasks today require robust 3D face
alignment and reconstruction on image sequences (e.g. videos), whereas our
current model does not take the advantage of having multiple images as input
that correspond to the same alignment target. In a sequence of images, the
view angle of the alignment target may change, or the occlusion region may
shift from image to image, resulting in more information provided than any
arbitrary image within the sequence. For future work, we might investigate the
possibility of developing a network that focuses on image sequences as inputs,
such as learning a pose-independent 3D face alignment for the target and
update the alignment for every image within the image sequence, to achieve
better alignment accuracy and robustness.

##  8 Conclusions

In conclusion, in this project, we investigated the state-of-the-art
researches in the field of 3D face alignment and reconstruction, specifically
focusing on 3D face alignment. We first implemented the 3DDFA network and
found that the proposed network can be computationally expensive,
timeconsuming and performs relatively unsatisfactory results. Therefore,
inspired by recent break through achieved by attention neural networks, we
proposed a network with state-of-the-art attention mechanisms, namely the
Spatial Group-wise Enhancement(SGE) and the Coordinate Attention Module(CA),
to improve the performance of the model when faces in images are partially
blocked, or when lighting in the image drastically changes. We also adopted a
joint loss of Weighted Parameter Distance Cost(WPDC) and Wing Loss to improve
convergence and stabilize parameter updates. Based on the test results on
benchmark datasests, our proposed CSANet achieves an improvement on accuracy
of 7.44% on the AFLW dataset and 10.16% on the AFLW2000-300D dataset. This
work can be applied to generate 3D face scans without human efforts and
expensive device.

Figure 6: Compared results of 3DDFA, DAMD, our proposed model from left to
right

## References

  * [1] Li, T., Xu, H., Zhang, H. & Wan, H. Detail 3D Face Reconstruction Based on 3DMM and Displacement Map. Journal Of Sensors. 2021 pp. 9921101 (2021,6), https://doi.org/10.1155/2021/9921101
  * [2] Zhu, X., Liu, X., Lei, Z. & Li, S. Face Alignment in Full Pose Range: A 3D Total Solution. IEEE Transactions On Pattern Analysis And Machine Intelligence. 41, 78-92 (2019)
  * [3] Howard, A., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M. & Adam, H. MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications. CoRR. abs/1704.04861 (2017), http://arxiv.org/abs/1704.04861
  * [4] Jiang, L., Wu, X. & Kittler, J. Dual Attention MobDenseNet(DAMDNet) for Robust 3D Face Alignment. CoRR. abs/1908.11821 (2019), http://arxiv.org/abs/1908.11821
  * [5] Blanz, V. & Vetter, T. A Morphable Model For The Synthesis Of 3D Faces. Seminal Graphics Papers: Pushing The Boundaries, Volume 2. (2023), https://doi.org/10.1145/3596711.3596730
  * [6] Blanz, V. & Vetter, T. Face recognition based on fitting a 3D morphable model. IEEE Transactions On Pattern Analysis And Machine Intelligence. 25, 1063-1074 (2003)
  * [7] Ruan, Z., Zou, C., Wu, L., Wu, G. & Wang, L. SADRNet: Self-Aligned Dual Face Regression Networks for Robust 3D Dense Face Alignment and Reconstruction. CoRR. abs/2106.03021 (2021), https://arxiv.org/abs/2106.03021
  * [8] Sauer, P., Cootes, T., Taylor, C. & Others Accurate Regression Procedures for Active Appearance Models.. Bmvc. pp. 1-11 (2011)
  * [9] Sung, J. & Kim, D. Adaptive active appearance model with incremental learning. Pattern Recognition Letters. 30, 359-367 (2009), https://www.sciencedirect.com/science/article/pii/S0167865508003504
  * [10] Cootes, T., Edwards, G. & Taylor, C. Active appearance models. IEEE Transactions On Pattern Analysis And Machine Intelligence. 23, 681-685 (2001)
  * [11] Cootes, T., Taylor, C., Cooper, D. & Graham, J. Active shape models-their training and application. Computer Vision And Image Understanding. 61, 38-59 (1995)
  * [12] Romdhani, S., Gong, S., Psarrou, A. & Others A Multi-View Nonlinear Active Shape Model Using Kernel PCA.. BMVC. 10 pp. 483-492 (1999)
  * [13] Zhou, Y., Zhang, W., Tang, X. & Shum, H. A bayesian mixture model for multi-view face alignment. 2005 IEEE Computer Society Conference On Computer Vision And Pattern Recognition (CVPR’05). 2 pp. 741-746 (2005)
  * [14] Saragih, J. & Goecke, R. A nonlinear discriminative approach to AAM fitting. 2007 IEEE 11th International Conference On Computer Vision. pp. 1-8 (2007)
  * [15] Guo, J., Zhu, X., Yang, Y., Yang, F., Lei, Z. & Li, S. Towards Fast, Accurate and Stable 3D Dense Face Alignment. CoRR. abs/2009.09960 (2020), https://arxiv.org/abs/2009.09960
  * [16] Liu, F., Zhao, Q., Liu, X. & Zeng, D. Joint Face Alignment and 3D Face Reconstruction with Application to Face Recognition. IEEE Transactions On Pattern Analysis And Machine Intelligence. 42, 664-678 (2020)
  * [17] Liu, Y., Jourabloo, A., Ren, W. & Liu, X. Dense Face Alignment. (2017)
  * [18] McDonagh, J. & Tzimiropoulos, G. Joint face detection and alignment with a deformable hough transform model. Computer Vision–ECCV 2016 Workshops: Amsterdam, The Netherlands, October 8-10 And 15-16, 2016, Proceedings, Part II 14. pp. 569-580 (2016)
  * [19] Tran, A., Hassner, T., Masi, I. & Medioni, G. Regressing Robust and Discriminative 3D Morphable Models with a very Deep Neural Network. CoRR. abs/1612.04904 (2016), http://arxiv.org/abs/1612.04904
  * [20] Zhu, X., Liu, X., Lei, Z. & Li, S. Face Alignment in Full Pose Range: A 3D Total Solution. CoRR. abs/1804.01005 (2018), http://arxiv.org/abs/1804.01005
  * [21] Blanz, V. & Vetter, T. A morphable model for the synthesis of 3D faces. Seminal Graphics Papers: Pushing The Boundaries, Volume 2. pp. 157-164 (2023)
  * [22] Blanz, V. & Vetter, T. Face recognition based on fitting a 3D morphable model. IEEE Transactions On Pattern Analysis And Machine Intelligence. 25, 1063-1074 (2003)
  * [23] Huber, P., Feng, Z., Christmas, W., Kittler, J. & Rätsch, M. Fitting 3d morphable face models using local features. 2015 IEEE International Conference On Image Processing (ICIP). pp. 1195-1199 (2015)
  * [24] Lee, Y., Lee, S., Park, K., Jo, J. & Kim, J. Single view-based 3D face reconstruction robust to self-occlusion. EURASIP Journal On Advances In Signal Processing. 2012 pp. 1-20 (2012)
  * [25] Paysan, P., Knothe, R., Amberg, B., Romdhani, S. & Vetter, T. A 3D face model for pose and illumination invariant face recognition. 2009 Sixth IEEE International Conference On Advanced Video And Signal Based Surveillance. pp. 296-301 (2009)
  * [26] Romdhani, S. & Vetter, T. Estimating 3D shape and texture using pixel intensity, edges, specular highlights, texture constraints and a prior. 2005 IEEE Computer Society Conference On Computer Vision And Pattern Recognition (CVPR’05). 2 pp. 986-993 vol. 2 (2005)
  * [27] Jourabloo, A. & Liu, X. Large-pose face alignment via CNN-based dense 3D model fitting. Proceedings Of The IEEE Conference On Computer Vision And Pattern Recognition. pp. 4188-4196 (2016)
  * [28] Tran, A., Hassner, T., Masi, I. & Medioni, G. Regressing robust and discriminative 3d morphable models with a very deep neural network. 2017 IEEE Conference On Computer Vision And Pattern Recognition (CVPR). pp. 1493-1502 (2017)
  * [29] Yi, H., Li, C., Cao, Q., Shen, X., Li, S., Wang, G. & Tai, Y. MMFace: A Multi-Metric Regression Network for Unconstrained Face Reconstruction. 2019 IEEE/CVF Conference On Computer Vision And Pattern Recognition (CVPR). pp. 7655-7664 (2019)
  * [30] Tran, L. & Liu, X. On learning 3d face morphable model from in-the-wild images. IEEE Transactions On Pattern Analysis And Machine Intelligence. 43, 157-171 (2019)
  * [31] Guo, Y., Zhang, J., Cai, L., Cai, J. & Zheng, J. Self-supervised CNN for Unconstrained 3D Facial Performance Capture from an RGB-D Camera. CoRR. abs/1808.05323 (2018), http://arxiv.org/abs/1808.05323
  * [32] Fan, X., Cheng, S., Huyan, K., Hou, M., Liu, R. & Luo, Z. Dual neural networks coupling data regression with explicit priors for monocular 3D face reconstruction. IEEE Transactions On Multimedia. 23 pp. 1252-1263 (2020)
  * [33] Feng, Y., Wu, F., Shao, X., Wang, Y. & Zhou, X. Joint 3D Face Reconstruction and Dense Alignment with Position Map Regression Network. CoRR. abs/1803.07835 (2018), http://arxiv.org/abs/1803.07835
  * [34] Ranjan, A., Bolkart, T., Sanyal, S. & Black, M. Generating 3D faces using Convolutional Mesh Autoencoders. CoRR. abs/1807.10267 (2018), http://arxiv.org/abs/1807.10267
  * [35] Jackson, A., Bulat, A., Argyriou, V. & Tzimiropoulos, G. Large Pose 3D Face Reconstruction from a Single Image via Direct Volumetric CNN Regression. CoRR. abs/1703.07834 (2017), http://arxiv.org/abs/1703.07834
  * [36] Hu, J., Shen, L. & Sun, G. Squeeze-and-Excitation Networks. Proceedings Of The IEEE Conference On Computer Vision And Pattern Recognition (CVPR). (2018,6)
  * [37] Park, J., Woo, S., Lee, J. & Kweon, I. BAM: Bottleneck Attention Module. CoRR. abs/1807.06514 (2018), http://arxiv.org/abs/1807.06514
  * [38] Woo, S., Park, J., Lee, J. & Kweon, I. Cbam: Convolutional block attention module. Proceedings Of The European Conference On Computer Vision (ECCV). pp. 3-19 (2018)
  * [39] Li, X., Hu, X. & Yang, J. Spatial group-wise enhance: Improving semantic feature learning in convolutional networks. ArXiv Preprint ArXiv:1905.09646. (2019)
  * [40] Hou, Q., Zhou, D. & Feng, J. Coordinate attention for efficient mobile network design. Proceedings Of The IEEE/CVF Conference On Computer Vision And Pattern Recognition. pp. 13713-13722 (2021)
  * [41] Feng, Z., Kittler, J., Awais, M., Huber, P. & Wu, X. Wing Loss for Robust Facial Landmark Localisation with Convolutional Neural Networks. CoRR. abs/1711.06753 (2017), http://arxiv.org/abs/1711.06753


* Utilizing the geometric nature and nonlinear relations, we have first derived the expected second-order gauge-transformation properties of our theoretical descriptions of the observable quantities and the intermediate quantities. With the theoretical descriptions in a general metric representation, we have next performed the second-order gauge-transformation of the metric perturbation variables in those theoretical descriptions. By comparing the direct gauge-transformation properties to the expectations from the nonlinear relations, we have explicitly verified the consistency of our theoretical expressions, among which only the expressions for the physical volume and area occupied by the source are gauge-invariant <cit.>.

* We have compared our results to those in the previous work and discussed the physical ingredients that are missed or ignored in the previous work.[As the current work was near completion, a work by <cit.> appeared on arXiv in which they present a nonlinear framework for computing the cosmological observables, but no second-order expressions or verification of gauge invariance are presented.] Among many differences, we have highlighted that the contributions at the observer position are typically neglected in literature and as a consequence all the quadratic terms that couple one contribution at the observer position and another contribution at the second order are also ignored. With any contributions missing, the second-order expressions in literature are incomplete and gauge-dependent.

Given the technical difficulties and theoretical subtleties in the second-order relativistic perturbation calculations of the cosmological observables, the second-order expressions in the previous work (including one of our own)
are generally considered incomplete with a few or several missing relativistic contributions. The explicit check of second-order gauge-invariance of the theoretical descriptions that we have performed for the first time constitutes a major step forward in establishing the complete second-order theoretical descriptions of the cosmological observables. With full degrees of freedom intact, our second-order expressions can be used to describe the cosmological observables in any theories of metric gravity. In the future, we will apply our second-order gauge-invariant formalism to compute the three-point correlation function of cosmological observables such as galaxy clustering, the luminosity distance, weak gravitational lensing, and cosmic microwave background anisotropies. Given the interest in the primordial non-Gaussianity, our complete second-order relativistic expressions will play a crucial role in providing accurate predictions for the higher-order statistics.

§ ACKNOWLEDGMENT

We acknowledge support by a Consolidator Grant of the European Research Council
(ERC-2015-CoG grant 680886).

§ GAUGE TRANSFORMATION OF THE LIGHT WAVE VECTOR AT THE OBSERVER POSITION

Here we present the consistency check for the components of the light wave vector $\hat k^\mu_o$ at the observer position.
We start by computing the second-order expectation for the gauge transformations of $\dnu_o$ and $\dea^\al_o$. Taking the limit for $P\to O$ of Eq. ($\ref{dnuP}$) and ($\ref{deaP}$) we obtain

_o:=(x=x_o)    =   _o+[ -d/dT + Td/dT' +T(1+-d/dT)

-  ℒ^ (_,+2 ℋT_,-d/dT_,  )+ 1/2(ℋ^2 - 3ℋ')T^2 -  2ℋTT'

+T'-T_,^-T' ]_o

^_o:=^(x=x_o)   =   ^_o+[d/d^-Td/d'^+T(n^+^α+d/d^)

-  ^β( ^_,+2 ℋ  T_,βn^α+d/d^_, )+ 1/2( ℋ^2 - 3 ℋ' ) T^2n^α
-2TT'n^-'^+^_,^-T'^]_o .
We emphasize that both sides of the equations are evaluated at the coordinate value $x^\mu_o$, while the coordinates of the reference position $O$ are different in the two charts.
An independent way to derive the gauge transformation of $\dnu_o$ can be obtained considering the perturbation $\dhnu$. Let us rewrite the nonlinear expression in Eq. ($\ref{dhnuexact}$) in the new coordinate system denoted by a tilde:
=   2Å+_-_+(1+2Å)(++ )+^( _-_ )+^_(1+)

+2^(_+_^) ,
where both sides are evaluated at the coordinates $\widetilde x_p^\mu$ of a generic point $P$ on the light path. Given our choice for the normalization constant in Eq. ($\ref{NC}$) and the relation ($\ref{dhnu}$), we readily see that the evaluation at $P=O$ implies $\widetilde\dhnu(\widetilde x_o)=0$. Hence, we can solve the equation above for the perturbation $\widetilde{\dnu}(\widetilde x_o$):
(x_o)   =   -[2Å+_-_++ +2Å(++ )+^( _-_ )

+2^(_+_^)](x_o) .
We are interested in the evaluation at $x^\mu_o$ since these coordinates correspond to the observer position in the original chart. Therefore, we expand the coordinates $\widetilde x^\mu_o$ around the original coordinates $x^\mu_o$ according to the coordinate transformation $\widetilde x^\mu(x)$ in Eq. ($\ref{xcsi}$). At second order we obtain
(x=x_o)   =   -

+Å _+_( ^-^)
+1/2 ^_+2_^](x= x_o)    ,
where we used the linear-order gauge transformation of $\dnu$. Employing the linear-order definition of $\dhnu$ from Eq. ($\ref{dhnuexact}$), we can see that the above equation yields
(x_o)   =   -
[ξ^μ_μ(+T )+Å+_-+Å-1/2Å^2

+Å _+_( ^-^)
+1/2 ^_+2_^](x_o)    ,
Replacing all the gauge transformations with the results of Section $\ref{sec:gauge}$ we finally derive
_o   =   _o+[ -d/dT + Td/dT' +T(1+-d/dT)
-  ℒ^ (_,+2 ℋT_,-d/dT_,  )

+ 1/2(ℋ^2 - 3ℋ')T^2 -  2ℋTT'
+T'-T_,^-T' ]_o
which agrees with the previous computation in Eq. ($\ref{dnuO}$).

We now want to check that our second-order expression for $\hat k^\mu_o$ gauge transforms consistently. Note that we have a second-order expression for $\dnu_o$ in Eq. ($\ref{dnu2o}$), but we do not have a complete second-order expression for $\dea^\al_o$ (see Eq. ($\ref{dea2o}$)). However, its relevant contribution to second-order is given by $S_{ij}$ in Eq. ($\ref{tetrade2}$), so we only need to prove the consistency of such expression.

To derive the second-order expectations we consider that at any spacetime point $x^\mu$ the tetrad components $[e_a]^\mu$ transform as a four-vector under diffeomorphisms of $\mM$. The corresponding gauge transformation can be read from Eq. ($\ref{gauge2}$). Using the definition in Eq. ($\ref{tetrade}$) for the space-like tetrad components, we obtain at second order:

[e_i]^α   =   [e_i]^α+1/a {
δ^α_i (ℋ T) + δ^β_i
^α_, +
δ^β_i  '^ (_β- _β)
- ^ α_,
[δ^γ_i (^_+C^[_,]) + ϵ^β_ik Ω^k  ]

+ ( T∂_η+ ^β∂_β- ℋ T )
[ δ^γ_i (^_+C^[_,]) + ϵ^α_ik Ω^k]
-ℋ T T'δ^α_i -  ℋ ^β T_, δ^α_i

- 1/2(ℋ' - ℋ^2) T^2 δ^α_i -T '^α_, δ^_i
-^ ^ α_, δ^_i + ℋT ^α_, δ^_i   } ,
where both sides are evaluated at the same coordinate value $x_o^\mu$. Taking the symmetric and antisymmetric part of the last equation and contracting the Kronecker deltas, we derive

S_ij-S_ij   =   -_ijT-_(i,j)+_(j'_i)-'_ijT+T_ij-1/2δ_ij^2T^2+1/2_ij'T^2

A_ij-A_ij   =   -_[i,j]+'_[i(_j]-_j])-T( C'_[i,j]-'_[i,j]+ε_ijk'^k )

+T ( C_[i,j]-_[i,j]+ ε_ijk^k) -^(C_[i,j]-_[i,j]+ ε_ijl^l_,)

according to the definition of the split in Eq. ($\ref{split}$). From the gauge transformation $\widetilde A_{ij}$ we can infer the linear-order expression of $A_{ij}$ in Eq. ($\ref{asim}$) which indeed transforms accordingly considering that $\Om^i$ is gauge-invariant and $C_\al$ at linear order gauge transforms as the vector part of $\CC_{\al\be}$:
C_=C_-L_ .
Finally, directly substituting all the gauge transformations in our second-order expression for $S_{ij}$ in Eq. ($\ref{tetrade2}$), we recover Eq. ($\ref{Sij}$), proving the validity of our expression.

§ EXPLICIT EXPRESSIONS OF THE MOST RELEVANT SECOND-ORDER QUANTITIES

In Section $\ref{sec:second}$ we derived all the second-order expressions related to the light propagation from a luminous source to the observer position. In Table $\ref{tab1}$ we provide the references to such equations.

tableReferences to the second-order equations given in Section $\ref{sec:second}$.

$\de g$     Eq. ($\ref{deg2}$)      $\de z$      Eq. ($\ref{dz2}$)

$\GG^\eta$ , $\GG^\al$     Eq. ($\ref{GG2}$)     $\dhnu$      Eq. ($\ref{dhnu2}$)
$\de u$      Eq. ($\ref{du2}$)     $\de\eta_o$      Eq. ($\ref{lapseconf}$)
$\de t_o$      Eq. ($\ref{dt2met}$)     $\DT_s$     Eq. ($\ref{deta2}$)
$\de x^\al_o$      Eq. ($\ref{dx2met}$)     $\DX^\al_s$     Eq. ($\ref{Dx2}$)
$\dnu_s$      Eq. ($\ref{dnuint}$)     $\drr$ , $\dtt$ , $\dpp$      Eq. ($\ref{dangoli2}$)
$\dea^\al_s$      Eq. ($\ref{deaint}$)     $\dDD$      Eq. ($\ref{dDD2}$)
$\dnu_o$      Eq. ($\ref{dnu2o}$)    $\de V$      Eq. ($\ref{DV2}$)
$\dea^\al_o$      Eq. ($\ref{dea2o}$)    $\de A$      Eq. ($\ref{DA2}$)

Of all these quantities, the most relevant are those related to cosmological observables, i.e., the redshift fluctuation $\de z$, the distortion in the source position $\drr,\dtt,\dpp$, and the perturbation in the physical volume $\de V$ and area $\de A$.

We now present the second-order expressions of these quantities in terms of metric perturbations and observer velocity. For clarity we will keep some of those quantities in the table at the linear order (i.e., they will not be expressed in terms of the metric), but we provide the reference to their expressions.
Note that the second-order antisymmetric part $A^j{}_i$ of the space-like tetrad perturbation $\de e^\al_i$ never appears.

Distortion in the source position $(\drr, \dtt, \dpp)$ in Eq. ($\ref{dangoli2}$)

-1/4C^_,(iC^)_, ) - Å-_++3/2Å^2

-_( ^+^)+Å( _-2)-^_+1/2_^+(^-^)C_[,]

+ε_i j( ^-^)n^i^j]|_+n_x_o^μ(_μx^)|_+n_∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t

-1/_z[(- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd['-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^
+_η_o+^μ_s(_μ)|_x̅_z ]

+1/2^_-^_-1/a^t_,)|_x̅_t ]

+1/2_t_o^2+n_[_^μ _μ( ^+n^)|_x̅_]^z_
+n_∫_0^_zd[ -^(^_,+n^_,)

-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')+n^('-^_,+Å'-2+_,

+'+2(Å' - )
- 2Å(Å'-2+_,
-_(α,) -_α')

-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,
+_z/2 ( ^2+sin^2 ^2 ) ,
=   1/_z_{

-1/4C^_,(iC^)_, )]|_+x_o^μ(_μx^)|_+[_^μ _μ( ^+n^)|_x̅_]^z_

+∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t
+∫_0^_zd[ -^(^_,+n^_,)

-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')
-/_z+1/2cossin ^2 .
The explicit expressions of the linear-order terms are given in Eqs. ($\ref{dxlin}$), ($\ref{lapseconf}$), ($\ref{dk1s}$), ($\ref{dk1o}$), ($\ref{dhnu2}$), ($\ref{Deta1}$), ($\ref{Dx1}$), ($\ref{dz2}$), ($\ref{deltar}$), ($\ref{deltat}$), ($\ref{deltap}$). According to Eq. ($\ref{dangoli2}$) the expression for $\dpp$ is analogous to that of $\dtt$ replacing $\ttt_\al$ with $\pp_\al/\sin\ttt$ and $-\frac{\drr}{\rbar_z}\dtt +\frac{1}{2}\cos\ttt\sin\ttt~\dpp^2$ with $-\frac{\drr}{\rbar_z}\dpp-\rbar_z\cot\ttt~\dtt\dpp$.

Fluctuation in the observed redshift $\de z$ in Eq. ($\ref{dz2}$)

z   =   _(x^μ_o(_μt)|_-1/2_t_o^2+∫_0^t̅_odt(-Å+1/2Å^2
+1/2^_-^_-1/a^t_,)|_x̅_t )

+(- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd['-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^
+_η_o+^μ_s(_μ)|_x̅_z .
The explicit expressions of the linear-order terms are given in Eqs. ($\ref{dxlin}$), ($\ref{lapseconf}$), ($\ref{dk1s}$), ($\ref{dk1o}$), ($\ref{dhnu2}$), ($\ref{Deta1}$), ($\ref{Dx1}$).

Perturbation in the physical volume $\de V$ occupied by the source in Eq. ($\ref{DV2}$)
=    3[ _(x^μ_o(_μt)|_-1/2_t_o^2+∫_0^t̅_odt(-Å+1/2Å^2
+1/2^_-^_-1/a^t_,)|_x̅_t )

+(- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd['-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^
+_η_o+^μ_s(_μ)|_x̅_z ]+Å+^α_α-12Å^2+1/2^α_α+Å ^α_α
^αβ_αβ+H_z_z/2 z( ^2+sin^2 ^2 )+3/2 ( ^2+sin^2 ^2 )

-1/4C^_,(iC^)_, ) - Å-_++3/2Å^2

-_( ^+^)+Å( _-2)-^_+1/2_^+(^-^)C_[,]

+ε_i j( ^-^)n^i^j]|_+n_x_o^μ(_μx^)|_+n_∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t

-1/_z[(- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd('-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^
+_η_o+^μ_s(_μ)|_x̅_z ]

+1/2^_-^_-1/a^t_,)|_x̅_t ]

+1/2_t_o^2+n_[_^μ _μ( ^+n^)|_x̅_]^z_
+n_∫_0^_zd[ -^(^_,+n^_,)
-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')+n^('-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

-1/4C^_,(iC^)_, )]|_+x_o^μ(_μx^)|_+[_^μ _μ( ^+n^)|_x̅_]^z_

+∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t
+∫_0^_zd[ -^(^_,+n^_,)

-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')

-1/4C^_,(iC^)_, )]|_+x_o^μ(_μx^)|_+[_^μ _μ( ^+n^)|_x̅_]^z_

+∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t
+∫_0^_zd[ -^(^_,+n^_,)

-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')
+-+ (+)

+Å+3/2 Å^2 + 12
+_+3 +3 ^2
-Å(2 _z-2κ+H_zz)

+(+3 )
+^μ_μ( 3 +-Å+_) .
The explicit expressions of the linear-order terms are given in Eqs. ($\ref{dxlin}$), ($\ref{lapseconf}$), ($\ref{dk1s}$), ($\ref{dk1o}$), ($\ref{dhnu2}$), ($\ref{Deta1}$), ($\ref{Dx1}$), ($\ref{dz2}$), ($\ref{deltar}$), ($\ref{deltat}$), ($\ref{deltap}$), ($\ref{deg2}$).

Perturbation in the physical area $\de A$ occupied by the source in Eq. ($\ref{DA2}$)
A   =    2[ _(x^μ_o(_μt)|_-1/2_t_o^2+∫_0^t̅_odt(-Å+1/2Å^2
+1/2^_-^_-1/a^t_,)|_x̅_t )

+(- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd['-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^
+_η_o+^μ_s(_μ)|_x̅_z ]- Å+3/2 Å^2 + 12

+Å+^α_α-12Å^2+1/2^α_α+Å ^α_α+12^α_α^_-
^αβ_αβ+^2+sin^2 ^2

-1/4C^_,(iC^)_, ) - Å-_++3/2Å^2

-_( ^+^)+Å( _-2)-^_+1/2_^+(^-^)C_[,]

+ε_i j( ^-^)n^i^j]|_+n_x_o^μ(_μx^)|_+n_∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t

-1/_z[(- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd('-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^
+_η_o+^μ_s(_μ)|_x̅_z ]

+1/2^_-^_-1/a^t_,)|_x̅_t ]

+1/2_t_o^2+n_[_^μ _μ( ^+n^)|_x̅_]^z_
+n_∫_0^_zd[ -^(^_,+n^_,)
-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')+n^('-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

-1/4C^_,(iC^)_, )]|_+x_o^μ(_μx^)|_+[_^μ _μ( ^+n^)|_x̅_]^z_

+∫_0^t̅_odt 1/a(^+Å ^-^x^_,)|_x̅_t
+∫_0^_zd[ -^(^_,+n^_,)

-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')

-1/4C^_,(iC^)_, )_n^i-n_∫_0^_zd[-'^+^^_,+Å^,α-'^α
-  _^  ,α + ^α_  ,

- 2 '^α_+
2 ^α_, - _^    ,α
+(2Å^,α - 2'^α-_^  ,α
+ ^α_  ,- 2 '^α_ )
-^(_^  ,α
- ^α_  ,+ 2 '^α_-4 ^α_(,γ)+2_γ^    ,α
)+ ^α(Å'+ _, + '_- 2 Å_,)

- 2^α(  2 _,
- _,+Å_,
-'_ -2_[,γ]
- 2_γ')]|_x̅_

+_-[ (- Å-_++3/2Å^2-_( ^+^)+Å( _-2)-^_
+1/2_^+(^-^)C_[,]+ε_i j( ^-^)n^i^j)_+∫_0^_zd('-^_,

+Å'-2+_,+'+2(Å' - )
- 2Å(Å'-2+_,

-_(α,) -_α')
-^(Å_,α-'_α+2_[,] -2'_
+ 2 _α,

+Å+_-_+Å-1/2Å^2+Å_+_( ^-^)
+1/2 ^_+2_^]

-1/_z( _+_/sin )-2κ(2/_z-)-2/_z+^2/^2_z-1/2^2-1/2(sin )^2

+1/_z+sin (_+_)+1/_z+z^2+2 z
+( 2 z+)(  -Å+2/_z-2κ-)
+^μ_μ( 2 z+-Å+_+_-) .
The explicit expressions of the linear-order terms are given in Eqs. ($\ref{dxlin}$), ($\ref{lapseconf}$), ($\ref{dk1s}$), ($\ref{dk1o}$), ($\ref{dhnu2}$), ($\ref{Deta1}$), ($\ref{Dx1}$), ($\ref{dz2}$), ($\ref{deltar}$), ($\ref{deltat}$), ($\ref{deltap}$), ($\ref{deg2}$).


# Network Contention-Aware Cluster Scheduling with Reinforcement Learning

Junyeol Ryu  Computer Science and Engineering
Seoul National University
Seoul, Korea
<EMAIL_ADDRESS>Jeongyoon Eo  Computer Science and Engineering
Seoul National University
Seoul, Korea
<EMAIL_ADDRESS>
###### Abstract

With continuous advances in deep learning, distributed training is becoming
common in GPU clusters. Specifically, for emerging workloads with diverse
amounts, ratios, and patterns of communication, we observe that network
contention can significantly degrade training throughput. However, widely used
scheduling policies often face limitations as they are agnostic to network
contention between jobs. In this paper, we present a new approach to mitigate
network contention in GPU clusters using reinforcement learning. We formulate
GPU cluster scheduling as a reinforcement learning problem and opt to learn a
network contention-aware scheduling policy that efficiently captures
contention sensitivities and dynamically adapts scheduling decisions through
continuous evaluation and improvement. We show that compared to widely used
scheduling policies, our approach reduces average job completion time by up to
18.2% and effectively cuts the tail job completion time by up to 20.7% while
allowing a preferable trade-off between average job completion time and
resource utilization.

###### Index Terms:

Scheduling, Machine learning, Reinforcement learning, Heterogeneous (hybrid)
systems

##  I Introduction

Distributed deep learning (DL) training is becoming increasingly prevalent in
GPU clusters. A recent analysis published by Alibaba [1] has reported that
over 80% of the total submitted DL training jobs111Combination of DL model to
train, GPU demand, and scheduled nodes. run on multiple GPUs spanning over
multiple nodes. This rate corresponds to approximately 5$\times$ increase in
five years from the previous report from Microsoft [2]. Also, emerging DL
training workloads present diverse amounts, ratios, and patterns of
communication. For instance, Fully Sharded Data-Parallel training (FSDP) [3,
4] features heavy communication, with at least 50% increased communication
cost compared to conventional data-parallel training [5]. Mixture of Experts
(MoE) [6] training implements a gating network-based routing mechanism between
distributed experts using AllToAll pattern that has higher communication cost
than traditional AllReduce pattern. Eventually, previous studies have shown
that GPU clusters can suffer significant performance degradation due to
conflicts in network communication when distributed training and emerging
workloads with diverse communication are common [7, 1, 8]. Ultimately, this
trend raises a new challenge to GPU cluster scheduling: mitigating performance
slowdown due to network contention.

Notwithstanding the new challenge, we notice that widely used scheduling
policies (e.g., LAS [9, 10] and SRTF [11]) are often agnostic to network
contention between jobs, hence are prone to significant degradation in jobs’
training throughput. For example, when FSDP and MoE training share networks,
they suffer up to 49.1% and 66.7% throughput degradation compared to isolated
cluster, respectively. Unfortunately, such unfavorable scheduling decisions
cannot be avoided under contention-agnostic policies. Yet, we observe that a
job experiences varying degrees of network contention based on the model and
the placement of the co-located222Scheduling jobs such that they become each
other’s node-sharing neighbors, sharing all or part of the allocated nodes,
and thus bandwidth of intra and inter-node networks. jobs. Accordingly, the
throughput degradation of the aforementioned example can be improved up to
21.6% and 19.5%, respectively, depending on how they are co-located. Building
upon this insight, we define contention sensitivity ($CS$) of a job as the
ratio of its ideal throughput to the degraded throughput when co-located with
another job (1).

$CS=\frac{Throughput_{\text{ideal}}}{Throughput_{\text{contention}}}$ (1)

In this paper, we propose a method for scheduling distributed DL jobs in GPU
clusters that minimizes network contention. Two notable challenges are as
follows: 1. efficiently capturing contention sensitivities, and 2. dynamically
adapting to diverse distributions of jobs and their contention sensitivities.
To address these challenges, we devise a reinforcement learning (RL)-based
approach to swiftly learn an effective scheduling policy by continuous
evaluation and improvement of scheduling decisions across diverse
distributions of jobs. Specifically, we make the following contributions:

  * •

We propose a novel design that translates the network contention problem in
cluster scheduling into an RL problem. We show that our design can efficiently
capture contention sensitivities of jobs and dynamically adapt scheduling
decisions across diverse distributions of jobs.

  * •

We present an end-to-end system for training scheduling policies with RL to
its deployment on GPU clusters. We provide two initial scheduling policies,
RL-base and RL-Hybrid, and implement mechanisms that execute the decisions of
the scheduling policy.

  * •

We evaluate our scheduling policies with a variety of DL training job traces
on a GPU cluster. RL-base outperforms LAS and SRTF by reducing average JCT by
up to 18.2% and effectively cutting tail JCT by up to 20.7%, while RL-Hybrid
achieves a preferable trade-off between average JCT and resource utilization.

  * •

We open our work at (https://github.com/gajagajago/deepshare) as a community
asset for future research in RL-based GPU cluster scheduling.

##  II Background

DL training in GPU clusters. Deep learning is a process to train a deep neural
network (DNN) with the purpose of analyzing and deriving valuable knowledge
from data [12]. The DNN is constructed with numerous layers and parameters.
During training, the DNN is tasked with making predictions and updating its
parameters based on the computation of errors in relation to actual outcomes.
Given its inherent computational intensity, training is usually performed on
high-capacity accelerators like GPUs. Due to extremely high price (e.g.,
NVIDIA A100 costs about $10k), organizations often build GPU clusters to be
shared by users and production groups [13]. Consequently, GPU clusters usually
employ a scheduler for the efficient allocation of cluster resources. These
schedulers predominantly operate under two key objectives: reducing Job
Completion Time (JCT) [14, 9, 15, 16, 17, 7] and enhancing resource
utilization [18, 9, 19, 13]. Thus, these two performance metrics form the
foundation of our RL formulation’s reward (Section IV).

RL for cluster scheduling. Reinforcement learning involves an agent that
learns to make better decisions directly from experience interacting with the
environment [20]. The agent learns by reinforcement, wherein it receives
rewards contingent on the quality of its decisions [21]. RL has recently
become an active area of research in machine learning [22, 23, 24, 25, 26],
and RL-based approaches have demonstrated great potential in various domains
including congestion control [27], video streaming [28, 29, 30], real-time
communication [31, 32, 33, 34, 35], and resource management [36, 21, 37, 38].
RL approaches are known to be especially well-suited to resource management
systems due to the followings:

  * •

Decisions made by the systems are highly repetitive, leaving abound of
training data (e.g., scheduling decisions and corresponding effects) to the RL
algorithm.

  * •

Reward can reflect complex objectives (e.g., JCT and utilization) which are
difficult to model analytically.

  * •

Agent are highly adaptive to constantly shifting or even previously unseen
circumstances.

DeepRM [21] and DeepPlace [37, 38] are two notable works that apply RL in
resource management. DeepRM represents the system state as a combination of
cluster resources and resource demands of jobs. In this scheme, cluster
resources are equipped with available time slots for job allocation, while
jobs demand distinct quantities of time slots for different resource types.
The action step revolves around matching jobs’ time slots with those of
resources, and the reward is designed to mirror the objective of minimizing
average slowdown. DeepPlace builds upon DeepRM and introduces a more intricate
reward that incorporates a blend of purpose-oriented cues, including factors
like resource contention penalty and under-utilization penalty.

Nevertheless, there exists a limitation when attempting to employ these
methodologies in scheduling distributed DL jobs in GPU clusters. First, since
these approaches primarily focus on host resources (such as CPU and DRAM) and
address small containerized microservices as their target jobs, their state
representation is inadequate for capturing distributed GPU jobs. Second, their
action scope remains limited to within a single node, making it impractical to
schedule distributed jobs spanning multiple nodes. Finally, due to the absence
of consideration for network contention between jobs, their reward structure
fails to encompass the performance impact of network contention. Therefore,
our main objective is to introduce an efficient RL-based solution for GPU
cluster scheduling that efficiently handles network contention problem
(Section III).

##  III Motivation

Task |  Model (Abbreviation) | Dataset | Trainable |  |  |
---|---|---|---|---|---|---
Parameters |  Average Bandwidth |  |  |  |  |
Consumption (MB/s) |  $\frac{Comm}{Comp}$ | Communication Patterns |  |  |  |
Graph |  GraphSage (GNN) [39] | Reddit | 0.33M |  24.63 | 0.57 | AllReduce
Image |  MobileNetV3 (IMG) [40] | ImageNet | 2.04M |  211.25 | 2.43 | AllReduce
Recommendation |  DLRM (DLRM) [41] | Criteo | 333.32M |  170.28 | 13.36 | AllReduce
Language |  Transformer-XL (LM) [42] | Wikitext-103 | 202.44M |  854.82 | 1.87 | AllReduce
|  GPT-2 (FSDP) [43] | Wikitext-2 | 184.89M |  2672.40 | 7.32 |  ReduceScatter, AllGather
|  GPT-2 (MoE) [43] | Wikitext-2 | 268.92M |  929.48 | 13.79 | AllToAll

Table I: Models used in this work. $\frac{Comm}{Comp}$ denotes communication-
to-computation ratio. Average bandwidth consumption and $\frac{Comm}{Comp}$ is
profiled in our evaluation environment (Section VI).

In this section, we introduce emerging DL training workloads encompassing
diverse communication characteristics. Then, we analyze their contention
sensitivities to draw insights for effective network contention-aware cluster
scheduling. We use these workloads (and their variations) to constitute the
jobs traces for evaluation (Section VI).

Communication characteristics of emerging DL training workloads (section III).
FSDP [4, 3] and MoE [6] exhibit high network bandwidth consumption and
communication-to-computation ratio. Specifically, FSDP and MoE demonstrate
respectively 3.01$\times$ and 5.67$\times$ higher communication-to-computation
ratio and 12.65$\times$ and 4.40$\times$ larger average bandwidth consumption
compared to a traditional image model training such as MobileNetV3 [40]. These
characteristics stem from their communication pattern, designed to achieve GPU
memory efficiency at the cost of increased communication. Concretely, FSDP
adds per-layer AllGather at forward pass, and per-layer AllGather and
ReduceScatter at backward pass compared to conventional data-parallel training
[4, 3]. MoE adds gated network-based AllToAll communication between
distributed experts [6].

DLRM [41] and Transformer-XL [42] display either a high communication-to-
computation ratio or high network bandwidth consumption. Especially, DLRM
exhibits a 5.49$\times$ higher communication-to-computation ratio but 19.3%
lower bandwidth consumption compared to MobileNetV3. Collective communications
account for a significant fraction of time in training DLRM at scale [44].
This is because recommendation tasks such as DLRM spend around 80% of total
training time on host resources [1] as sparse computation of element-wise
operators dominates. On the other hand, Transformer-XL exhibits a 23.04% lower
communication-to-computation ratio but 4.04$\times$ higher bandwidth
consumption compared to MobileNetV3. This is because the Transformer [45]
blocks employ the attention mechanism that requires larger amounts of
computation compared to the convolution-based mechanism of image models.

GraphSage [39] exhibits the lowest network bandwidth consumption and
communication-to-computation ratio, with 76.85% less communication-to-
computation ratio and 88.34% lower average bandwidth consumption compared to
MobileNetV3. This is because GraphSage training partitions the input graph
among nodes, where per-node graph preprocessing takes 30-90% of the training
time with involving only a little communication [1].

Figure 1: Contention sensitivity heat map. Darker value indicates higher
contention sensitivity, which signifies larger throughput degradation of the
target job according to the network contention from the co-located job.
Numbers behind the model name denote nodes and GPUs per node, respectively.

Contention sensitivity of a job varies by its co-located job. In Figure 1,
each pixel value depicts the contention sensitivity of the target job
according to the network contention from the co-located job. A pair of a
target model and a co-located model is represented as a $6\times 6$ grid (e.g.
the black boxes in Figure 1 shows the varying contention sensitivities of FSDP
and MoE, with respect to diverse GPU demands and node assignments). We observe
that some jobs (GNN and IMG) show consistent contention sensitivities
regardless of co-located jobs, whereas the others (DLRM, LM, FSDP, and MoE)
show high variability with regard to the model, GPU demands, and node
assignment of the co-located job. For example, when FSDP and MoE training are
co-located, they experience contention sensitivities of up to 1.96 and 3.00,
respectively, with varying degrees according to their co-location (the black
boxes in Figure 1). In contrast, when FSDP and MobileNetV3 training are co-
located, they exhibit moderate degrees of contention sensitivities with at
most 1.35 and 1.43, respectively (the white boxes in Figure 1).

We summarize our findings from this section as follows:

  * •

Jobs exhibit a variety of communication characteristics, which contribute to
their varying contention sensitivities when co-located with other jobs.

  * •

In this regard, scheduling based on efficiently captured contention
sensitivities and aimed at reducing expected contention will effectively
alleviate network contention.

##  IV RL Formulation

Requirements. Summarizing the key insights from previous sections, the main
requirements of an ideal network contention-aware GPU cluster scheduling
includes:

  * •

R1. Fast adaptation of its decisions to reflect the constantly changing
distribution of contention sensitivities as scheduled jobs change.

  * •

R2. Minimizing cluster-wide performance degradation due to contention,
achieving low average and tail JCT while maintaining high resource
utilization.

To satisfy the requirements, we propose our design for a network contention-
aware scheduling with RL. We explain our formulation of state, action, reward,
and training algorithm. We also present two initial scheduling policies
(agents in RL term) trained with RL, namely RL-base and RL-Hybrid. These
policies serve as benchmarks for evaluation (Section VI).

Figure 2: State representation. (Example for 4 nodes each with 8 GPUs, 3
scheduled jobs and 3 candidate jobs with varying in GPU demands)

State. To adapt the agent’s decisions such that it reflects the constantly
changing contention sensitivities (R1), we develop a state representation that
captures the cluster-wide co-location of jobs. Figure 2 illustrates an example
of our state design. We represent the cluster state as a two-dimensional
tensor of shape ¡Nodes, $2\times$GPUs per node¿. This fixed state design
satisfies a desirable attribute to be applied as input to neural network-based
RL training algorithms. The left half encodes the physical placement of
scheduled jobs. Each slot contains the profile of a job, if any job is
scheduled on the GPU corresponding to the slot. Conversely, the right half
captures the resource demands of candidate jobs in the waiting queue. For
example, if a candidate job is encoded in $slot_{i,j}$ of the right half, it
signifies that one feasible placement option for the job is distributing its
total $j*2^{i}$ demands across $2^{i}$ nodes with $j$ GPUs each. Since there
are various ways to partition a job’s demand, jobs with multiple GPU demands
can be encoded in multiple slots of the right half.

1

2

3

Input : $Trace$: job trace, $Alg$: RL algorithm

Parameter : $W_{1},W_{2}$: weights, $T$: scheduling interval, $F$: policy
checkpoint path

4

5$C\leftarrow$ Init empty cluster environment

6 $Q\leftarrow$ Waiting jobs in $Trace$

7 $Policy\leftarrow$ Init policy with $Alg$

8 while  _! $C$.empty or !$Q$.empty_ do

9 Train(_$C$ , $Q$, $Policy$_)

10

11 end while

12 Save $Policy$ to $F$

13

14 Function  _Round(_$C$ , $Q$, $Policy$_)_:

15 $State\leftarrow$ Get state from $C$ and $Q$

16 $Actions\leftarrow$ $Policy$ computes action from $State$

17 for  _$job,nodes\in Actions$_ do

18 Schedule $job$ to $nodes$

19

20 end for

21 Sleep for $T$

22 $Reward\leftarrow$ ComputeReward(_$C$_)

23 Update $Policy$ with $Reward$

24

25

26

27 Function  _ComputeReward(_$C$_)_:

28 $CS$, $Util\leftarrow$ 0

29 for  _$job\in C.jobs$_ do

30 $CS_{\text{job}}\leftarrow$ Profile contention sensitivity of $job$

31 Update $CS$ with $CS_{\text{job}}$

32

33 end for

34 for  _$node\in C.nodes$_ do

35 for  _$GPU\in node$_ do

36 $Util_{\text{GPU}}\leftarrow$ Profile utilization of $GPU$

37 Update $Util$ with $Util_{\text{GPU}}$

38

39 end for

40

41 end for

42 return $-W_{1}*CS+W_{2}*Util$

43

44

Algorithm 1 RL Training Algorithm

Action. Choosing an arbitrary set of candidate jobs from the waiting queue can
lead to a large action space of up to $2^{Q}$ when the queue length is $Q$. To
reduce the size of the action space for faster convergence of the training
algorithm, $K$ candidate jobs with different GPU demands are picked starting
from the head of the waiting queue (forming the right half of the state
representation). Besides, to keep an action space of constant size without
invalid actions, our action space for the candidate jobs is a set of lists of
nodes to schedule the targets. $\emptyset$ denotes that the agent decided not
to schedule any jobs in the current round. For example, if the agent returns
[Node1, Node2] as an action for job type 1 (Figure 2), it is assigning the job
to Node1 and Node2 with two GPUs each as job type 1 is encoded in $slot_{2,2}$
on the right half of the state representation. One possibly beneficial
extension of the action space can be to migrate a job if there is a better
placement option or preempt if the current co-location is detrimental. But for
the current prototype, we only support preemption based on a predefined
threshold on contention sensitivity and plan to incorporate this feature in
the action space soon.

Reward. To minimize cluster-wide performance degradation due to contention
(R2), our reward penalizes an increase in cluster-wide average contention
sensitivity of scheduled jobs and provides incentives on higher cluster-wide
average GPU utilization (2). The two terms articulate the reward because
reduced average job contention sensitivity is positively related to reduced
JCT, and high GPU utilization corresponds to enhanced resource utilization,
which are the two pivotal performance metrics of GPU cluster schedulers
(Section II). Weights ${W_{1}}$ and $W_{2}=1-W_{1}$ are imposed on contention
sensitivity (CS) and utilization (Util) term. Its values can be tailored to
reflect the relative preference between average contention sensitivity and
utilization, as the two terms typically have a trade-off relationship. For
instance, co-locating as many jobs as possible to maximize resource
utilization may exacerbate average contention sensitivity, leading to more
network contention experienced by the co-located jobs.

$Reward=-W_{1}*CS+W_{2}*Util$ (2)

Training algorithm. Algorithm 1 illustrates the training algorithm. The agent
(Policy) is trained with a neural network-based RL algorithm (Alg) using job
traces (Trace) on a simulated GPU cluster environment. The agent makes
scheduling decisions until all jobs have finished (Lines 1-1). For every
scheduling round, the agent is given the state representation (Line 1) and
chooses an action (Line 1). Jobs are scheduled as depicted in the action
(Lines 1-1). After waiting for the scheduling interval (Line 1), reward is
computed with respect to contention sensitivities of the scheduled jobs and
resource utilization of the cluster (Line 1) and is returned to the agent to
adapt its policy using the neural network-based RL algorithm (Alg) (Line 1).
By iterating through numerous scheduling rounds in the training process, the
agent continuously explores diverse job co-location options and adapts its
scheduling decisions to effectively mitigate cluster-wide network contention.
When the training finishes, the agent is saved into a file as a trained
scheduling policy (Line 1).

Benchmark policies. We present two initial scheduling policies trained with
RL, namely RL-base and RL-Hybrid. These policies serve as benchmarks for
evaluation (Section VI).

  * •

RL-base makes scheduling decisions only based on the action of the trained
policy. It may decide not to schedule jobs even when the cluster has enough
resources to avoid an increase in performance degradation due to contention.

  * •

RL-Hybrid performs decision-level multiplexing of the trained policy and a
simple rule-based policy. 333Such hybrid design is in line with the recent
trend of incorporating the decisions of rule-based policies to guide or aid
the RL-based policy to achieve both the RL-based one’s high adaptivity and the
rule-based one’s stability and interpretability [33, 34, 35, 46]. It usually
follows the decisions of the trained policy, but when it faces an $\emptyset$,
i.e. decision not to schedule, it follows a safety rule of trying greedy
scheduling to prevent low utilization.

##  V System

Figure 3: System overview. The policy is trained on a trace-driven simulated
GPU cluster environment using RL and deployed on the actual GPU cluster.

###  V-A Framework Overview

We build an end-to-end system from training a scheduling policy with RL on a
simulated cluster environment to its deployment on GPU clusters. Figure 3
illustrates the overall system architecture. The system is composed of two
parts: Trainer and Scheduler. Trainer implements a simulated cluster
environment and the training algorithm to train a scheduling policy with RL
(Section IV). Scheduler deploys the trained policy on GPU clusters and
provides mechanisms to execute the scheduling decisions.

###  V-B Trainer

Trainer provides OpenAI gym [47]-based simulated cluster environment of nodes
and GPUs. It also provides a common interface for a set of off-the-shelf RL
algorithms in widely used stable-baselines3 [48]. Hence, researchers can
easily customize the RL formulation and train their own policies. Trainer
requires only job trace (e.g., Microsoft’s Philly trace [2] and Alibaba’s PAI
trace [1]) to start training on simulated cluster environment. Training an
agent on an episode of 256 jobs completes in less than a minute.

###  V-C Scheduler

Scheduler is implemented on top of Slurm [49], an open-sourced Linux cluster
manager, and uses HDFS [50] as the checkpoint store. Controller is implemented
on top of Slurm control daemon (slurmctld) by adding global data structures
for cluster state and waiting queues, an interface for communicating with the
scheduler and the node agents (ComputeObs() and Schedule()), and
PreemptionManager to support checkpointing through asynchronous preemption
protocol. Scheduler provides an interface to load the checkpoint of the
trained policy. It also supports a set of widely used scheduling policies,
where users can configure one via script. NodeAgent is implemented on top of
Slurm daemon (slurmd) with CheckpointManager, Profiler built on top of PyTorch
profiler [51], and ProcessManager for initializing a node process group used
for synchronization in distributed DL training.

Figure 4 illustrates the in-depth workflow of Scheduler. The trained policy is
plugged into Core and deployed on the GPU cluster. Core executes round-based
scheduling with a predefined round duration knob. For every round,
ComputeObs() encodes the cluster and the waiting queue state as a state and
sends it to Core( 1). Following the scheduling decision produced by the policy
( 2), Controller dequeues the target jobs to schedule from the waiting queue (
3). Then Schedule() spawns NodeAgent for the job (one per each node allocated
to the job), and submits the job training script ( 4). NodeAgent builds the
process group for the job and delivers the profiled job stats to Controller
after executing an online profiling for predefined iterations ( 5). Monitoring
the contention sensitivity of jobs from the profiles, Controller sends a
preemption signal to node agents with jobs with higher contention sensitivity
than a predefined knob. NodeAgent waits for CheckpointManager to save the
checkpoint of the job to distributed file system (DFS)-based job checkpoint
store and returns its status to PreemptionManager asynchronously ( 6). To
support re-scheduling, PreemptionManager maintains a preemption lookup table.
Thus, whether a job has been preempted can be identified by querying the job
identifier on the lookup table, and the scheduler loads from the checkpoint
store to resume from the latest checkpoint ( 7).

Figure 4: Detailed workflow of Scheduler.

##  VI Evaluation

We evaluate the performance of our approach using the two benchmark policies,
RL-base and RL-Hybrid (Section IV).

Environment. Experiment conducts on a homogeneous cluster of four nodes, each
equipped with eight NVIDIA TITAN XP GPUs connected over 16GBps PCIe Gen3. Each
node has one Mellanox MT27700 family ConnectX-4 NIC and is interconnected with
10Gbps Ethernet and 40Gbps InfiniBand.

We believe our evaluation environment of 32 GPUs covers a large section of
important use cases: Microsoft reports that 93.7% of all jobs submitted to
their internal DGX-2 clusters required at most 32 GPUs in the second half of
2021 [52]. Also, despite the moderate scale of the evaluation environment, we
note that our design can be applied to large-scale clusters without any
limitation.

Traces. To evaluate different scenarios with varying degrees of communication
intensities, we generate four traces with different communication intensities
according to the categorization in Section III. 444Normal, heavy, medium, and
low communication intensive traces follow the model distribution of
GNN:IMG:DLRM:LM:FSDP:MoE=1:1:1:1:1:1, 1:1:1:1:4:4, 1:1:4:4:1:1, and
4:4:1:1:1:1, respectively. Each trace contains ten job sets each with 256
randomly shuffled jobs. Each job is randomly assigned a GPU demand of up to
32, and given a total number of samples to train, which is configured to set
its expected training time in an isolated cluster as one hour. The jobs’
configurations (e.g., batch size, parameter size, number of layers, and
gradient accumulation steps) are initialized randomly to evaluate
extensibility to unseen cases.

Branches of RL-base and RL-Hybrid. We use five branches of RL-base and RL-
Hybrid each by sweeping the range of the weights in the reward ($W_{1}$,
$W_{2}$) (Section IV) to demonstrate trade-off relationship between average
contention sensitivity and utilization. 555We train each branch for 10 sets of
jobs, each with 20 episodes that each contains 256 jobs. In Figure 5 to Figure
8, A, B, C, D, and E correspond the branch with (0.3, 0.7), (0.4, 0.6), (0.5,
0.5), (0.6, 0.4), and (0.7, 0.3) of the weights, respectively.

Figure 5: CDF of JCT. (Trace: Normal, Branch: B) Figure 6: Proportion of
average contention sensitivity. (Trace: Normal, Branch: B) Figure 7: KDE of
utilization. (Trace: Normal, Branch: B)

Results. Figure 5 shows the cumulative distribution function (CDF) of JCT.
Comparing average JCT to SRTF and LAS, RL-base (RL-Hybrid) achieves 15.4%
(12.1%) and 18.2% (15.1%) reduction, respectively. The large improvement in
average JCT is the result of reduced cluster-wide contention sensitivity.
Figure 6 shows the proportion of average contention sensitivity (i.e., average
of contention sensitivities that jobs experience during training). RL-base and
RL-Hybrid show a concentration of average contention sensitivity at a lower
degree than SRTF and LAS. Clearly, this demonstrates that scheduling decisions
of RL-base and RL-Hybrid successfully reduce cluster-wide average contention
sensitivity which leads to an eventual large reduction in cluster-wide average
JCT. For p90 JCT, RL-base and RL-Hybrid equally achieve 16.4% and 20.7%
reduction compared to SRTF and LAS, respectively. Even though RL-Hybrid shows
the smaller degree of average JCT reduction ($\approx 3\%$) than RL-base, RL-
Hybrid trades such cost to a large improvement in cluster-wide utilization.
Figure 7 shows the kernel density estimation (KDE) of utilization. RL-Hybrid
displays the highest utilization, emphasizing the preferable trade-off between
average JCT and utilization. The difference mainly stems from their policy
discrepancies; when the trained policy decides $\emptyset$, i.e., not to
schedule, RL-Hybrid tries to improve utilization by multiplexing to greedy
scheduling.

The reduced tail JCT demonstrates that scheduling with RL effectively curbs
scheduling decisions that trigger the worst contention situations.

(a) Normal communication traces

(b) Heavy communication traces

(c) Medium communication traces

(d) Low communication traces

Figure 8: Performance comparison for varying traces. The light red and red
colored dots are different branches of RL-base and RL-Hybrid, respectively.
Utilization denotes the ratio of used GPUs to total GPUs in the cluster.

Figure 8 shows the results according to traces with varying degrees of
communication intensities. In normal, high, and medium communication traces,
we observe a trade-off relationship between average JCT and utilization among
RL-base, SRTF, and LAS (blue dashed lines). RL-base’s A and B trained with
average contention sensitivity term weight ($w_{1}$) 0.3 and 0.4 achieve
relatively balanced average JCT and utilization. On the other hand, RL-base’s
C and D fail to reach the preferable trade-off. RL-base’s E trained with the
highest penalization of the average contention sensitivity achieves the best
average JCT at the cost of much lower utilization compared with SRTF and LAS
which opt for maximizing utilization while not taking contention into account.
RL-Hybrid moves the original dots of RL-base to the higher right end of the
graph. This means that when cluster job distribution follows normal, high, or
medium communication intensity conditions, adhering to the action of the
trained RL policy in common cases while conforming to the decisions produced
by the rule-based safety condition to prevent low utilization in rare cases
leads to a large improvement in utilization at the expense of a relatively
small increase in average JCT. Conversely, in low communication traces, we
observe that SRTF and LAS perform similarly to RL-Hybrid because the
relatively low level of contention sensitivities of jobs makes the impact of
network contention negligible.

##  VII Related Work

GPU cluster schedulers for DL. GPU clusters employ a scheduler for efficient
allocation of cluster resources. The common goals of the schedulers include
reducing average JCT [14, 9, 15, 16, 17, 7], improving utilization [18, 9, 19,
13], or providing fairness [53, 54]. Optimus [14] estimates training speed and
assigns more resources to jobs with higher marginal gain. Gandiva [18] employs
job time-slicing to provide early feedback and executes a greedy approach for
job migration and packing. Tiresias [9] introduces least attained service
metric to prevent starvation of late arriving jobs. Themis [53] introduces
finish time fairness and promotes fairness between jobs. Gandivafair
guarantees fairness between users through allocating proportional resource
shares. Antman [19] segments jobs into prioritized and opportunistic jobs to
apply different scheduling policies. Gavel [15] proposes a policy-agnostic
scheduling mechanism and schedules based on throughput metric that is
comparable across heterogeneous accelerator types. AFS [16] opts for balance
between the shortest job first policy and resource efficiency when allocating
idle resources to jobs. Pollux [17] co-optimizes system throughput and job-
level performance metrics. Synergy [13] allocates host resources (e.g., CPU
and DRAM) disproportionately according to jobs’ sensitivity to resources. Muri
[7] enables multiple jobs to be packed on the same set of resources by
interleaving different stages of different jobs. However, our work presents
the first approach that identifies the network contention sensitivity of jobs
and optimizes the cluster objectives by efficiently mitigating network
contention.

RL for systems. Using RL for cluster scheduling allows rapid adaptation by
learning to schedule DL training jobs that experience different levels of
network contention sensitivities, where the distribution of jobs with
different model, GPU demand, and placement combinations changes constantly.
The potential of the RL-based approach in high adaptivity to constantly
changing or even unseen conditions has been already shown in various domains.
In TCP congestion control (CC), PCC-RL [27] shows that RL-based TCP CC
algorithm can successfully learn to distinguish different types of network
losses that hand-optimized TCP CUBIC [55] cannot capture. In adaptive video
streaming, Pensieve [28] demonstrates that RL-based adaptive bitrate selection
for video streaming can improve quality-of-experience (QoE) such as reduced
frame stalling while improving bandwidth utilization. Merina [30] shows that
using meta-RL allows fast adaptation in unseen throughput dynamics, further
improving QoE across a wide range of network throughput patterns. Our work
applies RL in resource management system, and present the first approach to
propose network contention-aware GPU cluster sheduling with RL.

##  VIII Conclusion

We present a novel design that translates the network contention problem in
GPU cluster scheduling into an RL problem. Our RL formulation trains the
scheduling policies to efficiently tackle the challenge of continuously
changing contention sensitivities of jobs in GPU clusters. We build an end-to-
end system that can train scheduling policies with RL and deploy on GPU
clusters. Our evaluation show that RL-based scheduling policies achieve
reduced average and tail JCT by up to 18.2% and 20.7% compared to the widely
used LAS and SRTF scheduling policies, and allows preferable trade-off of
large improvement of utilization with small cost in average JCT. Our work is
open-sourced at https://github.com/gajagajago/deepshare for future research in
RL-based GPU cluster scheduling.

## References

  * [1] Q. Weng, W. Xiao, Y. Yu, W. Wang, C. Wang, J. He, Y. Li, L. Zhang, W. Lin, and Y. Ding, “MLaaS in the Wild: Workload Analysis and Scheduling in Large-Scale Heterogeneous GPU Clusters,” in _19th USENIX Symposium on Networked Systems Design and Implementation (NSDI 22)_ , 2022, pp. 945–960.
  * [2] M. Jeon, S. Venkataraman, A. Phanishayee, J. Qian, W. Xiao, and F. Yang, “Analysis of Large-Scale Multi-Tenant GPU Clusters for DNN Training Workloads,” in _2019 USENIX Annual Technical Conference (USENIX ATC 19)_ , 2019, pp. 947–960.
  * [3] M. Ott, S. Shleifer, M. Xu, P. Goyal, Q. D. Duval, and V. Caggiano. (2021) Fully Sharded Data Parallel: faster AI training with fewer GPUs. [Online]. Available: https://engineering.fb.com/2021/07/15/open-source/fsdp/
  * [4] Y. Zhao, R. Varma, C.-C. Huang, S. Li, X. Min, and A. D. Desmaison. (2022) Introducing PyTorch Fully Sharded Data Parallel (FSDP) API. [Online]. Available: https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/
  * [5] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, “Zero: Memory optimizations toward training trillion parameter models,” in _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 2020, pp. 1–16.
  * [6] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean, “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer,” _arXiv preprint arXiv:1701.06538_ , 2017.
  * [7] Y. Zhao, Y. Liu, Y. Peng, Y. Zhu, X. Liu, and X. Jin, “Multi-resource interleaving for deep learning training,” in _Proceedings of the ACM SIGCOMM 2022 Conference_ , 2022, pp. 428–440.
  * [8] S. Rajasekaran, M. Ghobadi, G. Kumar, and A. Akella, “Congestion control in machine learning clusters,” in _Proceedings of the 21st ACM Workshop on Hot Topics in Networks_ , 2022, pp. 235–242.
  * [9] J. Gu, M. Chowdhury, K. G. Shin, Y. Zhu, M. Jeon, J. Qian, H. Liu, and C. Guo, “Tiresias: A GPU cluster manager for distributed deep learning,” in _16th USENIX Symposium on Networked Systems Design and Implementation (NSDI 19)_ , 2019, pp. 485–500.
  * [10] M. Nuyens and A. Wierman, “The foreground–background queue: a survey,” _Performance evaluation_ , vol. 65, no. 3-4, pp. 286–307, 2008.
  * [11] M. E. Crovella, R. Frangioso, and M. Harchol-Balter, “Connection scheduling in web servers,” Citeseer, Tech. Rep., 1999.
  * [12] P. P. Shinde and S. Shah, “A review of machine learning and deep learning applications,” in _2018 Fourth international conference on computing communication control and automation (ICCUBEA)_. IEEE, 2018, pp. 1–6.
  * [13] J. Mohan, A. Phanishayee, J. Kulkarni, and V. Chidambaram, “Looking beyond GPUs for DNN scheduling on Multi-Tenant clusters,” in _16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)_ , 2022, pp. 579–596.
  * [14] Y. Peng, Y. Bao, Y. Chen, C. Wu, and C. Guo, “Optimus: an efficient dynamic resource scheduler for deep learning clusters,” in _Proceedings of the Thirteenth EuroSys Conference_ , 2018, pp. 1–14.
  * [15] D. Narayanan, K. Santhanam, F. Kazhamiaka, A. Phanishayee, and M. Zaharia, “Heterogeneity-Aware cluster scheduling policies for deep learning workloads,” in _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_ , 2020, pp. 481–498.
  * [16] C. Hwang, T. Kim, S. Kim, J. Shin, and K. Park, “Elastic resource sharing for distributed deep learning,” in _18th USENIX Symposium on Networked Systems Design and Implementation (NSDI 21)_ , 2021, pp. 721–739.
  * [17] A. Qiao, S. K. Choe, S. J. Subramanya, W. Neiswanger, Q. Ho, H. Zhang, G. R. Ganger, and E. P. Xing, “Pollux: Co-adaptive cluster scheduling for goodput-optimized deep learning,” in _15th USENIX Symposium on Operating Systems Design and Implementation (OSDI 21)_ , 2021.
  * [18] W. Xiao, R. Bhardwaj, R. Ramjee, M. Sivathanu, N. Kwatra, Z. Han, P. Patel, X. Peng, H. Zhao, Q. Zhang _et al._ , “Gandiva: Introspective cluster scheduling for deep learning,” in _13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)_ , 2018, pp. 595–610.
  * [19] W. Xiao, S. Ren, Y. Li, Y. Zhang, P. Hou, Z. Li, Y. Feng, W. Lin, and Y. Jia, “AntMan: Dynamic Scaling on GPU Clusters for Deep Learning,” in _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_ , 2020, pp. 533–548.
  * [20] R. S. Sutton and A. G. Barto, _Reinforcement learning: An introduction_. MIT press, 2018.
  * [21] H. Mao, M. Alizadeh, I. Menache, and S. Kandula, “Resource management with deep reinforcement learning,” in _Proceedings of the 15th ACM workshop on hot topics in networks_ , 2016, pp. 50–56.
  * [22] V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu, “Asynchronous methods for deep reinforcement learning,” in _International conference on machine learning_. PMLR, 2016, pp. 1928–1937.
  * [23] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller, “Playing atari with deep reinforcement learning,” _arXiv preprint arXiv:1312.5602_ , 2013.
  * [24] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski _et al._ , “Human-level control through deep reinforcement learning,” _nature_ , vol. 518, no. 7540, pp. 529–533, 2015.
  * [25] J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz, “Trust region policy optimization,” in _International conference on machine learning_. PMLR, 2015, pp. 1889–1897.
  * [26] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot _et al._ , “Mastering the game of Go with deep neural networks and tree search,” _nature_ , vol. 529, no. 7587, pp. 484–489, 2016.
  * [27] N. Jay, N. Rotman, B. Godfrey, M. Schapira, and A. Tamar, “A deep reinforcement learning perspective on internet congestion control,” in _International Conference on Machine Learning_. PMLR, 2019, pp. 3050–3059.
  * [28] H. Mao, R. Netravali, and M. Alizadeh, “Neural adaptive video streaming with pensieve,” in _Proceedings of the conference of the ACM special interest group on data communication_ , 2017, pp. 197–210.
  * [29] T. Huang, C. Zhou, X. Yao, R.-X. Zhang, C. Wu, B. Yu, and L. Sun, “Quality-aware neural adaptive video streaming with lifelong imitation learning,” _IEEE Journal on Selected Areas in Communications_ , vol. 38, no. 10, pp. 2324–2342, 2020.
  * [30] N. Kan, Y. Jiang, C. Li, W. Dai, J. Zou, and H. Xiong, “Improving Generalization for Neural Adaptive Video Streaming via Meta Reinforcement Learning,” in _Proceedings of the 30th ACM International Conference on Multimedia_ , 2022, pp. 3006–3016.
  * [31] J. Fang, M. Ellis, B. Li, S. Liu, Y. Hosseinkashi, M. Revow, A. Sadovnikov, Z. Liu, P. Cheng, S. Ashok _et al._ , “Reinforcement learning for bandwidth estimation and congestion control in real-time communications,” _arXiv preprint arXiv:1912.02222_ , 2019.
  * [32] A. Zhou, H. Zhang, G. Su, L. Wu, R. Ma, Z. Meng, X. Zhang, X. Xie, H. Ma, and X. Chen, “Learning to coordinate video codec with transport protocol for mobile video telephony,” in _The 25th Annual International Conference on Mobile Computing and Networking_ , 2019, pp. 1–16.
  * [33] H. Zhang, A. Zhou, J. Lu, R. Ma, Y. Hu, C. Li, X. Zhang, H. Ma, and X. Chen, “OnRL: improving mobile video telephony via online reinforcement learning,” in _Proceedings of the 26th Annual International Conference on Mobile Computing and Networking_ , 2020.
  * [34] H. Zhang, A. Zhou, Y. Hu, C. Li, G. Wang, X. Zhang, H. Ma, L. Wu, A. Chen, and C. Wu, “Loki: improving long tail performance of learning-based real-time video adaptation by fusing rule-based models,” in _Proceedings of the 26th Annual International Conference on Mobile Computing and Networking_ , 2021\.
  * [35] B. Wang, Y. Zhang, S. Qian, Z. Pan, and Y. Xie, “A hybrid receiver-side congestion control scheme for web real-time communication,” in _Proceedings of the 12th ACM Multimedia Systems Conference_ , 2021, pp. 332–338.
  * [36] G. Singh, R. Nadig, J. Park, R. Bera, N. Hajinazar, D. Novo, J. Gómez-Luna, S. Stuijk, H. Corporaal, and O. Mutlu, “Sibyl: Adaptive and extensible data placement in hybrid storage systems using online reinforcement learning,” in _Proceedings of the 49th Annual International Symposium on Computer Architecture_ , 2022, pp. 320–336.
  * [37] S. Mitra, S. S. Mondal, N. Sheoran, N. Dhake, R. Nehra, and R. Simha, “Deepplace: Learning to place applications in multi-tenant clusters,” in _Proceedings of the 10th ACM SIGOPS Asia-Pacific Workshop on Systems_ , 2019, pp. 61–68.
  * [38] S. S. Mondal, N. Sheoran, and S. Mitra, “Scheduling of time-varying workloads using reinforcement learning,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 35, no. 10, 2021, pp. 9000–9008.
  * [39] W. Hamilton, Z. Ying, and J. Leskovec, “Inductive representation learning on large graphs,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [40] A. Howard, M. Sandler, G. Chu, L.-C. Chen, B. Chen, M. Tan, W. Wang, Y. Zhu, R. Pang, V. Vasudevan _et al._ , “Searching for mobilenetv3,” in _Proceedings of the IEEE/CVF international conference on computer vision_ , 2019, pp. 1314–1324.
  * [41] M. Naumov, D. Mudigere, H.-J. M. Shi, J. Huang, N. Sundaraman, J. Park, X. Wang, U. Gupta, C.-J. Wu, A. G. Azzolini _et al._ , “Deep learning recommendation model for personalization and recommendation systems,” _arXiv preprint arXiv:1906.00091_ , 2019.
  * [42] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov, “Transformer-xl: Attentive language models beyond a fixed-length context,” _arXiv preprint arXiv:1901.02860_ , 2019.
  * [43] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever _et al._ , “Language models are unsupervised multitask learners,” _OpenAI blog_ , vol. 1, no. 8, p. 9, 2019.
  * [44] J. A. Yang, J. Park, S. Sridharan, and P. T. P. Tang, “Training deep learning recommendation model with quantized collective communications,” in _Conference on Knowledge Discovery and Data Mining (KDD)_ , 2020.
  * [45] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [46] S. Abbasloo, C.-Y. Yen, and H. J. Chao, “Classic meets modern: A pragmatic learning-based congestion control for the internet,” in _Proceedings of the Annual conference of the ACM Special Interest Group on Data Communication on the applications, technologies, architectures, and protocols for computer communication_ , 2020, pp. 632–647.
  * [47] G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba, “Openai gym,” _arXiv preprint arXiv:1606.01540_ , 2016.
  * [48] A. Raffin, A. Hill, A. Gleave, A. Kanervisto, M. Ernestus, and N. Dormann, “Stable-baselines3: Reliable reinforcement learning implementations,” _The Journal of Machine Learning Research_ , vol. 22, no. 1, pp. 12 348–12 355, 2021.
  * [49] A. B. Yoo, M. A. Jette, and M. Grondona, “Slurm: Simple linux utility for resource management,” in _Job Scheduling Strategies for Parallel Processing: 9th International Workshop, JSSPP 2003, Seattle, WA, USA, June 24, 2003. Revised Paper 9_. Springer, 2003, pp. 44–60.
  * [50] K. Shvachko, H. Kuang, S. Radia, and R. Chansler, “The hadoop distributed file system,” in _2010 IEEE 26th symposium on mass storage systems and technologies (MSST)_. Ieee, 2010, pp. 1–10.
  * [51] “PyTorch Profiler,” 2023, https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html.
  * [52] A. Shah, V. Chidambaram, M. Cowan, S. Maleki, M. Musuvathi, T. Mytkowicz, J. Nelson, O. Saarikivi, and R. Singh, “TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches,” in _20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23)_ , 2023, pp. 593–612.
  * [53] K. Mahajan, A. Balasubramanian, A. Singhvi, S. Venkataraman, A. Akella, A. Phanishayee, and S. Chawla, “Themis: Fair and efficient GPU cluster scheduling,” in _17th USENIX Symposium on Networked Systems Design and Implementation (NSDI 20)_ , 2020, pp. 289–304.
  * [54] S. Chaudhary, R. Ramjee, M. Sivathanu, N. Kwatra, and S. Viswanatha, “Balancing efficiency and fairness in heterogeneous GPU clusters for deep learning,” in _Proceedings of the Fifteenth European Conference on Computer Systems_ , 2020, pp. 1–16.
  * [55] S. Ha, I. Rhee, and L. Xu, “CUBIC: a new TCP-friendly high-speed TCP variant,” _ACM SIGOPS operating systems review_ , vol. 42, no. 5, pp. 64–74, 2008.