Graphics Capsule: Learning Hierarchical 3D Face Representations
from 2D Images

The connections of the human brain are thought to be sparse and hierarchical [1, 4, 15, 9], which inspires the design of capsule networks to present the objects with dynamic parse trees. Given inputs, capsule networks [12, 26, 13, 18, 27, 35] will encode the images to a set of low-level capsules, which describe the local entities of the objects, and then assemble them into higher-level capsules to describe more complicated entities. The parameters of capsules are usually with explicit meanings, which enables the interpretability of neural networks. Recently, some capsule networks have been proposed to explore the hierarchy of objects. SCAE [18] proposes to describe the objects with a set of visualizable templates through unsupervised learning. However, SCAE can only handle simple 2D objects like digits. HP-Capsule [35] extends SCAE to tackle human faces, which proposes subpart-level capsules and uses the compositions of subparts to present the variance of pose and appearance. Due to the limitation of 2D representations, HP-Capsule can only tackle faces with small poses. Sabour et al. [27] propose to apply the capsule network to human bodies, but it needs optical flow as additional information to separate the parts. In this paper, we propose graphics capsules to learn the hierarchical 3D representations from unlabeled images.

We evaluate the graphics capsule performance on the unsupervised face segmentation task. Several methods have been proposed for this challenging task. DFF [7] proposes to use non-negative matrix factorization upon the CNN features to discover semantics, but it needs to optimize the whole dataset during inference. Choudhury et al. [6] follow a similar idea, which uses k-means to cluster the features obtained by a pre-trained network. SCOPS [14] and Liu et al. [20] propose to constrain the invariance of images between TPS transformation. However, their methods rely on the concentration loss to separate parts, leading to similar silhouettes of different parts. HP-Capsule [35] proposes a bottom-up schedule to aggregate parts from subparts. The parts of the HP-Capsule rely on the learning of subpart-part relations, which is unstable when tackling faces with large poses. Compared with these methods, our IGC-Net can provide interpretable 3D representations of the parts, which are with salient semantics and keep semantic consistency across the in-the-wild faces with various poses.

Learning to recover the 3D face from 2D monocular images has been studied for years. Following the analysis-by-synthesis strategy, many methods [32, 39, 5, 8] propose to estimate the parameters of the 3D Morphable Model [24], which describes the faces with a uniform topology pre-defined by humans. Recently, several works [34, 37, 38] have been proposed to only use the symmetric character of faces to learn 3D face reconstruction. Under the graphics decomposition, these methods achieve promising results. Inspired by them, we propose the graphics capsule to learn the hierarchical 3D face representations from images, which provides insight into how neural networks understand faces by learning to decompose them into a set of parts.

To learn a hierarchical 3D representation from unlabeled images, we propose an Inverse Graphics Capsule Network (IGC-Net), whose capsules are composed of interpretable CG descriptions, including a depth map $\mathbf{D}\in\mathbb{R}^{H\times W}$, an albedo map $\mathbf{A}\in\mathbb{R}^{C\times H\times W}$ and 3D pose parameters $\mathbf{p}\in\mathbb{R}^{1\times 6}$ (rotation angles and translations). Our IGC-Net is applied to human faces, which have been widely used to investigate the human vision system due to their similar topology structures and complicated appearances. The overview of IGC-Net is shown in Figure 1. Following a bottom-up schedule, a CNN-based image encoder first encodes the input image $\mathbf{I}$ into the shape and the albedo embeddings $\mathbf{f}_{s}$ and $\mathbf{f}_{a}$:

$\displaystyle\mathbf{f}_{s},\mathbf{f}_{a}$

$\displaystyle=\mathrm{ImageEncoder}(\mathbf{I}).$

(1)

Then a Graphics Decomposition Module (GDM) is employed to decompose the global embeddings into a set of part-level embeddings, which can be further decoded into interpretable graphics capsules:

$\displaystyle\begin{aligned} \{\hat{\mathbf{e}}_{s}^{1},...,\hat{\mathbf{e}}_{s}^{M}\},\{\hat{\mathbf{e}}_{a}^{1},...,\hat{\mathbf{e}}_{a}^{M}\}&=\mathrm{GDM}(\mathbf{f}_{s},\mathbf{f}_{a}),\\ \{\mathbf{D}_{p}^{m},\mathbf{A}_{p}^{m},\mathbf{p}_{p}^{m}\}&=\mathrm{GraphicsDecoder}(\hat{\mathbf{e}}_{s}^{m},\hat{\mathbf{e}}_{a}^{m}),\\ \end{aligned}$

