summaries: This paper presents a method for 3D scene reconstruction from a single image using implicit surface representations such as occupancy or SDF. The authors propose to incorporate loss functions on the spatial gradients to provide dense supervision in the 3D space in the case where 3D labels may be incomplete (e.g. open 3D meshes) or not well-defined everywhere. Experiments are performed on ShapeNet and ScanNet show that the proposed method can achieve competitive performance on single-image scene reconstruction tasks.

summaries: This paper describes novel loss functions for learning to predict an implicit 3D scene representation from a single image.  They argue that when working with real scan data of scenes (rather than single objects) it is difficult to generate accurate occupancy or signed distance function (SDF) ground truth as would be required for supervised learning.  Instead, they propose to only use occupancy or SDF supervision near the surfaces of objects; elsewhere, they rely on constraints on the gradient of the occupancy or SDF adapted from Gropp et al. 2020.  They perform a thorough evaluation on several benchmark datasets and compare against state-of-the-art competing methods.  They show that they outperform competing methods, even though in some cases their method has access to less supervisory data.    They also perform an ablation study to show the importance of various parts of the loss function.

summaries: This paper presents a new method to learn implicit 3D scene reconstructions from single image input. The main improvement is a closed-form Differentiable Gradient Sampling.  By taking spatial gradient into  consideration, the proposed method can apply back-propagation of the  loss on spatial gradients to feature maps and allow the training for the case of without dense 3D supervision.

summaries: In this paper, the authors propose a new method for single view 3D reconstruction.  A conditional (image feature prior) implicit representation framework is proposed to reconstruct 3D scene from a single view. In this paper, the authors propose that feature gradient is essential for watertight reconstruction and propose a differentiable gradient sampling method for the formulation. Experiments have been performed on both synthetic and real datasets. Superior results have been presented.

abstractText: Implicit neural shape functions, e.g. occupancy fields or signed distance functions (SDF), are promising 3D representations for modeling clean mesh surfaces. Existing approaches that use these representations for single-view object reconstruction require supervision signals from every location in the scene, posing difficulties when extending to real-world scenarios where ideal watertight geometric training data is difficult to obtain. In such cases, constraints on the spatial gradient of the implicit field, rather than the value itself, can provide a training signal, but this has not been employed as a source of supervision for single-view reconstruction in part due to the difficulties of differentiably sampling a spatial gradient from a feature map. In this paper, we derive a novel closed-form Differentiable Gradient Sampling (DGS) solution that enables backpropagation of the loss on spatial gradients to the feature maps, thus allowing training on large-scale scenes without dense 3D supervision. As a result, we demonstrate single view implicit surface reconstructions on real-world scenes via learning directly from a scanned dataset. Our model performs well when generalizing to unseen images from Pix3D or downloaded directly from the internet (Fig. 1). We comfortably advance the state-of-the-art with results on several established datasets including ShapeNet and ScannetV2; extensive quantitative analysis confirms that our proposed DGS module plays an essential role in achieving this performance improvement. Full codes are available in MaskedURL. View 0 View 1 View 2 View 3 View 4 View 5 Figure 1: Given a single RGB image as the input (the first column in each row), our model can predict its 3D implicit surface reconstruction (shown in six novel views in the last six columns). The test images for the first 3 rows are downloaded from the Internet and the last 2 rows are from the pix3d dataset. 1 Under review as a conference paper at ICLR 2022

authors:
Martı́n Abadi
Paul Barham
Jianmin Chen
Zhifeng Chen
Andy Davis
Jeffrey Dean
Matthieu Devin
Sanjay Ghemawat
Geoffrey Irving
Michael Isard

title: Tensorflow: A system for largescale machine learning

venue: In 12th {USENIX} symposium on operating systems design and implementation ({OSDI}

authors:
Matan Atzmon
Yaron Lipman

title: Sal: Sign agnostic learning of shapes from raw data

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

authors:
Aljaž Božič
Pablo Palafox
Justus Thies
Angela Dai
Matthias Nießner

title: Transformerfusion: Monocular rgb scene reconstruction using transformers

venue: arXiv preprint arXiv:2107.02191,

authors:
Angel Chang
Angela Dai
Thomas Funkhouser
Maciej Halber
Matthias Niessner
Manolis Savva
Shuran Song
Andy Zeng
Yinda Zhang

title: Matterport3d: Learning from rgb-d data in indoor environments

venue: arXiv preprint arXiv:1709.06158,

authors:
Angel X Chang
Thomas Funkhouser
Leonidas Guibas
Pat Hanrahan
Qixing Huang
Zimo Li
Silvio Savarese
Manolis Savva
Shuran Song
Hao Su

title: Shapenet: An information-rich 3d model repository

venue: arXiv preprint arXiv:1512.03012,

authors:
Zhiqin Chen
Hao Zhang

title: Learning implicit fields for generative shape modeling

venue: In IEEE Computer Vision and Pattern Recognition (CVPR),

authors:
Christopher B Choy
Danfei Xu
JunYoung Gwak
Kevin Chen
Silvio Savarese

title: 3d-r2n2: A unified approach for single and multi-view 3d object reconstruction

venue: In European conference on computer vision,

authors:
Michael G Crandall
Pierre-Louis Lions

title: Viscosity solutions of hamilton-jacobi equations

venue: Transactions of the American mathematical society,

authors:
Brian Curless
Marc Levoy

title: A volumetric method for building complex models from range images

venue: In Proceedings of the 23rd annual conference on Computer graphics and interactive techniques,

authors:
Angela Dai
Angel X Chang
Manolis Savva
Maciej Halber
Thomas Funkhouser
Matthias Nießner

title: Scannet: Richly-annotated 3d reconstructions of indoor scenes

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

authors:
Haoqiang Fan
Hao Su
Leonidas J Guibas

title: A point set generation network for 3d object reconstruction from a single image

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

authors:
Amos Gropp
Lior Yariv
Niv Haim
Matan Atzmon
Yaron Lipman

title: Implicit geometric regularization for learning shapes

venue: arXiv preprint arXiv:2002.10099,

authors:
T Groueix
M Fisher
VG Kim
BC Russell
M Aubry

title: Atlasnet: a papier-mâché approach to learning 3d surface generation (2018)

venue: arXiv preprint arXiv:1802.05384,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

authors:
Max Jaderberg
Karen Simonyan
Andrew Zisserman

title: Spatial transformer networks. Advances in neural information processing

authors:
Yue Jiang
Dantong Ji
Zhizhong Han
Matthias Zwicker

title: Sdfdiff: Differentiable rendering of signed distance fields for 3d shape optimization

authors:
Hiroharu Kato
Yoshitaka Ushiku
Tatsuya Harada

title: Neural 3d mesh renderer

authors:
Diederik Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Katrin Lasinger
René Ranftl
Konrad Schindler
Vladlen Koltun

title: Towards robust monocular depth estimation: Mixing datasets for zero-shot cross-dataset transfer

authors:
Manyi Li
Hao Zhang

title: D2im-net: Learning detail disentangled implicit fields from single images

authors:
Chen-Hsuan Lin
Chaoyang Wang
Simon Lucey

title: Sdf-srn: Learning signed distance 3d object reconstruction from static images

venue: arXiv preprint arXiv:2010.10505,

authors:
Shichen Liu
Tianye Li
Weikai Chen
Hao Li

title: Soft rasterizer: A differentiable renderer for image-based 3d reasoning

venue: The IEEE International Conference on Computer Vision (ICCV),

authors:
William E Lorensen
Harvey E Cline

title: Marching cubes: A high resolution 3d surface construction algorithm

venue: ACM siggraph computer graphics,

authors:
Lars Mescheder
Michael Oechsle
Michael Niemeyer
Sebastian Nowozin
Andreas Geiger

title: Occupancy networks: Learning 3d reconstruction in function space

authors:
Ben Mildenhall
Pratul P. Srinivasan
Matthew Tancik
Jonathan T. Barron
Ravi Ramamoorthi
Ren Ng

title: Nerf: Representing scenes as neural radiance fields for view synthesis, 2020

authors:
Zak Murez
Tarrence van As
James Bartolozzi
Ayan Sinha
Vijay Badrinarayanan
Andrew Rabinovich

title: Atlas: End-to-end 3d scene reconstruction from posed images

venue: In Computer Vision– ECCV 2020: 16th European Conference,

authors:
Michael Niemeyer
Lars Mescheder
Michael Oechsle
Andreas Geiger

title: Differentiable volumetric rendering: Learning implicit 3d representations without 3d supervision

authors:
Michael Oechsle
Songyou Peng
Andreas Geiger

title: Unisurf: Unifying neural implicit surfaces and radiance fields for multi-view reconstruction

venue: arXiv preprint arXiv:2104.10078,

authors:
Jeong Joon Park
Peter Florence
Julian Straub
Richard Newcombe
Steven Lovegrove

title: Deepsdf: Learning continuous signed distance functions for shape representation

authors:
Adam Paszke
Sam Gross
Francisco Massa
Adam Lerer
James Bradbury
Gregory Chanan
Trevor Killeen
Zeming Lin
Natalia Gimelshein
Luca Antiga

title: Pytorch: An imperative style, highperformance deep learning library

venue: In Advances in neural information processing systems,

authors:
Stefan Popov
Pablo Bauszat
Vittorio Ferrari

title: Corenet: Coherent 3d scene reconstruction from a single rgb image

venue: In European Conference on Computer Vision,

authors:
René Ranftl
Katrin Lasinger
David Hafner
Konrad Schindler
Vladlen Koltun

authors:
Shunsuke Saito
Zeng Huang
Ryota Natsume
Shigeo Morishima
Angjoo Kanazawa
Hao Li

title: Pifu: Pixel-aligned implicit function for high-resolution clothed human digitization

venue: In Proceedings of the IEEE International Conference on Computer Vision,

authors:
Vincent Sitzmann
Michael Zollhöfer
Gordon Wetzstein

title: Scene representation networks: Continuous 3d-structure-aware neural scene representations

venue: arXiv preprint arXiv:1906.01618,

authors:
Vincent Sitzmann
Julien Martel
Alexander Bergman
David Lindell
Gordon Wetzstein

title: Implicit neural representations with periodic activation functions

venue: Advances in Neural Information Processing Systems,

authors:
Shuran Song
Fisher Yu
Andy Zeng
Angel X Chang
Manolis Savva
Thomas Funkhouser

title: Semantic scene completion from a single depth image

authors:
Jiaming Sun
Yiming Xie
Linghao Chen
Xiaowei Zhou
Hujun Bao

title: Neuralrecon: Real-time coherent 3d reconstruction from monocular video

authors:
Xingyuan Sun
Jiajun Wu
Xiuming Zhang
Zhoutong Zhang
Chengkai Zhang
Tianfan Xue
Joshua B Tenenbaum
William T Freeman

title: Pix3d: Dataset and methods for single-image 3d shape modeling

authors:
Shubham Tulsiani
Tinghui Zhou
Alexei A Efros
Jitendra Malik

title: Multi-view supervision for single-view reconstruction via differentiable ray consistency

authors:
Peng Wang
Lingjie Liu
Yuan Liu
Christian Theobalt
Taku Komura
Wenping Wang

title: Neus: Learning neural implicit surfaces by volume rendering for multi-view reconstruction

venue: arXiv preprint arXiv:2106.10689,

authors:
Olivia Wiles
Georgia Gkioxari
Richard Szeliski
Justin Johnson

title: Synsin: End-to-end view synthesis from a single image

authors:
Qiangeng Xu
Weiyue Wang
Duygu Ceylan
Radomir Mech
Ulrich Neumann

title: Disn: Deep implicit surface network for high-quality single-view 3d reconstruction

venue: In Advances in Neural Information Processing Systems,

authors:
Lior Yariv
Yoni Kasten
Dror Moran
Meirav Galun
Matan Atzmon
Ronen Basri
Yaron Lipman

title: Multiview neural surface reconstruction by disentangling geometry and appearance

authors:
Lior Yariv
Jiatao Gu
Yoni Kasten
Yaron Lipman

title: Volume rendering of neural implicit surfaces

venue: arXiv preprint arXiv:2106.12052,

authors:
Wang Yifan
Felice Serena
Shihao Wu
Cengiz Öztireli
Olga Sorkine-Hornung

title: Differentiable surface splatting for point-based geometry processing

venue: ACM Transactions on Graphics (TOG),

authors:
Wei Yin
Jianming Zhang
Oliver Wang
Simon Niklaus
Long Mai
Simon Chen
Chunhua Shen

title: Learning to recover 3d scene shape from a single image

authors:
Alex Yu
Vickie Ye
Matthew Tancik
Angjoo Kanazawa

title: pixelnerf: Neural radiance fields from one or few images

venue: arXiv preprint arXiv:2012.02190,

title: 2020) to evaluate occupancy prediction of the 128×128×

title: 2020), the grid cube is a unit cube that spans

heading: 1 INTRODUCTION

text:
Recent studies have shown promising learning capabilities of implicit models for 3D geometry representations (Park et al., 2019; Mescheder et al., 2019; Mildenhall et al., 2020; Sitzmann et al., 2019; Gropp et al., 2020; Saito et al., 2019; Yu et al., 2020). Rather than explicitly representing 3D geometry with a textured mesh (Kato et al., 2018) or point cloud (Fan et al., 2017), implicit representations express a function of points in the space. Among many recent promising implicit models, the occupancy field (Mescheder et al., 2019; Popov et al., 2020; Chen & Zhang, 2019b) and the Signed Distance Field (SDF) (Park et al., 2019; Xu et al., 2019; Li & Zhang, 2021) are particularly suitable for high resolution mesh surface reconstruction. Given a query point (x, y, z) ∈ R3 in the space, the occupancy function returns the occupancy status of the point o(x, y, z) ∈ {+,−}, while the SDF value encompasses the closest distance from a point to the surface of the given 3D geometry, as well as its sign for its occupancy. The mesh surface is defined by the classification boundary of the occupancy field or the zero level-set of the SDF (Lorensen & Cline, 1987).
Despite the strong representational power of these models and their success on predicting single view objects, learning to predict complex 3D geometry, such as scenes, from a real world image remains a challenge for this class of methods. Obtaining the occupancy label, or the sign, from noisy non-watertight meshes is non-trivial (Atzmon & Lipman, 2020) (Fig. 2). Furthermore, producing the label on the fly from a large number of triangle meshes, as encountered in typical real-world scenes with complex geometry, can be computationally prohibitive (Straub et al., 2019; Chang et al., 2017; Dai et al., 2017; Song et al., 2017). Storing sufficiently dense query points with pre-computed occupancies or SDF for complex geometries, or meta data of meshes like octrees or hashing, is a possibility, but incurs severe engineering and computational burdens on storage and runtime loading (Yu et al., 2021). The challenge is exacerbated for predicting scenes in particular, since unlike objects scenes in the view frustrum typically cannot be enclosed by a pre-defined range of space and do not have a categorical canonical coordinate system (Popov et al., 2020).
Recent studies on fitting a single scene geometry (Gropp et al., 2020; Sitzmann et al., 2020) have pointed out the value of the 3D spatial gradient of the implicit field. For instance, the gradients of the SDF on surface points are actually the surface normals and the 2-norm of the spatial gradients for any point in the space is 1, known as the Eikonal equation (Crandall & Lions, 1983). Similarly, for the occupancy case, we propose to regularize the gradient of the occupancy field to be zero away from the surface (Fig. 3(c)). Thus, one can train the full occupancy field or SDF with only the surface point clouds, thanks to these constraints on the gradients of the implicit shape models, even when the value of the occupancy or the SDF itself is not available everywhere. However, existing approaches (Gropp et al., 2020; Sitzmann et al., 2020) have only demonstrated this benefit in the cases of geometry only inputs, where a network is fit to a particular scene per model without any conditioning on input images. When the locally sampled image features are used for the feedforward prediction, it is necessary to derive a differentiable gradient sampling solution over the feature pixel map.
In this paper, we tackle this challenge by deriving a closed-form differentiable gradient sampling (DGS) solution for learning a single-view 3D implicit reconstruction model (Fig. 4). Our novel propagation of the loss gradient of the spatial gradient back to the feature maps (Eq. 6) serves in addition to the existing spatial gradient sampling (Eq. 5) and loss gradient back-propagation (Eq. 4) used in existing deep learning frameworks (Paszke et al., 2019; Abadi et al., 2016). The resulting end-to-end learning scheme with supervision over the spatial gradients opens the potential to train a model that can generalize to unseen test cases of complex scenes with only surface point cloud supervision — a typical setting for real-world scene data — without requiring ground truth perquery-point labels.
Our contributions include a framework to propagate spatial gradients through a spatial feature sampling procedure, and a novel closed-form DGS solution. We further show that DGS enables training a single-image 3D implicit reconstruction model that can generalize to unseen scenes, with only imperfect surface annotations as the supervision (e.g. non-watertight meshes). Finally, experiments on synthetic (ShapeNet Chang et al. (2015)) and real-scanned data (ScannetV2 (Dai et al., 2017)) indicate that we clearly advanced state-of-the-art accuracy. To the best of our knowledge, DGS-enabled shape inference provides the first single-view implicit shape reconstruction on real scene datasets which can generalize accurately to unseen scenes from different datasets or domains (see Figure 1).
2 RELATED WORKS
3D Implicit Representations Among all the 3D representations, implicit models are advantageous for arbitrarily high resolution modeling (unlike voxels with fixed resolution and no detailed surface modeling) and easy learning (unlike mesh with varying topological structures). Implicit models typically learn to establish a mapping between a query point and the prediction of the point. Groueix et al. proposed to learn a mapping from the uv texture map to the 3D surface point. More recent works (Park et al., 2019; Mescheder et al., 2019; Xu et al., 2019; Popov et al., 2020; Chen & Zhang, 2019a) focus on mapping the query point coordinates to the signed distance field or the occupancy field. In addition to these pure geometry modeling, recent works like (Saito et al., 2019) also model the surface texture via learning a mapping from the surface point to
the RGB value, or use the RGB loss as the supervison signal (Niemeyer et al., 2020; Yariv et al., 2020). In contrast to defining textures explicitly on the surface, Mildenhall et al. (2020) proposed the volumetric rendering representation which maps the point coordinate and the viewing angle to volumetric textures. Oechsle et al. (2021); Wang et al. (2021); Yariv et al. (2021) use volume renderings for fitting high quality geometry of single scenes, but these are not generalizable to unseen scenes. In multi-view stereo setting, Murez et al. (2020); Sun et al. (2021); Božič et al. (2021) proposed to accumulate the multi-view predictions in a TSDF voxel volume. We believe our proposed DGS module is necessary when extending most of the existing implicit surface representations into feed-forward models that are conditioned on input image(s), and training these models with a loss on the gradient of the predicted field. In this work, we focus on pure geometrical reconstructions from single view, as an example of the application of the proposed DGS module.
Differentiable Operations End-to-end deep learning requires all modules in a computation path to be differentiable. Many differentiable modules are proposed to serve a particular functionality. The spatial transformer layer (Jaderberg et al., 2015) is one of the early attempts for differentiably sam-
pling the feature map. Kato et al. (2018) proposed to render the mesh into image in a differentiable manner so that we can optimize the mesh or the camera parameters to fit to a particular image. A similar idea has also been applied for rendering point clouds (Wiles et al., 2020; Yifan et al., 2019). Tulsiani et al. (2017) and Mildenhall et al. (2020) proposed to compute the ray rendering in an accumulating manner. One common feature of these differentiable operations is a strategy to soften the existing explicit modeling (Liu et al., 2019), and replace the gradient of the loss with approximate signals, making sure those signals provide guidance in the right direction. Our differentiable gradient sampler also seeks to soften the spatial gradient between the adjacent pixels.

heading: 3 LEARNING FRAMEWORK

text:
Problem definition and notations. We aim to train a feed-forward deep model for predicting the 3D implicit surface reconstruction conditioned on a single image. We denote the input RGB image as I ∈ Rm×n×3. Like most feed-forward models, our model contains an image encoder. We denote the extracted 2D features as ϕ. Our model f̂Θ(·), parameterized by Θ, takes in the image feature ϕ and predicts for each 3D location (x, y, z) the implicit value f̂Θ(ϕ;x, y, z). We denote the ground truth value as f̃(ϕ;x, y, z). For the occupancy field, the implicit field value f represents the occupancy probability. We denote the predicted occupancy probability as ôΘ(ϕ;x, y, z) ∈ [0, 1] and the ground truth binary occupancy label as õ(ϕ;x, y, z) ∈ {“ − ”, “ + ”}, where “−” represents unoccupied space and “+” vice versa. For SDF, the implicit field value f represents the signed distance between the query point and its projection on the surface. We denote the predicted and the ground truth signed distance as ŝΘ(ϕ;x, y, z) ∈ R and s̃(ϕ;x, y, z) ∈ R respectively. Background. Typical fully supervised training (Xu et al., 2019; Popov et al., 2020) imposes the loss to each individual sampled query points by comparing the predicted field value f̂Θ(ϕ;x, y, z) with the ground truth value f̃(ϕ;x, y, z). For occupancy predictions, we compute the loss as∑
{I,PI}∈D ∑
(x,y,z)∈PI BCE(ôΘ(ϕ;x, y, z), õΘ(ϕ;x, y, z)) (Fig. 3(a)), where “BCE” represents the binary cross entropy loss, D represents the whole single-view image training set, and PI represents the set of all the possible query point sampling locations (x, y, z) within the view frustum. For SDF, we compute the loss as ∑ {I,PI}∈D ∑ (x,y,z)∈PI |ŝΘ(ϕ;x, y, z)− s̃(ϕ;x, y, z)|.

heading: 3.1 PROPOSED FRAMEWORK