(2)

where $\hat{\mathbf{e}}_{s}^{m}$ is the shape embedding of the $m$th part, $\hat{\mathbf{e}}_{a}^{m}$ is the corresponding albedo embedding, $M$ is the number of part capsules, and $\Theta_{p}^{m}:\{\mathbf{D}_{p}^{m},\mathbf{A}_{p}^{m},\mathbf{p}_{p}^{m}\}$ is a graphics capsule that describes a part with depth, albedo, and 3D pose. Afterward, the part capsules are assembled according to their depth to generate the global object capsule:

	$\displaystyle\mathbf{V}^{m}(i,j)$	$\displaystyle=\mathbf{1}_{m=\mathop{\mathrm{argmin}}\limits_{n}(\mathbf{D}^{n}_{p}(i,j))},$		(3)
	$\displaystyle\mathbf{D}_{o}$	$\displaystyle=\sum\nolimits_{m}\mathbf{V}^{m}\odot\mathbf{D}_{p}^{m},$
	$\displaystyle\mathbf{A}_{o}$	$\displaystyle=\sum\nolimits_{m}\mathbf{V}^{m}\odot\mathbf{A}_{p}^{m},$
	$\displaystyle\mathbf{p}_{o}$	$\displaystyle=\frac{1}{M}\sum\nolimits_{m}\mathbf{p}_{p}^{m},$

where $V^{m}_{i,j}$ is the visibility map of the $m$th part capsule at the position $(i,j)$, $\odot$ is the element-wise production, and $\Theta_{o}:\{\mathbf{D}_{o},\mathbf{A}_{o},\mathbf{p}_{o}\}$ is the object capsule. During assembly, a capsule is visible at $(i,j)$ only when its depth is smaller than the others. The part-level depth and albedo maps are multiplied with their visibility maps and aggregated as one, respectively, and the object pose is the average of part poses. In the object capsule, both the depth $\mathbf{D}_{o}$ and albedo $\mathbf{A}_{o}$ are defined in the canonical space, and the pose $\mathbf{p}_{o}$ is used to project the 3D object to the image plane. Finally, by estimating the lighting $\mathbf{l}$ with another module similar to  [34], the recovered image $\hat{\mathbf{I}}$ is generated by the differentiable rendering $\Lambda$ [16]:

$\displaystyle\hat{\mathbf{I}}=\Lambda(\mathbf{D},\mathbf{A},\mathbf{p},\mathbf{l}).$

(4)

When training IGC-Net, we can minimize the distance between the input image $\mathbf{I}$ and the reconstructed image $\hat{\mathbf{I}}$ following the analysis-by-synthesis strategy, so that the network parameters can be learned in an unsupervised manner.

Humans can decompose an object into a set of parts and construct a hierarchy by just observation. To realize this ability in neural networks, we propose the Graphics Decomposition Module (GDM) to decompose the global embedding of the image into a set of semantic-consistent part-level descriptions. The illustration of GDM is shown in Figure 2.

Taking shape decomposition as an example, GDM maintains $M$ shape basis $\{\mathbf{W}_{s}^{m}\}$ as the implicit part templates. Given the global embeddings $\mathbf{f}_{s}$ extracted in Eqn. 1, GDM performs cross attention between the global embedding and the basis to get $M$ disentangled $D$ dimensional embeddings:

$\displaystyle\begin{aligned} \mathbf{e}_{s}^{m}&=\mathbf{f}_{s}\mathbf{W}_{s}^{m},~{}~{}~{}m=1,...,M.\\ \end{aligned}$

(5)

To further reduce the entanglement between $\{\mathbf{e}_{s}^{m}\}$ and generate independent part-level embeddings, an $M$-way one-hot attention vector is generated for each of the $D$ dimensions, by deploying that only one embedding can preserve its value and the others are set to $0$ at each dimension. This dimension attention is formulated as:

$\displaystyle\begin{aligned} &\hat{\mathbf{e}}_{s}^{m}=\mathbf{e}_{s}^{m}\odot\mathbf{M}_{[m,:]},\\ &\mathbf{M}_{[:,d]}={\rm hard\_softmax}([\mathbf{e}_{s}^{1}(d),\mathbf{e}_{s}^{2}(d),...,\mathbf{e}_{s}^{M}(d)]),\\ &{\rm hard\_softmax}(\mathbf{e})=\frac{\mathbf{e}}{\sum_{i}\mathbf{e}(i)}\odot~{}{\rm onehot}(\frac{\mathbf{e}}{\sum_{i}\mathbf{e}(i)}),\end{aligned}$

(6)

where $\mathbf{M}_{M\times D}$ is the attention matrix, whose $m$th row is $\mathbf{M}_{[m,:]}$ and $d$th column is $\mathbf{M}_{[:,d]}$, $\mathbf{e}_{s}^{m}(d)$ is the $d$th dimension of the embedding $\mathbf{e}_{s}^{m}$, ${\rm onehot}(\cdot)$ is the one-hot operation, and $\hat{\mathbf{e}}_{s}^{m}$ is the final part-level shape embedding. The same pipeline is applied to the albedo embeddings, where the only difference is that the attention $\mathbf{M}$ is copied from the shape embeddings, which ensures that the shape and the albedo information are decomposed synchronously.

By incorporating both shape and albedo information as cues, GDM successfully decomposes parts from objects under varied poses and appearances, ensuring that each part capsule represents a semantic-consistent part.

When training IGC-Net with unlabelled images, we employ the following constraints to learn the hierarchical 3D representations effectively:

The final loss functions to train IGC-Net are combined as:

$\displaystyle\begin{aligned} \mathcal{L}=&\mathcal{L}_{rec}+\lambda_{per}\mathcal{L}_{per}+\lambda_{contra}\mathcal{L}_{contra}\\ &+\lambda_{sparse}\mathcal{L}_{sparse}+\lambda_{bg}\mathcal{L}_{bg},\\ \end{aligned}$

(12)

where $\lambda_{contra}$, $\lambda_{sparse}$ and $\lambda_{bg}$ are the hyper-parameters to balance different loss functions.

Due to the 3D representations embedded in the graphics capsules, IGC-Net successfully builds the hierarchy of in-the-wild faces that are captured under varied illuminations and poses, shown in Figure 3. By incorporating shape and albedo information as cues, the face images are naturally decomposed into six semantic-consistent parts: background, eyes, mouth, forehead, nose, and cheek, without any human supervision. Each part is described with a specific graphics capsule, which is composed of a set of interpretable parameters including pose, view-independent shape, and view-independent albedo. These parts are assembled by their depth to generate the object capsules as the object-centered representations, building a bottom-up face hierarchy. We also try to discover other numbers of facial parts by controlling $M$ and get reasonable results, shown in the supplementary material.

To show the potential of IGC-Net, we extend our method to image collections of cat faces. Compared with human faces, cat faces are more challenging as cats have more varied textures than humans. The results are shown in Figure 4. It can be seen that the cats are split into background, eyes, ears, nose, forehead, and other skins.

The graphics capsules learned by IGC-Net provide a face hierarchy with explicit graphics descriptions, which gives a plausible way to materialize Marr’s theory [23, 22] that the purpose of vision is building hierarchical 3D representations of objects for recognition. In this section, we apply IGC-Net to validate the advantages of such hierarchical 3D descriptions and uncover the face recognition mechanism of neural networks.

To execute the quantitative and qualitative evaluation, we treat the silhouettes of parts as segment maps and apply them to the unsupervised face segmentation task. Note that there is no ground truth for the unsupervised part-level segmentation, the key of this task is to evaluate the semantic consistency of the parsing manners. The following experiments show the superiority of our method.

The basis of building the hierarchy of objects is to learn the parts with explicit semantics and keep semantic consistency across different samples. In this section, we perform the ablation study to show the importance of the one-hot operation in the GDM (see Eqn. 6) and the semantic constraint $\mathcal{L}_{contra}$ (see Eqn. 9) for discovering meaningful parts. Figure 9 shows the qualitative ablation study on CelebA. In the second row of Figure 9, it can be seen that, without the one-hot operation to prevent the information leakage of different parts, the semantic consistency across samples will be damaged. The third row of Figure 9 shows that the contrastive semantic constraint $\mathcal{L}_{contra}$ is important for the discovery of parts with salient semantics. Without such constraint, the segmentation of the important organs such as the eyes will have ambiguity. These conclusions are also validated by the quantitative ablation study shown in Table 3.