text:
As obtaining the full labels f̃(ϕ;x, y, z) for query points (x, y, z) from all the locations PI is nontrivial for complex geometry from real-world (Fig. 2) and due to additional engineering constraints on runtime loading and memory bottleneck for the sufficintly dense pre-computed query point occypancy labels, we propose to incorporate the loss w.r.t. the spatial gradients. For occupancy,
L = ∑
{I,PI}∈D
( ∑
(x,y,z)∈PI−P0I
λor∥∇x,y,z ôΘ(ϕ;x, y, z)∥
+ ∑
(x,y,z)∈P0+I
BCE(ôΘ(ϕ;x, y, z), “ + ”) + ∑
(x,y,z)∈P0−I
BCE(ôΘ(ϕ;x, y, z), “− ”)) (1)
where P0+I and P 0− I denotes inward and outward near-surface query points, and λor represents the loss weight of the occupancy geometrical regularization. The facing direction of the mesh surface (determining inward or outward) can be determined by normals (if available), or via rendering the surface in all views as in the RGBD captures (Dai et al., 2017) (with the surface facing toward camera as the “outward” side). The three terms in Eq. 1 correspond to the three types of losses in Fig. 3(c) (“∇”, “+”, “−”) respectively. For SDF, we set the loss similar to Gropp et al. (2020):
L = ∑
{I,PI}∈D
( ∑
(x,y,z)∈PI
λsr|∥∇x,y,z ŝΘ(ϕ;x, y, z)∥2 − 1|+ ∑
(x,y,z)∈P0I
|ŝΘ(ϕ;x, y, z)|
+ ∑
(x,y,z)∈P0I
λsn∥∇x,y,z ŝΘ(ϕ;x, y, z)−∇x,y,z s̃(ϕ;x, y, z)∥) (2)
where P0I denotes the query points on the ground truth surface only, and λsr and λsn represents the loss weight for the signed distance geometric regularization term and normal term respectively. The three terms represents the Eikonal regularization (Crandall & Lions, 1983), surface zero-SDF loss, and the surface normal loss respectively, with the last term optional (Gropp et al., 2020). Note now
the loss functions in Eq. 1, 2 no longer require the implicit ground truth label for query points far away from the ground truth surface, enabling training with the real-world imperfect scanned data. In both cases, the availability of surface normals is optional.
Practically, we found only our proposed loss for the occupancy version (Fig. 3(c), Eq. 1) demonstrates capability to be applicable for the real-world single-view feed-forward scenario, and use Eq. 1 rather than Eq. 2 in all of our experiments. First, learning with the Eikonal term (Eq. 2) suffers severely from its sensitivity of model initialization. We found either the SDF or the Truncated SDF (TSDF) representation model cannot easily converge during training. Specifically, we found it poses numerical difficulties when the model is learned to predict a large distance value or a constant truncation value for the majority of the query points in the air that is far from any surfaces. This becomes a major problem when we extend from the single objects scenario (Mescheder et al., 2019; Popov et al., 2020) or single scene fitting (Gropp et al., 2020). Second, our loss function (Eq. 1) can significantly save memory footprint during training and enable large batch size training, which is crucial in our learning framework. Unlike in Eq. 2 where all the query points require spatial gradient computation, which leads to 3× higher of memory footprint, in our loss function (Eq. 1), only the non-surface query points do. We set the query point batch size of the non-surface points (512 in practice) to be much smaller than the critical near-surface points (4096 in practice), enabling learning with the image batch size of 32 in our real-scene training.

heading: 3.2 SAMPLING WITH DIFFERENTIABLE GRADIENTS

text:
A distinct difference compared to existing works is that our spatial gradient ∇x,y,z f̂Θ(ϕ;x, y, z) is also conditioned on the pixels in ϕ. This is in stark contrast to existing works. Gropp et al. (2020) devised the model fΘ fits to a single scene, and the spatial gradient ∇x,y,z f̂Θ(x, y, z) can be conveniently computed because it does not involve the sampling procedure. Jiang et al. (2020); Lin et al. (2020) used a non-spatial global feature for inference and hence bypassed sampling. In our learning framework, the spatial gradient computation must undergo the sampling procedure.
We name our gradient computation involving the sampling operation as the Pixel Conditioned Gradients and derive a closed-form solution, Differentiable Gradient Sampling (DGS), for handling forward and backward propagation. Figure 4(a) provides an illustration of our training pipeline. Each layer in our network tracks both the response of the layer as well as its spatial gradient w.r.t. (x, y, z). While it is well established to track the layer-wise spatial gradient for fully-connected (FC) layers or convolutions in existing works (Gropp et al., 2020), tracking the spatial gradients and back-propagating the loss to the feature map pixels ϕ through the sampling module has not been studied. To this end, we derive the closed-form sampling scheme for tracking and propagating the spatial gradients ∂ϕ∂x,y,z through the sampling layer.
Background - 2D Differentiable Sampling. Differentiably sampling pixel values from a grid of 2D feature map with the given pixel locations (i, j) is a common operation. Throughout our paper, we define the pixel coordinates in the Normalized Device Coordinate (NDC) system that ranges from -1 to 1. As illustrated in Fig. 4(d), given the feature map ϕ and the sampling locations (i, j), the resulting sampled value ϕ(i, j) is
ϕ(i, j) = (1− α)((1− β)ϕA + βϕB) + α((1− β)ϕC + βϕD) (3)
Please refer to Fig. 4(c) for the definitions of α, β and ϕA, ϕB , ϕC , ϕD. Without loss of generality we use bilinear interpolation. During training, the gradient from the loss can be back-propagated via
∂L ∂ϕA = (1− α)(1− β) ∂L ∂ϕ(i, j) . (4)
2D Differentiable Gradient Sampling. Our learning framework (Sec. 3.1, Fig. 4(a)) requires the extension of the sampling capability from just the feature value response ϕi,j to its spatial gradient ∇i,jϕ(i, j). During the forward and the backward propagation, both the sampled feature response ϕi,j and its spatial gradient ∇i,jϕ(i, j) are recorded for further propagation (Fig. 4(e)):
∂ϕ(i, j)
∂i = (1− β)(ϕC − ϕA) + β(ϕD − ϕB) h . (5)
Hence, during the forward pass, we compute the spatial gradient via Eq. 5 in addition to the existing value sampling (Eq. 3). During the backward pass, we compute the loss gradient over the spatial gradient via
∂L ∂ϕA = (1− α)(1− β) ∂L ∂ϕ(i, j) + (−1− β h ) ∂L ∂ ∂ϕ(i,j)∂i + (−1− α w ) ∂L ∂ ∂ϕ(i,j)∂j , (6)
where w and h are the width and height of a pixel in NDC, s.t. w = 2/W and h = 2/H for the feature map with the size W ×H . 3D Differentiable Gradient Sampling. We now extend the sampling to 3D. We model the camera as the pin-hole camera. For any point (x, y, z) in the camera space, we seek for its projected 2D locations (i, j) based on the camera focal length f0 via i = yf0/z, j = xf0/z. The feed forward pass is
∂ϕ(x, y, z) ∂x = ∂ϕ(i, j) ∂j · f0 z , ∂ϕ(x, y, z) ∂y = ∂ϕ(i, j) ∂i · f0 z ,
∂ϕ(x, y, z) ∂z = ∂ϕ(i, j) ∂i · (−yf0 z2 ) + ∂ϕ(i, j) ∂j · (−xf0 z2 ).
(7)
During the backward propagation procedure, the DGS accumulates the gradient via
∂L
∂ϕA =
∂L ∂ϕ(x, y, z) · (1− α)(1− β) + ∂L ∂ ∂ϕ(x,y,z)∂x · (−1− α w ) f0 z + ∂L ∂ ∂ϕ(x,y,z)∂y · (−1− β h ) f0 z
+ ∂L ∂ ∂ϕ(x,y,z)∂z · ((1− β h )( yf0 z2 ) + ( 1− α w )( xf0 z2 )).
(8)

heading: 4 EXPERIMENTS

heading: 4.1 LEARNING FROM SYNTHETIC DATA (SHAPENET)

text:
Dataset. Following Popov et al. (2020), we conduct shapeNet experiments with the low-realism and the high-realsim renderings. For low-realism, we use the renderings of various models from Choy et al. (2016). Following Mescheder et al. (2019), we split the dataset into the training set of 30661 models, the validation set of 4371 models, and the test set of 8751 models. During testing, we follow Mescheder et al. (2019) to only test the first rendering of each model. Similar to all the prior works, we use the same 13 classes to report any performance via grouping all the test cases into the same class. For high-realism, we follow the split and evaluation protocol of the single-object scenes Popov et al. (2020) - we only use the 13 categories used by Choy et al. (2016), filter out repeated samples, and finally construct the data with 666,565 models from the training set, 96,084 from the validation set, and 189,748 from the test set. Following Popov et al. (2020), we only evaluate the first 1% test cases of the test set (1898 samples).
Comparison with state-of-the-art approaches. We choose to compare to recent representative state-of-the-art approaches, including OccNet (Mescheder et al., 2019), DISN (Xu et al., 2019),
CoReNet (Popov et al., 2020) and D2IM-Net (Li & Zhang, 2021). Quantitative results are reported in Tab. 3 for the low-realism evaluation setting, and in Tab. 1 for the high-realism setting, where our results outperform all the state-of-the-art results in all the settings. Please refer to Sec. A for qualitative comparison, implementation details and the evaluation metrics.
For our low-realism evaluation setting, all the baseline approaches reported their results in the papers. For OccNet and DISN, the reported results are based on knowing the category canonical view prior, which demonstrates considerable privilege with respect to accuracy Popov et al. (2020). Hence, we mark the results from the literature as OccNet-Privilege and DISN-Privilege (OccNetPriv. and DISN-Priv. for short). Compared to their privileged setting, our results demonstrate superior results even without the category canonical view prior privilege. To further provide the the baseline results where these two approaches are without the category canonical view prior, we retrain their models with the released codes, and report the results in the “OccNet” and “DISN” row respectively. The results further provide evidence that the category canonical view prior demonstrates privilege on accuracy, as observed in Popov et al. (2020).
For our high-realism evaluating setting, since the results for Occnet and DISN are not directly available from the paper, we retrain their model from the released codes (without the category canonical view prior) on the high-realism data provided by Popov et al. (2020). We also retrain CoReNet using their released codes and obtain results better than reported in their original paper (61.5 as mIoU instead of 59.5). Our model still consistently outperforms all the baselines on almost all the categories. We emphasize that our comparison to CoReNet is a direct ablation comparison where we directly follow their experimental setup and model size. Furthermore, our approach only utilizes the closeto-surface occupancy labels, which is a subset of the whole occupancy labels for all the query point. Our model still outperforms the baseline and provides convincing evidence that our supervision loss function demonstrate better results than the original CoReNet approach. We attribute the superiority of our approach to our particular focus toward the difficult close-to-surface region similar to hard negative mining. Ablation Study. We further conduct ablation studies to further validate the importance of derived DGS module, via attempting work-around training methods without DGS. i) NoGrad - To test the performance of the baseline when only the surface data are provided, we train this ablation model in exactly the same way compared to its original model, with the only exceptions that only the near-surface points are equipped with training labels, and we mask out all the training labels from the non-surface query points (where it is not necessary to backpropagate the loss gradient of the spatial gradient using DGS). To further evaluate how the rate of known voxels affect the learning performance, we enlarge the near-surface region and evaluate when the rate is 10%, 30% and 50%. An increase in the performances among these baselines would indicate the importance of knowing more voxel labels if our proposed gradient loss (Eq. 1) is not imposed. ii) FixedE - We train with both the near-surface as well as the spatial gradient supervision, without DGS - meaning the loss gradient of the spatial gradient would not back-propagate to any module before the sampling module - in our case, the feature encoder network. Note all the other losses without gradient sampling can still back-propagated to the feature maps. We report the ablation results in Tab. 1. Both experiments are conducted with the high-realism ShapeNet data. Once again, we comfortably outperform all the baselines, validating the essential roles of DGS in our learning framework.

heading: 4.2 LEARNING FROM REAL SCANNED DATASETS (SCANNETV2)

text:
Dataset. We use ScannetV2 (Dai et al., 2017) for training and evaluating the performance of the models on the real images. We follow the standard training / testing split as used in Sun et al. (2021) and Murez et al. (2020), where 1513 scenes are for trainval (with 1201 for training and 312 for validation), and 100 for testing. Each scene is provided with multiple image capture as well as the associated camera pose. We train the models with all the views given in the training / validation set (2423872 frames in total, after filtering out frames with invalid extrinsic poses), while for testing, we select 10 frames with different extrinsic poses for each test scene. Practically, since all the frames of the scenes are in the form of video clips, with adjacent frames associated with similar extrinsic cameras poses, we select the 10 frames for each frame via extracting every 100 frames from each scene video (e.g. Frame 1, 101, 201, ..., 901), resulting in 1000 frames in total in our test set (10 frames per scene, with 100 test scenes).
Metrics. We use the same evaluation metrics following Sun et al. (2021) and Murez et al. (2020). Since we are the first, to our knowledge, to evaluate single view 3D implicit reconstruction on the ScannetV2 benchmark, and here we only evaluate the geometries within the camera view frustum rather than the whole scene geometries as in Sun et al. (2021); Murez et al. (2020). In addition, we also only evaluate the geometries in front of the amodal depth for each pixel ray. We define the amodal depth for a pixel ray to be delineated by the minimum between the closest structure-category surface (e.g. walls and doors, etc) and the farthest surface. In practice, in order to accommodate the evaluation of the surfaces right on the amodal depth, we slack the evaluation scope with a factor λ (1.05 in our case) multiplied with the amodal depth. This evaluation protocol would be equivalent to the “single-layered” protocol used in Sun et al. (2021), within our single-view scenario. Due to the inherent ambiguity of the scaling and shifting of the predicted 3D single view geometry (Ranftl et al., 2019; Yin et al., 2021), we follow them by regressing the best scale and shift comparing the predicted depth with the ground truth depth. For approaches that predict 3D surfaces, the rendered depth map from the predicted mesh would be used for calculating the scale and shift. Baselines. Since we demonstrate the first attempt to conduct the implicit reconstruction of singleview scenes, very few exiting works provide a direct baseline performance to our task. For OccNet (Mescheder et al., 2019), since its image feature extraction is not local, and its gradient propa-
gation does not require sampling, we use its direct application with Gropp et al. (2020) “OccNet + GeoReg” as one baseline. We train DISN (Xu et al., 2019) with the TSDF voxelization labels (Murez et al., 2020; Sun et al., 2021; Božič et al., 2021). For CoReNet (Popov et al., 2020), we stick to its own voxelization and internal filling toolbox for obtaining the training labels. Since CoReNet can only predict geometries within the fixed range of space, we tried our best to pick the best cube location based on the dataset statistics. We also incorporate depth approaches (Ranftl et al., 2019; Yin et al., 2021) for comparison by finetuning their weights on ScannetV2. Lastly, we compare with the two ablation models NoGrad and FixedE as introduced in Sec. 4.1. Results. We provide quantitative comparison in Tab. 2 and qualitative comparison in Fig. 5 respectively. Our model outperforms all state-of-the-art approaches as well as the ablation models. Compared to the synthetic data scenario in Tab. 3 and 1, our model demonstrates even larger advantages when compared to existing state-of-the-art approaches, as our motivation stems from addressing learning from imperfect 3D labels directly from the real scan data. Compared to single-view depth prediction approaches trained on massive data (Lasinger et al., 2019; Yin et al., 2021), our approach does not prevail on “Acc” and “Prec”. This is due to the fact that these two metrics only project the predicted surface to the ground truth surface, giving advantages to approaches that only predict the visible surface. Our approach still prevails for other metrics (Chamfer Distance and F1), which are considered as the most important metrics (Božič et al., 2021; Sun et al., 2021; Murez et al., 2020). Generalization to Unseen Scenes (Pix3d and Open-Domain Images). We further test our model without further finetuning directly on unseen scenes for evaluating the generalizability of the learned model. We tried our model on Pix3d Sun et al. (2018) as well as test images downloaded directly from the internet. We provide a detailed qualitative comparison in Fig. 6 and more results in Fig. 1. The results further indicate our learning framework exhibits promise for unseen scene generalization. Failure Cases. We provide two representative failure cases of our approach in Fig. 7. The first case (the first row) demonstrates difficulties in predicting the floor occupancy as a result of the noisy and non-watertight mesh during training. The second case (the second row) shows that our model cannot identify objects (e.g. chairs) and well predicts the invisible part of them.

heading: 5 CONCLUSIONS

text: We have presented our learning framework for real-world 3D implicit surface reconstruction from a single view image. Owing to our unique learning framework that directly train the model from the raw scan data and our novel occupancy loss function over the gradients, we are able to go beyond the existing works for single-objects reconstruction or single view fitting. Thanks to our differentiable gradient sampling module we enable efficient and memory-economic end-to-end training and achieved scene-level single view reconstruction for the first time, to our knowledge.

heading: A ADDITIONAL DETAILS AND RESULTS ON SHAPENET

text:
Metrics. We use the Intersection-of-Union (IoU) for evaluating the model performances. For lowrealism renderings, we follow the protocol of IoU benchmarking from Mescheder et al. (2019) that evaluating occupancy of the specified 100K query points in the space. For high-realism renderings, we follow the protocol from Popov et al. (2020) to evaluate occupancy prediction of the 128×128× 128 grid. Following Popov et al. (2020), the grid cube is a unit cube that spans from −0.5 to 0.5 for x and y, and from f/2 to 1 + f/2 for z in the camera coordinate system.
Implementation details of our model (DGS). We build our model on top of one representative state-of-the-art model - CoReNet (Popov et al., 2020) for its 3D volumetric implicit modeling capability. This is because the space of the geometry predictions are always within the unit cube volume. We maintain exactly the same encoder and decoder architecture (ResNet50 (He et al., 2016) and the ray traced skip connection respectively), optimization parameters (e.g. the learning rate and the epsilon of Adam Kingma & Ba (2014)) and the prediction format (per-query occupancy probability). We set λor to be 0.01 in our loss function (Eq. 1).
Qualitative Results. We provide qualitative results in Fig. 9-13 as the additional illustration of the performance of all the approaches on ShapeNet.

heading: B ADDITIONAL RESULTS ON SCANNETV2

text:
Implementation details for DGS. Since the DISN variant architectures (Xu et al., 2019; Niemeyer et al., 2020; Yu et al., 2020) demonstrate higher degress of flexibility of representing any point in the space, in contrast to 3D convolutions that are fixed for a particular range of space (Popov et al., 2020), we build our model based on the former architectures. Similar to Niemeyer et al. (2020); Yu et al. (2020), we use 5 residual blocks in the fully connected part with the first 3 blocks receiving the 2D features and its spatial gradients. Our image encoder uses the same RexNext101 architecture (Xie et al., 2017) as in Yin et al. (2021) and starts the training with the pre-trained weights. We use a batch size of 32 images with 2048 near-surface points and 512 non-surface points per image. Similar to the ShapeNet experiments, we set λor to be 0.01. We found in the real-world setting, this parameter setting is crucial for the convergence of the learning procedure.
Quantitative Results for 2.5D evaluation metrics. We provide the quantitative 2.5D evaluation results in Tab. 4. Our approach constantly outperforms all other implicit based baselines. While our approach falls short slightly when compared to the existing depth-based approaches, we claim that our approach is not directly trained with the massive depth training data. Our results indicate that our depth performance is still on par with the state-of-the-art depth prediciton approaches.
Additional qualitative results. We provide additional qualitative results on ScannetV2 Fig. 14-18.

heading: C ANALYSIS WITH NUMERICAL GRADIENT APPROXIMATION

text:
As part of our proposed learning framework with the loss imposed on the spatial gradient, the differentiable gradient sampling module (Sec. 3.2) has a numerical alternative where we can perturb the query points to get its simulated numerical gradients. We further compare the closed-form performance with the numerical counterpart in Tab. 5. Please find qualitative comparison in Fig. 9-13.
We provide a closer look at our comparison to the simulated gradient baseline. We found in our experiments that the simulated gradient baseline demonstrates relatively severe convergence difficulties. As shown in Fig. 8, the simulated gradient model (red) does not observe loss drop in the first 10k training iterations. Despite its subsequent loss drop, which indicates the simulated gradient model can still learn to predict the geometry in moderate accuracy, its converged training loss is still higher than our DGS (blue). This aligns with our benchmarking evaluations in Tab. 5 that the closed-form variant (DGS) achieve better performance.
Table 5: Intersection over Union % (IoU ↑) benchmarking comparison between the closed-form solution and the numerical gradients on the high-realism ShapeNet.
Category Craft Rifle Disp. Lamp Spk. Box Chair Bench Car Plane Sofa Table Phone Mean
Numerical 62.4 68.3 63.7 46.2 60.6 54.9 61.0 49.0 70.7 61.1 69.2 55.6 77.2 61.5 Closed-form 63.3 70.4 65.7 49.0 61.2 57.0 62.1 50.9 70.1 62.8 70.3 56.9 78.3 62.9
0 0.5 1 1.5 2 2.5 3 Iterations 105
0.3
0.4
0.5
0.6
0.7
0.8
Th e
To ta
l L os
s Numerical Closed-Form
Figure 8: Convergence Analysis for the comparison between the closed-form (blue) and the numerical counterpart (red). Notably, the numerical counterpart does not observe loss drop in the first 10k iterations.

heading: D ADDITIONAL GENERALIZABILITY QUALITATIVE RESULTS

text: We provide more qualitative results on our model generalizing to unseen images downloaded from the Internet in Fig. 19. Each result is visualized in 6 views. Once again, these additional visual results further indicate that our model demonstrates good generalizability capability for handling unseen indoor scene images.

heading: E EVALUATION OF THE GENERALIZABILITY TO MATTERPORT3D DATASET

text:
To quantatively evaluate our generaizability to new datasets, we provide our results on the Matterport3D dataset (Chang et al., 2017) using our model trained from ScannetV2 as used in Sec. 4.2.
The Matterport3d dataset (Chang et al., 2017) collects indoor scenes of 90 full houses. Textured meshes are provided along with the scanned real images. Since the Matterport3D dataset does not provide an official train/test split, and our model is not trained on Matterport3d, we evaluate our model over all the 90 scenes in the dataset. In particular, we use the 129600 images (without the elevated pitch views where the camera is looking up at the ceiling) from the dataset, use the 10 first images from each scene (resulting in 900 test images in total), and test their reconstruction performance in the same way as we tested on ScannetV2. We provide the quantative evaluation result in Tab. 6. We can see from the table that the results align with our evaluation in Tab. 2 that our model demonstrates clear advantages. We also noticed that the AdelaiDepth Yin et al. (2021) baseline (finetuned on ScannetV2) performance is closer to our results (F1) score compared to our ScannetV2 evaluation. This is probably due to the fact that AdelaiDepth was directly pre-trained from a massive number of training images and potentially gives it some privilege when generalize to a new dataset.

heading: F DETAILS OF THE FORMULATION AND DERIVATION OF DGS

text:
In this section, we provide the full formulation as well as the derivation procedure of our DGS in this section as the supplementary to Sec. 3.2.
In Eq. 4, we provided the backward gradient for differentiable sampling for Pixel A. The backward gradient for Pixel B, C and D are
∂L ∂ϕB = (1− α)β · ∂L ∂ϕ(i, j)
∂L ∂ϕC = α(1− β) · ∂L ∂ϕ(i, j)
∂L ∂ϕD = αβ · ∂L ∂ϕ(i, j)
(9)
In Eq. 5, we provided the forward gradient in the vertical pixel direction. The forward gradient of the horizontal pixel direction is
∂ϕ(i, j)
∂j = (1− α)(ϕB − ϕA) + α(ϕD − ϕC) w
(10)
In Eq. 6, we provided the backward gradient for only Pixel A. The backward gradient for Pixel B, C and D are
∂L ∂ϕB = (1− α)β · ∂L ∂ϕ(i, j) + (−β h ) · ∂L
∂ ϕ(i,j)∂i +
1− α w · ∂L ∂ ϕ(i,j)∂j
∂L ∂ϕC = α(1− β) · ∂L ∂ϕ(i, j) + 1− β h · ∂L ∂ ϕ(i,j)∂i + (−α w ) · ∂L
∂ ϕ(i,j)∂j ∂L
∂ϕD = αβ · ∂L ∂ϕ(i, j) + β h · ∂L ∂ ϕ(i,j)∂i + α w · ∂L ∂ ϕ(i,j)∂j
(11)
In Eq. 8, we provided the backward gradient for only Pixel A in the 3D setting. The backward gradient for Pixel B, C and D are
∂L
∂ϕB =
∂L ∂ϕ(x, y, z) · (1− α)β + ∂L ∂ ∂ϕ(x,y,z)∂x · 1− α w · f0 z + ∂L ∂ ∂ϕ(x,y,z)∂y · (−β h ) · f0 z
+ ∂L ∂ ∂ϕ(x,y,z)∂z · (β h · yf0 z2 − 1− α w · xf0 z2 )
∂L
∂ϕC =
∂L ∂ϕ(x, y, z) · α(1− β) + ∂L ∂ ∂ϕ(x,y,z)∂x · (−α w ) · f0 z +
∂L ∂ ∂ϕ(x,y,z)∂y · 1− β h · f0 z
+ ∂L ∂ ∂ϕ(x,y,z)∂z · (−1− β h · yf0 z2 + α w · xf0 z2 )
∂L
∂ϕD =
∂L ∂ϕ(x, y, z) · αβ + ∂L ∂ ∂ϕ(x,y,z)∂x · α w ) · f0 z +
∂L ∂ ∂ϕ(x,y,z)∂y · β h · f0 z
+ ∂L ∂ ∂ϕ(x,y,z)∂z · (−β h · yf0 z2 − α w · xf0 z2 )
(12)
As an illustration of the derivation for Eq. 8, we can obtain the backward gradient via
∂L
∂ϕA =
∂L ∂ϕ · ∂ϕ ∂ϕA + ∂L ∂ ∂ϕ∂x · ∂ ∂ϕ∂x ∂ϕA + ∂L ∂ ∂ϕ∂y · ∂ ∂ϕ∂y ∂ϕA + ∂L ∂ ∂ϕ∂z · ∂ ∂ϕ∂z ∂ϕA
(13)
The four partial derivatives with respect to ϕA in Eq. 13 can be determined based on the sampling rules. For example, based on Eq. 3, we can quickly obtain the first derivative with respect to ϕA via
∂ϕ
∂ϕA = (1− α)(1− β) (14)
For the other three derivatives with respect to ϕA in Eq. 13, we can substitute Eq. 5, 10 into Eq. 7, and obtain them via
∂ ∂ϕ∂x ∂ϕA = −f0 z · 1− α w
∂ ∂ϕ∂y ∂ϕA = −f0 z · 1− β h
∂ ∂ϕ∂z ∂ϕA = yf0 z2 · 1− β h + xf0 z2 · 1− α w
(15)
We can then obtain Eq. 8 via substituting Eq. 14, 15 into Eq. 13.

title: DIFFERENTIABLE GRADIENT SAMPLING FOR LEARNING IMPLICIT 3D SCENE RECONSTRUCTIONS


summaries: This work studies the problem of efficient pretraining of large-scale models for language and vision representations, namely the issue of significant memory requirements for models with billions to trillions of parameters. Authors propose two modifications: first, to reduce the memory load and improve convergence at the initial stage of training, they suggest to train a multilayer model with shared parameters and then unshare them. Second, to maximize GPU utilization with offloading, authors develop a method for granular CPU offloading, which keeps larger chunks of the model in GPU memory. When combined, the proposed methods allow the authors to train a 10 trillion parameter model on 512 GPUs.

summaries: The authors propose to train very large neural language models via a "Sharing-Delinking" paradigm. The proposed method first trains a model with weights shared across layers. In this way, the model appears to be smaller and it can fit into fewer GPUs. At some point, the authors delink the weights, and continue training the model in the conventional way.  The authors also propose a granular CPU offloading mechanism to save CPU memory.

summaries: The paper proposes a technique called Pseudo-to-Real (P2R) for reducing the computational and time requirements of training massive (or Giant) models with trillions of parameters. The key idea of P2R is a two phase training approach for Giant models. The first phase involves training a smaller version of the model (a.k.a., Pseudo-Giant) which is obtained by making all layers share parameters. The second phase involves training the Giant model after initializing with Pseudo-Giant weights. The paper further proposes Granular CPU offloading which is to offload some but not all model parameters to CPU memory to reduce GPU memory consumption. Finally, the paper provides some evaluation results to demonstrate P2G.

summaries: The paper proposed an interesting strategy to reduce the training time for large scale language models consisting of stacking layers with identical structures. Users first trained the models with shared parameters across the layers, then relax the tie constraints so that parameters at different layers are updated differently.  The paper showed some empirical evidence that the proposed strategy converged faster given a limited training time budget and demonstrated the feasibility of training a 10T model.

abstractText: Recent expeditious developments in deep learning algorithms, distributed training, and even hardware design for large models have enabled training extreme-scale models, say GPT-3 and Switch Transformer possessing hundreds of billions or even trillions of parameters. However, under limited resources, extreme-scale model training that requires enormous amounts of computes and memory footprint suffers from frustratingly low efficiency in model convergence. In this paper, we propose a simple training strategy called “Pseudo-to-Real” for high-memoryfootprint-required large models. Pseudo-to-Real is compatible with large models with architecture of sequential layers. We demonstrate a practice of pretraining unprecedented 10-trillion-parameter model, an order of magnitude larger than the state-of-the-art, on solely 512 GPUs within 10 days. Besides demonstrating the application of Pseudo-to-Real, we also provide a technique, Granular CPU offloading, to manage CPU memory for training large model and maintain high GPU utilities. Fast training of extreme-scale models on a decent amount of resources can bring much smaller carbon footprint and contribute to greener AI.

authors:
Mandeep Baines
Shruti Bhosale
Vittorio Caggiano
Naman Goyal
Siddharth Goyal
Myle Ott
Benjamin Lefaudeux
Vitaliy Liptchinsky
Mike Rabbat
Sam Sheiffer
Anjali Sridhar
Min Xu

title: Fairscale: A general purpose modular pytorch library for high performance and large scale training

authors:
Rishi Bommasani
Drew A Hudson
Ehsan Adeli
Russ Altman
Simran Arora
Sydney von Arx
Michael S Bernstein
Jeannette Bohg
Antoine Bosselut
Emma Brunskill

title: On the opportunities and risks of foundation models

venue: arXiv preprint arXiv:2108.07258,

authors:
Tom B. Brown
Benjamin Mann
Nick Ryder
Melanie Subbiah
Jared Kaplan
Prafulla Dhariwal
Arvind Neelakantan
Pranav Shyam
Girish Sastry
Amanda Askell
Sandhini Agarwal
Ariel Herbert-Voss
Gretchen Krueger
Tom Henighan
Rewon Child
Aditya Ramesh
Daniel M. Ziegler
Jeffrey Wu
Clemens Winter
Christopher Hesse
Mark Chen
Eric Sigler
Mateusz Litwin
Scott Gray
Benjamin Chess
Jack Clark
Christopher Berner
Sam McCandlish
Alec Radford
Ilya Sutskever
Dario Amodei

title: Language models are few-shot learners

venue: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems

authors:
Mark Chen
Jerry Tworek
Heewoo Jun
Qiming Yuan
Henrique Ponde
Jared Kaplan
Harri Edwards
Yura Burda
Nicholas Joseph
Greg Brockman

title: Evaluating large language models trained on code

venue: arXiv preprint arXiv:2107.03374,

authors:
Rewon Child
Scott Gray
Alec Radford
Ilya Sutskever

title: Generating long sequences with sparse transformers

venue: arXiv preprint arXiv:1904.10509,

authors:
Alexis Conneau
Kartikay Khandelwal
Naman Goyal
Vishrav Chaudhary
Guillaume Wenzek
Francisco Guzmán
Edouard Grave
Myle Ott
Luke Zettlemoyer
Veselin Stoyanov

title: Unsupervised cross-lingual representation learning at scale

venue: In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: Pre-training of deep bidirectional transformers for language understanding

venue: Association for Computational Linguistics. 1https://github.com/huggingface/transformers

authors:
Ming Ding
Zhuoyi Yang
Wenyi Hong
Wendi Zheng
Chang Zhou
Da Yin
Junyang Lin
Xu Zou
Zhou Shao
Hongxia Yang
Jie Tang

title: Cogview: Mastering text-to-image generation via transformers

authors:
Jianbo Dong
Zheng Cao
Tao Zhang
Jianxi Ye
Shaochuang Wang
Fei Feng
Li Zhao
Xiaoyong Liu
Liuyihan Song
Liwei Peng
Yiqun Guo
Xiaowei Jiang
Lingbo Tang
Yin Du
Yingya Zhang
Pan Pan
Yuan Xie

title: EFLOPS: algorithm and system co-design for a high performance distributed training platform

venue: In IEEE International Symposium on High Performance Computer Architecture,

authors:
Li Dong
Nan Yang
Wenhui Wang
Furu Wei
Xiaodong Liu
Yu Wang
Jianfeng Gao
Ming Zhou
Hsiao-Wuen Hon

title: Unified language model pre-training for natural language understanding and generation

venue: Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems

authors:
William Fedus
Barret Zoph
Noam Shazeer

title: Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity

venue: arXiv preprint arXiv:2101.03961,

authors:
Xianyan Jia
Le Jiang
Ang Wang
Jie Zhang
Xinyuan Li
Wencong Xiao
Yong Li
Zhen Zheng
Xiaoyong Liu
Wei Lin

title: Whale: Scaling deep learning model training to the trillions

venue: arXiv preprint arXiv:2011.09208,

authors:
Jared Kaplan
Sam McCandlish
Tom Henighan
Tom B Brown
Benjamin Chess
Rewon Child
Scott Gray
Alec Radford
Jeffrey Wu
Dario Amodei

title: Scaling laws for neural language models

authors:
Dmitry Lepikhin
HyoukJoong Lee
Yuanzhong Xu
Dehao Chen
Orhan Firat
Yanping Huang
Maxim Krikun
Noam Shazeer
Zhifeng Chen

title: Gshard: Scaling giant models with conditional computation and automatic sharding

authors:
Mike Lewis
Yinhan Liu
Naman Goyal
Marjan Ghazvininejad
Abdelrahman Mohamed
Omer Levy
Veselin Stoyanov
Luke Zettlemoyer

title: BART: Denoising sequence-to-sequence pretraining for natural language generation, translation, and comprehension

authors:
Junyang Lin
Rui Men
An Yang
Chang Zhou
Ming Ding
Yichang Zhang
Peng Wang
Ang Wang
Le Jiang
Xianyan Jia

title: M6: A chinese multimodal pretrainer

venue: arXiv preprint arXiv:2103.00823,

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized BERT pretraining approach

authors:
Ilya Loshchilov
Frank Hutter

title: Decoupled weight decay regularization

venue: In 7th International Conference on Learning Representations,

authors:
Stephen Merity
Caiming Xiong
James Bradbury
Richard Socher

title: Pointer sentinel mixture models

venue: In 5th International Conference on Learning Representations,

authors:
Deepak Narayanan
Mohammad Shoeybi
Jared Casper
Patrick LeGresley
Mostofa Patwary
Vijay Anand Korthikanti
Dmitri Vainbrand
Prethvi Kashinkunti
Julie Bernauer
Bryan Catanzaro

title: Efficient large-scale language model training on gpu clusters

venue: arXiv preprint arXiv:2104.04473,

authors:
Denis Paperno
Germán Kruszewski
Angeliki Lazaridou
Ngoc Quan Pham
Raffaella Bernardi
Sandro Pezzelle
Marco Baroni
Gemma Boleda
Raquel Fernández

title: The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)

authors:
David A. Patterson
Joseph Gonzalez
Quoc V. Le
Chen Liang
Lluis-Miquel Munguia
Daniel Rothchild
David R. So
Maud Texier
Jeff Dean

title: Carbon emissions and large neural network training

venue: arXiv preprint arXiv:2104.10350,

authors:
Alec Radford
Karthik Narasimhan
Tim Salimans
Ilya Sutskever

title: Improving language understanding by generative pretraining

venue: URL https://s3-us-west-2.amazonaws.com/openai-assets/researchcovers/ languageunsupervised/language understanding paper

authors:
Alec Radford
Jeffrey Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners, 2019

authors:
Alec Radford
Jong Wook Kim
Chris Hallacy
Aditya Ramesh
Gabriel Goh
Sandhini Agarwal
Girish Sastry
Amanda Askell
Pamela Mishkin
Jack Clark
Gretchen Krueger
Ilya Sutskever

title: Learning transferable visual models from natural language supervision

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

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J Liu

title: Exploring the limits of transfer learning with a unified text-to-text transformer

venue: Journal of Machine Learning Research,

authors:
Samyam Rajbhandari
Jeff Rasley
Olatunji Ruwase
Yuxiong He

title: Zero: Memory optimizations toward training trillion parameter models

venue: In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis,

authors:
Samyam Rajbhandari
Olatunji Ruwase
Jeff Rasley
Shaden Smith
Yuxiong He

title: Zeroinfinity: Breaking the gpu memory wall for extreme scale deep learning

venue: arXiv preprint arXiv:2104.07857,

authors:
Aditya Ramesh
Mikhail Pavlov
Gabriel Goh
Scott Gray
Chelsea Voss
Alec Radford
Mark Chen
Ilya Sutskever

title: Zero-shot text-to-image generation

venue: arXiv preprint arXiv:2102.12092,

authors:
Jeff Rasley
Samyam Rajbhandari
Olatunji Ruwase
Yuxiong He

title: Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters

venue: In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining,

authors:
Jie Ren
Samyam Rajbhandari
Reza Yazdani Aminabadi
Olatunji Ruwase
Shuangyan Yang
Minjia Zhang
Dong Li
Yuxiong He

title: Zero-offload: Democratizing billion-scale model training

venue: arXiv preprint arXiv:2101.06840,

authors:
Stephen Roller
Sainbayar Sukhbaatar
Arthur Szlam
Jason Weston

title: Hash layers for large sparse models

venue: arXiv preprint arXiv:2106.04426,

title: Turing-nlg: A 17-billion parameter language model by microsoft

venue: Microsoft Research Blog,

authors:
Noam Shazeer
Azalia Mirhoseini
Krzysztof Maziarz
Andy Davis
Quoc V. Le
Geoffrey E. Hinton
Jeff Dean

title: Outrageously large neural networks: The sparsely-gated mixtureof-experts layer

authors:
Noam Shazeer
Youlong Cheng
Niki Parmar
Dustin Tran
Ashish Vaswani
Penporn Koanantakool
Peter Hawkins
HyoukJoong Lee
Mingsheng Hong
Cliff Young
Ryan Sepassi
Blake A. Hechtman

title: Mesh-tensorflow: Deep learning for supercomputers

venue: In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems

authors:
Mohammad Shoeybi
Mostofa Patwary
Raul Puri
Patrick LeGresley
Jared Casper
Bryan Catanzaro

title: Megatron-lm: Training multi-billion parameter language models using model parallelism

authors:
Yu Sun
Shuohuan Wang
Shikun Feng
Siyu Ding
Chao Pang
Junyuan Shang
Jiaxiang Liu
Xuyi Chen
Yanbin Zhao
Yuxiang Lu

title: Ernie 3.0: Large-scale knowledge enhanced pre-training for language understanding and generation

venue: arXiv preprint arXiv:2107.02137,

authors:
Alex Wang
Amanpreet Singh
Julian Michael
Felix Hill
Omer Levy
Samuel R. Bowman

title: GLUE: A multi-task benchmark and analysis platform for natural language understanding

authors:
Ruibin Xiong
Yunchang Yang
Di He
Kai Zheng
Shuxin Zheng
Chen Xing
Huishuai Zhang
Yanyan Lan
Liwei Wang
Tie-Yan Liu

title: On layer normalization in the transformer architecture

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

authors:
Linting Xue
Noah Constant
Adam Roberts
Mihir Kale
Rami Al-Rfou
Aditya Siddhant
Aditya Barua
Colin Raffel

title: mT5: A massively multilingual pre-trained text-to-text transformer

venue: In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,

authors:
An Yang
Junyang Lin
Rui Men
Chang Zhou
Le Jiang
Xianyan Jia
Ang Wang
Jie Zhang
Jiamang Wang
Yong Li

title: M6-t: Exploring sparse expert models and beyond

venue: arXiv preprint arXiv:2105.15082,

authors:
Yang You
Jing Li
Sashank J. Reddi
Jonathan Hseu
Sanjiv Kumar
Srinadh Bhojanapalli
Xiaodan Song
James Demmel
Kurt Keutzer
Cho-Jui Hsieh

title: Large batch optimization for deep learning: Training BERT in 76 minutes

venue: In 8th International Conference on Learning Representations,

authors:
Wei Zeng
Xiaozhe Ren
Teng Su
Hui Wang
Yi Liao
Zhiwei Wang
Xin Jiang
ZhenZhang Yang
Kaisheng Wang
Xiaoda Zhang

title: Pangu-α: Large-scale autoregressive pretrained chinese language models with auto-parallel computation

venue: arXiv preprint arXiv:2104.12369,

authors:
Minjia Zhang
Yuxiong He

title: Accelerating training of transformer-based language models with progressive layer dropping

authors:
Zhengyan Zhang
Yuxian Gu
Xu Han
Shengqi Chen
Chaojun Xiao
Zhenbo Sun
Yuan Yao
Fanchao Qi
Jian Guan
Pei Ke

title: Cpm-2: Large-scale cost-effective pre-trained language models

venue: arXiv preprint arXiv:2106.10715,

authors:
Zhengyan Zhang
Xu Han
Hao Zhou
Pei Ke
Yuxian Gu
Deming Ye
Yujia Qin
Yusheng Su
Haozhe Ji
Jian Guan

title: Cpm: A large-scale generative chinese pre-trained language model

venue: AI Open,

authors:
Zhu Zhang
Jianxin Ma
Chang Zhou
Rui Men
Zhikang Li
Ming Ding
Jie Tang
Jingren Zhou
Hongxia Yang

title: M6-ufc: Unifying multi-modal controls for conditional image synthesis

venue: arXiv preprint arXiv:2105.14211,

authors:
Yukun Zhu
Ryan Kiros
Richard S. Zemel
Ruslan Salakhutdinov
Raquel Urtasun
Antonio Torralba
Sanja Fidler

title: Aligning books and movies: Towards story-like visual explanations by watching movies and reading books

venue: In 2015 IEEE International Conference on Computer Vision,

heading: 1 INTRODUCTION

text:
Foundation models with self-supervised learning on big data have become an emerging paradigm of artificial intelligence systems [2], as they mostly possess high transferability to a wide range of downstream tasks and even multiple modalities. The scale of foundation models across domains, including natural language processing, computer vision, and cross-modality representation learning, have been growing tremendously from millions to trillions of parameters [7; 23; 24; 26; 33; 37; 3; 29; 14; 11; 16; 47; 46; 42; 44; 38] thanks to the concurrent advancement in distributed training framework [20; 27; 31; 28; 1; 12; 30] and hardware design, and these studies have made a demonstration of the neural scaling law [13]. However, the training of these transformerbased models incurs high financial costs and even environmental damage due to the massive carbon footprint and thus training extreme-scale models under a decent amount of resources but with high efficiency should be a fundamental goal for both the research and industrial communities to achieve, which promotes the progress of greener AI [22; 34].
Generally there are two tracks of research in large-scale pretraining, dense models and sparse expert models respectively. A typical case of large-scale dense models is GPT-3 [3], a 175-billionparameter transformer model trained with 10, 000 GPUs for months, incurring striking financial and environmental costs. Researchers have been searching for methods to training large-scale models with a decent amount of costs. Solutions include effective management of memory with gradient and optimizer state partitioning [27] or more efficient model parallelism and pipeline parallelism [37; 20; 12]. A series of following studies apply those techniques to realize fast training of 10-billion-parameter transformers with hundreds of GPUs in 1 − 2 months. [16; 46; 44; 38] Sparse expert models with large model capacity are capable of fast training owing to the combination of data parallelism and expert parallelism [35; 14; 11; 32], and it is even accessible to train a 1-trillion-parameter transformer with no more than 500 GPUs [42].
Be there as it may, a question emerges in our mind: is it possible to train an extreme-scale model with only a decent amount of resources, e.g., training a 10-trillion-parameter model with around 500 GPUs? Tackling the problem requires the utilization of external memory except for GPU memory, for instance, CPU memory or even NVMe storage [31; 28]. These methods resolve the problem of
high memory footprint, but instead, their extra cost is low training efficiency caused by the frequent swap in-and-out between memories.
In this paper, we provide a solution to training large models that require high memory footprint, and we demonstrate a successful practice of pretraining an unprecedented extreme-scale model with over 10 trillion parameters, an order of magnitude larger than the previous state-of-the-arts [11; 42]. The whole pretraining was conducted on solely 512 NVIDIA-V100 GPUs and lasted around 10 days. A simple and effective training strategy called “Pseudo-to-Real” enables the sharing and delinking of parameters. This training strategy is compatible with architectures built by stacking layers of an identical structure, including dense models like GPT [23; 24; 3], BERT [7], or sparse expert models like M6 [16; 42]. It is essentially a two-stage training. It first trains a relatively small model but with a computation graph of a large one with the utilization of cross-layer parameter sharing, and we name it “Pseudo Giant”. Then it builds a correspondingly large model and delinks the parameters of the shared layer for second-stage model initialization. In this way, we achieve fast convergence in the first stage as the training costs much less memory and speeds up with large batches. Parameter sharing that addresses the communication overhead improves training speed as well. The second-stage training is responsible for the final convergence for better performance.
We unlock the secret of pretraining an unprecedented extreme-scale model with over 10 trillion parameters on limited resources of 512 GPUs. Compared with the previous M6-T on around 500 GPUs, we do not have a significant increase in computation resources but level up the model scale by an order of magnitude. Besides the application of the “Pseudo-to-Real” training strategy, we provide a faster offloading mechanism for both management of CPU memory for parameter storage and utility of GPUs. We successfully train the M6-10T within 10 days to reach strong performance in log perplexity evaluation and outperform the baseline M6-T.
Contributions at a glance are below:
• We illustrate the training difficulty of extreme-scale models on limited resources and provide a simple but effective solution called “Pseudo-to-Real”. Upstream and downstream evaluation demonstrates the effectiveness of the strategy.
• We further demonstrate a successful practice of pretraining a 10-trillion-parameter model on 512 GPUs and reach an outstanding performance within 10 days.

heading: 2 RELATED WORK

text:
Large-Scale Pretraining In recent years, pretrained language models with growing magnitudes of parameters have been proposed, keeping to raise the validated upper limit of scaling law for model capacity w.r.t the number of parameters [13]. Earlier milestones of extreme-large models come from GPT-2 [24] and Megatron-LM [37], which demonstrates that scaling the transformer model up to billions of parameters can result in improvement on language modeling benchmarks [19; 21]. TuringNLG [33], as a successor, implements a 17-billion-parameter transformer and achieves further lower perplexity. Similar phenomena are also observed on classification language tasks by T5 model [26]. The GPT-3 [3] pushes the boundary of model scale to 1,750 billion parameters and demonstrates its striking effectiveness on downstream tasks in even zero-shot settings. Furthermore, large-scale pretraining has recently demonstrated success in other fields, including pretraining on other languages or cross-lingual pretraining [41; 46; 44; 6], cross-modal pretraining [25; 29; 16; 8; 48] and code generation [4]. Along with the benefit from increasing the model scale, the concern of unaffordable pretraining cost in time, computation resource and energy keeps emerging [22; 2], resulting in the strong demand for more efficient and greener large-scale pretraining [34].
Methods to Train Larger Models Faster Our work is about designing algorithm to train extremescale models efficiently. Researchers have been demonstrating different types of methods to reach this objective. To reduce the computational cost during training, some studies introduce sparsity to the model. Mixture-of-Experts (MoE) [35] was proposed and was reintroduced in Mesh Tensorflow [36]. It shows that MoE with sparsity can significantly improve the model scale efficiently without increasing computation, and the models achieved state-of-the-art performance in language modeling and machine translation. GShard [14] extended it to a tremendously large scale of 600B parameters with the sophisticated collaborated design of model architecture and demonstrated its effectiveness
across over 100 languages. Similarly, Switch Transformer reached 1.6 trillion parameters and showed its strong performance on NLU tasks. Those models with high sparsity are computationally efficient, and therefore though they possess large model capacity they still can be trained with high efficiency. Most other studies still focus on training dense models to validate the scaling law, and thus the emerged problem is the distributed training of large models. The most influential distributed framework should be DeepSpeed that proposed ZERO [27] that partitions optimizer states and gradients to multiple GPU devices, and ZERO-offload [31] as well as ZERO-Infinity [28] can even offload parameters to CPU memory and NVMe storage. The memory footprint management makes training extremely large models on limited resources possible. However, such offloading mechanisms still have some defects that they may fail to fully utilize the fast hardware. For example, when offloading all parameters to the CPU, the GPU memory can be idle. In this work, we tackle this issue by proposing a granular offloading mechanism that can determine which parameters to be offloaded.
In addition to reducing the amount of computation in a single iteration, another route to speedup training is to reduce the needed iterations for model convergence. Child et al. [5] and Xiong et al. [40] propose to put forward the layernorm operations in transformer blocks for more stabilized and faster convergence. You et al. [43] employs a layer-wise adaptive optimizer to enable super-large pretraining batches. Zhang & He [45] progressively increases the layer dropping rate in a stochastic manner which significantly speedups pretraining.

heading: 3 APPROACH

heading: 3.1 MODEL ARCHITECTURE

text:
Choice in model architecture depends on several factors. First, the architecture should contain a sequence of stacking layers, as the sequential structure enables cross-layer parameter sharing. We prefer a simple encoder or decoder architecture, instead of an encoder-decoder framework where there are cross attentions that bring extra parameters and incur difficulties in activation checkpointing. Second, a model of such architecture should be compatible with different types of downstream tasks, including understanding and generation, and it is more preferable that it can be compatible with multiple modalities. Third, as we mention that dense models and sparse expert models are two main tracks of large-scale pretraining, we prefer the model that can flexibly become whether dense or sparse expert models. Therefore, we select M6 [16] as an option, and we evaluate the effects of our method on M6 of different scales and types.
M6 is built with stacking transformer layers, which includes self attention and feed-forward neural nets (FFN). For the transformation from dense models to sparse expert models, we should only replace FFN layers with the Mixture-of-Expert (MoE) layers. MoE consists of multiple experts, which are usually FFNs distributed on different devices. A gating network decides the dispatching and combining behaviors of each token and thus tokens can be processed in diverse devices. Such mechanism is a combination of data parallelism and expert parallelism, and thus it is highly efficient though with large model capacity. For the training, to realize the learning of both understanding and generation, the model is trained with text denoising and language modeling on plain text data and with image-based text denoising and image captioning on multimodal data. The model is compatible with different types of downstream tasks and can process information of multiple modalities.

heading: 3.2 PSEUDO-TO-REAL

text: This section demonstrates the details of “Pseudo-to-Real” two-stage training strategy that enables fast training of high-memory-footprint-required transformer models. The strategy consists of two stages. The first stage trains a model with many fewer parameters but with a large computation graph (“Pseudo Giant”), and the second stage trains a corresponding large model (“Real Giant”) initialized with the delinked weights of the shared layer. Thus we name the strategy “Pseudo-to-Real”, and the general idea is demonstrated in Figure 1.

heading: 3.2.1 “PSEUDO” STAGE: LAYER-WISE PARAMETER SHARING

text:
The core of “Pseudo” stage is to train a Pseudo-Giant that shares parameters across layers. Cross-layer parameter sharing has proved successful in maintaining satisfactory performance while keeping a much smaller amount of parameters. We introduce it to training an extreme-scale model, and we hypothesize that the first-stage training can gain benefits from cross-parameter sharing as it can address communication overhead and it consumes much less memory footprint. Also, as Pseudo Giant with much fewer parameters is not bounded by memory, it can be trained with large batches for acceleration.
Suppose we build an M6 model withL layers that share parameters across all layers. The Pseudo Giant though consists of a computation graph of a L-layer transformer, its number of weight parameters and optimizer states should be 1/L of those of the original one. As to the gradients, we can accumulate the gradients of each layer in the backward computation process, and therefore the amount of gradients becomes 2/L of the original one. Such saving in memory enables much faster training with larger batches. Also, it is capable to use fewer resources even due to lower memory consumption.
This can also be applied to sparse expert models, as their architecture is stacking transformer layers. However, different from dense models, the MoE partition their weights to all devices, and each token is distributed to a selected expert by a routing function. Thus the models can trained with data parallelism across all devices and expert parallelism where parameters at each device are activated. This mechanism though enlarges model capacity significantly without much efficiency loss, yet it limits the flexible usage of GPU resources in different stages, say training and inference. Due to the application of data parallelism, we can implement dense models on different numbers of GPU devices at the training and inference stage. However, as mentioned above, the combination of data parallelism and expert parallelism requires the identical amount of devices at different stages.
To tackle this problem, we specifically design methods for expert merging and partitioning. Due to cross-layer parameter sharing, at the Pseudo stage it is capable of implementing more experts on each device, and it is available to use fewer GPUs while keeping the total number of experts. We name it expert merging for simplicity. At the Real stage, the delinking of parameters with the same amount o devices will cause out-of-memory errors. We delink the parameters and partition experts to more devices, e.g., from 256 to 512, and thus there are fewer experts on each device where memory is sufficient. Thus we name it expert partitioning. Therefore, now it is available to use different number of GPU devices at different stages for training sparse expert models.

heading: 3.2.2 “REAL” STAGE: DELINKING THE SHARED PARAMETERS

text:
We name a large model without cross-layer parameter sharing “Real Giant”, in comparison with Pseudo Giant. Both models share an identical computation graph, but possess different numbers
of parameters. Given a Pseudo Giant fully trained until convergence, we apply the delinking of cross-layer shared parameters to accelerate Real Giant training. There is no need to train a large model from scratch. The model can start its convergence from low perplexity.
Embedding initialization can be directly restored, but the layer weights should be treated specially. In practice, there is only one layer of weights θshared in Pseudo Giant, and there are L layers of weights {θ1, θ2, · · · , θL} in Real Giant. Thanks to their identical structure, each layer of Real Giant can be initialized with θshared. Without further training, this model is equivalent to a fully-trained Pseudo Giant.
This extremely simple training strategy is highly beneficial for the high-memory-footprint-required large models, especially extreme-scale models like the 10-trillion-parameter M6. While the first stage of training saves much time for faster convergence, we can use a decent amount of computational resources in this stage as lower efficiency in this stage becomes acceptable. Therefore, in the practice of training an extreme-scale M6, we apply CPU offloading to utilize CPU memory. Therefore, we can use a limited amount of resources, e.g., 512 GPUs, to train an unprecedented 10-trillion-parameter model efficiently, which is an order of magnitude larger than the state-of-the-arts.

heading: 3.2.3 TIMING FOR SWITCHING

text:
A question naturally emerges: when should we switch from the Pseudo stage to the Real stage? As mentioned above, the greatest advantage of Pseudo stage for training is the significantly faster convergence speed. Yet the performance of Pseudo Giant is bounded by its limitation in the number of parameters. Training Pseudo Giant until convergence apparently incurs much waste of time.
In practice, we present a simple strategy to determine the training step to switch from Pseudo to Real based on the convergence speed. During the training of the Pseudo stage, we evaluate a training step in a fixed interval by attempting to transfer it into the Real stage and training for a small while (e.g., 30 minutes). After that, we will revert the model parameters to the evaluated step and continue the training of the Pseudo stage for the same training time as the Real stage. If the decreasing speed of loss in the Real stage surpasses that of the Pseudo stage, we determine the evaluated training step as the best switching point for the next-stage training.

heading: 3.3 EXPERIMENTS

text: In this section, we demonstrate experiments to evaluate the effectiveness of the training strategy by evaluating the model quality and observing the performance improvement.

heading: 3.3.1 PERFORMANCE EVALUATION

text:
We aim to validate two hypotheses: 1. Pseudo-to-Real training paradigm can help the model reach competitive performance with the one trained from scratch without parameter sharing; 2. Pseudo-toReal can effectively accelerate convergence in training of large-scale models.
We aim at discovering the model quality brought by different pretraining strategies. “Pseudo (P)” refers to training with parameter sharing across layers, and “Pseudo-to-Real (P2R)” refers to the proposed two-stage training of a sharing-delinking paradigm. We pretrain the model on BookCorpus [49] and English Wikipedia [7], which are corpora with around 16GB of plain texts. Following Radford et al. [23] and Lewis et al. [15], we use a vocabulary of around 50, 000 subwords. Each
sample consists of sentences from an identical passage, and we use a sequence length of 512 and correspondingly truncate or pad the sequence. Following the common practice in pretraining [7; 17], we apply AdamW optimizer [18] for optimization. To determine the most suitable learning rate of the two stages for fast convergence, we have made some preliminary tests and finally used the peak learning rate of 2e-4 for “Pseudo” stage and 8e-5 for “Real” stage, respectively. We use a cosine decaying mechanism for learning rate scheduling and a warm-up ratio of 0.001. We pretrain the models until convergence and transfer them to downstream tasks.
For comprehensive analysis, the experiments include natural language understanding and generation tasks. Notably, for NLU tasks, we follow Devlin et al. [7] and validate transfer effects on GLUE [39]. For natural language generation, we specify zero-shot language modeling and text summarization to evaluate both models’ upstream and downstream quality on generation. Specifically, we conduct experiments on WikiText-103 and Gigaword for both tasks.
Experimental results on NLU tasks are demonstrated in Table 1. For better comparison, we also present the experimental results of BERT with a similar amount of parameters, and it shows that M6 can achieve a better performance over the baseline. From Table 1, we find that Pseudo-only training brings a significant performance downgrade in the 8 tasks. It is fair to believe that the limited amount of parameters hinders downstream performance, though with a large computation graph. However, in comparison, we find that P2R can help the model perform similarly to the one trained from scratch without parameter sharing. Specifically, M6 with P2R can even outperform the baseline on 6 tasks, including SST-2, MNLI, QQP, MRPC, RTE, and STS-B. We further validate their performance on zero-shot language modeling and text sumarization tasks. Table 2 demonstrate the model performance on both tasks, and we additionally add Megatron-LM [37] and [10] for better comparison. Similarly, Pseudo-only training leads to far worse performance, and P2R is also able to help the model achieve similar performance with Real training.

heading: 3.3.2 EFFICIENCY EVALUATION

text:
Next, we evaluate the efficiency of training strategies by observing model performance under the condition of identical computation budget.
Experimental Setup To satisfy the requirements of high memory footprint where the two-staged training can make a difference in training efficiency, we conduct experiments on a 1.4B-parameter model. We present the details of model configuration in Table 3.
For both Pseudo and Real Giant, we train them with a total batch size of 6, 144. In practice, we use a micro-batch size of 32 and a gradient accumulation step of 4, and we train our models on 48 NVIDIA-V100 GPU devices. We compare the model performance with a limited budget, where models have been trained for the same duration of time. Correspondingly, “Pseudo” has been trained for around 25, 000 steps, Real has been trained for around 4, 700 steps, and “P2R” has been trained for 12, 000 steps.
Model Performance Table 3 demonstrates the training speed of the models. We report their training speed on 48 NVIDIA-V100 GPU devices with their consumed samples per second. Training Real Giant model from scratch is highly time-consuming. Pseudo Giant training has an advantage of around 5 times of convergence speed over Real Giant training. We also observe the loss convergence of both P2R and Real Giant trained from scratch. In Figure 2 we present the development of pretraining language modeling loss, which is the log perplexity, on the time basis. The log perplexity of P2R decreases much faster than that of the Real Giant, with an advantage larger than 0.3.
Experimental results on downstream tasks presented in Table 4 are consistent with our hypothesis that Pseudo-to-Real training can speed up training effectively. In the setup of limited budget, the M6 model trained with Pseudo-to-Real can outperform the Pseudo Giant and Real Giant in both language modeling and text generation.

heading: 4 TOWARDS A 10-TRILLION-PARAMETER MODEL

text: Previously, training large-scale models brings tons of challenges to the collaboration algorithm design, distributed training, as well as hardware design, etc. Training a GPT-3 of 175B parameters with a combination of data parallelism on over 500GB of data should cost around 300 GPU-years. Later with the emergence of partitioning on optimizer states, gradients, and even weights, GPU memory can be fully utilized without performance degradation. Now we can even use the CPU memory or even NVMe storage to store the parameters, but we have to bear the costs of efficiency. Therefore, we attempt to tackle the difficulty of extreme-scale model training from the perspective of algorithm design and thus we apply the aforementioned Pseudo-to-Real training strategy to train an extreme-scale model.

heading: 4.1 MODEL SETUP

text:
We design a 10-trillion-parameter M6 model with the combination of existing methods and proposed strategies to demonstrate a case of how to train an extreme-scale model efficiently. In comparison with the previous studies of trillion-parameter models [11; 42], this one is almost 10 times larger. To efficiently utilize the memory, we adopt Mixture-of-Experts and we replace every FFN layer with the memory-efficient MoE layers. Notably, we remove the auxiliary loss that consumes memory and
demonstrates little effects on model quality, and we follow Yang et al. [42] to apply expert prototyping for improved model quality and training stability. To be more specific, the hidden size dmodel is 1024 and the intermediate size dff is 9984. The number of model layers is 48. For each MoE layer, there are 10240 experts distributed on multiple devices. We use 80 prototypes of experts based on our experience in preliminary experiments. The training batch size per GPU device is 8. We implement models on EFLOPs, a distributed training platform with an advanced server architecture and a new network topology [9]. Specifically our models are trained on a cluster of 8-GPU workers connected by RDMA networks with a bandwidth of 100Gb. The CPU memory of each worker is around 750GB. The expert distribution and Granular Offloading strategies are supported by Whale framework[12].

heading: 4.2 GRANULAR CPU OFFLOADING

text:
To utilize CPU memory for larger models with fewer resources, we apply the Granular CPU offloading which has a higher efficiency compared with the conventional CPU offloading. Previously we note that conventional offloading mechanisms offload all parameters, which may fail to effectively utilize GPU. To improve the efficiency of CPU offloading, we propose a new CPU offloading mechanism called Granular CPU offloading. The training process is composed of phases including “Forward (Fn)”, “Backward(Bn)” and “Apply (An)”. Offloading all model parameters to CPU in Fn and Bn requires loading parameters from CPU to GPU memory twice. Activation checkpointing that brings recomputation needs the parameters loaded in Bn. An requires the gradients offloaded from GPU memory to CPU memory. Assume the model parameter size is W, the above processes bring parameter movement of size 4 ∗W. In offloading, PCIE is the bottleneck of the whole training process. We observe that when offloading all parameters with recomputation, the GPU memory is idle. We can fill up the GPU memory by selective offloading, leaving part of the model in GPU memory. In this way, the model can be accelerated by reducing across-device memory copy. In our preliminary experiment, with the setting of training a 48-layer 78B-parameter M6 model on 8 NVIDIA-V100 GPU devices, the step-time costs 89 seconds when fully offloading the parameters into CPU memory. In comparison, granularly offloading the first 24 layers into CPU memory and leaving the remaining 24 layers on GPU reduces the step-time to only 45 seconds. The significant difference in training step-time indicates that the time-cost of parameter movement between CPU and GPU will dominate the training step-time when offloading is employed, thus the granularity of offloading and the utilization of GPU should be seriously considered in extreme-scale pretraining. In addition, offloading the whole model can result in OOM error in the CPU when the model is extremely large.
With Granular CPU offloading, we successfully implement a 10-trillion-parameter M6 model on solely 512 NVIDIA-V100 GPUs. Furthermore, at the Pseudo stage, we can train a Pseudo Giant with a computation graph of 10 trillion parameters only with 256 GPU devices without the utilization of CPU memory for offloading. Thus in our practice, we train a Pseudo Giant with only 256 GPUs, and
then partition the experts and redistribute them to 512 GPUs. This saves the usage of GPU resources, which is more resource-efficient and also more environmentally friendly.

heading: 4.3 ANALYSIS

text:
Training Efficiency We pretrain the M6-10T with Pseudo-to-Real training strategy for around 10 days, and we additionally train an M6-10T from scratch without the strategy for around 3 days for comparison. The Real stage training can be facilitated without Out-of-Memory errors with the help of Granular CPU offloading, but its step-time is only around 180s. In contrast, the step-time of Pseudo stage is only 14s without the cost of offloading, which greatly boosts the training efficiency of M6-10T P2R. The M6-10T with P2R has been trained for 15k steps, but the one from scratch has been trained for solely 1.3k steps. We record the log perplexity of both models trained on the M6-Corpus on the time basis in Figure 4(a). Results show that within the same length of time M6-10T with P2R can outperform the one trained from scratch by a large margin.
Convergence Analysis We have trained M6-10T with Pseudo-to-Real strategy for around 10 days, and the model converges to a low level of log perplexity based on the upstream evaluation on the M6-Corpus. For better comparison, we also show the convergence performance of the 1T-parameter M6-T model proposed in the previous work. As shown in Figure 4(b), the observation is consistent with our intuition that the model with a larger capacity can converge faster on the sample basis, and it should achieve the best performance on language modeling. What leaves open is whether it can positively lead to better downstream performance concerning different types of downstream tasks. Finetuning extreme-scale models should be difficult and there is still much room for us to discover the potential of extreme-scale models.

heading: 5 CONCLUSION AND FUTURE WORK

text:
Pseudo-to-Real training strategy is a simple and effective way to train large-scale models that are highly memory consuming, and it is also essential to training extremely large models with limited resources with significantly higher training efficiency. We unlock pretraining unprecedented extremescale models with 10 trillion parameters with limited resources of 512 GPUs in 10 days. Besides the application of Pseudo-to-Real training strategy, we further provide Granular CPU offloading to enhance GPU utility while breaking the GPU memory wall with a cost in efficiency. The advances take a leap towards extreme-scale model training beyond implementation on limited resources. With only a few GPU cards, training large models with tens or hundreds of parameters has become accessible to many researchers. We believe that this can motivate low carbon dioxide production and encourages the progress of green AI.
ETHICS STATEMENT
This work is highly concerned with large-scale language models and multimodal pretrained models. These models have been pretrained on broad data of plain texts and image-text pairs, which might contain harmful information, such as hate speech, terrorism, pornography, etc. We have put much efforts to remove these kinds of data in our datasets by quality evaluation on texts and images. However, this problem cannot be eliminated and ignored, and it is common in the pretraining community. For those models that are not trained on commonly-used public datasets, we will carefully release the model checkpoints before careful evaluation, and also limit the access to avoid misconduct.
REPRODUCIBILITY STATEMENT
This work is generally reproducible. Following the description in Section 3, researchers can easily implement the training strategy on the codebases for pretraining, including Huggingface Transformer 1, Fairseq 2, etc.