Source-linked AI summary
Stacked Capsule Autoencoders
Adam R. Kosiorek, Sara Sabour, Yee Whye Teh, Geoffrey E. Hinton
TL;DR
SCAE addresses the challenge of representing objects through geometrically organized parts while handling viewpoint changes. It uses stacked part and object capsule autoencoders to reconstruct images and part poses from learned relationships. Object-capsule presence vectors support unsupervised classification, reaching 55% on SVHN and 98.7% on MNIST.
Problem
Conventional feature replication across non-translational viewpoint changes can create cumbersome, high-dimensional feature maps.
Method
SCAE uses a part capsule autoencoder followed by an object capsule autoencoder to infer part poses, organize parts into objects, and reconstruct them through learned geometric relationships.
Results
55% on SVHN and 98.7% on MNIST are achieved for unsupervised classification, improving to 67% and 99%, respectively, with fewer than 300 learned parameters.
Takeaways & Limitations
Object-capsule presence vectors form class-related clusters, enabling competitive unsupervised object classification without classification-focused training objectives.
Takeaways & Limitations
During inference, discovered objects are not used top-down to refine the presences or poses of the parts.
Abstract
from arXiv · showhide
Objects are composed of a set of geometrically organized parts. We introduce an unsupervised capsule autoencoder (SCAE), which explicitly uses geometric relationships between parts to reason about objects. Since these relationships do not depend on the viewpoint, our model is robust to viewpoint changes. SCAE consists of two stages. In the first stage, the model predicts presences and poses of part templates directly from the image and tries to reconstruct the image by appropriately arranging the templates. In the second stage, SCAE predicts parameters of a few object capsules, which are then used to reconstruct part poses. Inference in this model is amortized and performed by off-the-shelf neural encoders, unlike in previous capsule networks. We find that object capsule presences are highly informative of the object class, which leads to state-of-the-art results for unsupervised classification on SVHN (55%) and MNIST (98.7%). The code is available at https://github.com/google-research/google-research/tree/master/stacked_capsule_autoencoders
1 Introduction
SCAE addresses viewpoint-related limitations of conventional feature replication by learning object–part coordinate relationships and stacking part and object capsule autoencoders. Its object-capsule presence vectors cluster by class, supporting unsupervised classification.
- Viewpoint changes make replicating features across scale, orientation, and other affine degrees of freedom cumbersome and high-dimensional.
- Object–part coordinate transformations provide a viewpoint-invariant representation of object geometry, unlike object–viewer coordinates that change with viewpoint.
- SCAE first segments images into parts, infers their poses, and reconstructs images with affine-transformed templates before organizing parts into objects.
- Object capsules combine object–viewer relationships with object–part relationships to predict mixtures that reconstruct part poses.
- 55% on SVHN and 98.7% on MNIST are reported for unsupervised classification, improving to 67% and 99% with fewer than 300 learned parameters.
2 Stacked Capsule Autoencoders (SCAE)
SCAE learns objects as geometrically organized parts through two stacked autoencoding stages. It first discovers and reconstructs part templates, then organizes their poses into sparsely activated object capsules.
- Constellation Capsule Autoencoder: CCAE learns to assign two-dimensional input points to up to unknown numbers of transformed constellations without predefined shapes.It models point sets as arrangements of familiar constellations under independent similarity transforms.
- Object Capsule Autoencoder: The model trains without supervision by maximizing part-capsule likelihood, assigning parts through mixture responsibilities and learning sparse object capsules.Empirical results show unsupervised instance-level segmentation of points from different constellations, while OCAE sparsity requires additional regularization.
- Part Capsule Autoencoder: The PCAE infers part poses, presences, and features from images, then reconstructs images by arranging affine-transformed learned templates.Part templates include an alpha channel for occlusion, while special features can modify template appearance.
- Design constraints: The approach constrains each part capsule to one use per image and does not use discovered objects top-down to refine part presences or poses during inference.OCAE derivatives still refine the lower-level part encoder during training.
- Part Capsule Autoencoder: PCAE training learns stroke-like object-part templates, including MNIST strokes and double-stroke patterns in sobel-filtered SVHN.The PCAE is trained by maximizing the image likelihood.
- Object Capsule Autoencoder: The OCAE encodes part poses, features, templates, and presences into object capsules whose viewer relationships predict candidate part poses.Each object capsule has a feature vector, presence probability, and object-viewer relationship; candidate parts use object-part relationships and mixture predictions.
3 Evaluation
The evaluation tests geometric capsule representations on point constellations and images, finding strong unsupervised classification and viewpoint generalization, with component-specific ablation effects.
- 3.1 Discovering Constellations: 2.8% error was achieved by CCAE, compared with 26% for the baseline, on unsupervised segmentation of independently transformed point constellations.The authors attribute the order-of-magnitude reduction to an inductive bias for modelling geometric relationships in this toy setup.
- 3.2 Unsupervised Class Discovery in Images: Object-capsule presence vectors form tight clusters with large separation between different digits, supporting unsupervised class assignment in image data.The evaluation assigns class labels to presence vectors using clustering and linear prediction methods.
- 3.2 Unsupervised Class Discovery in Images: SCAE achieves state-of-the-art unsupervised classification on MNIST and SVHN but under-performs on CIFAR10 because it cannot model backgrounds.The image experiments use PCAE templates and OCAE object capsules, with results reported in Table 1.
- 3.3 Ablation study: 92.2 ± 0.59% accuracy on AFFNIST exceeds the 90.12% result of Sparse Unsupervised Capsules for viewpoint generalization.The model is trained on padded-and-translated 40 × 40 MNIST and evaluated on the AFFNIST test split.
- 3.3 Ablation study: Sparsity losses improve performance, part-capsule logit noise appears critical, affine decoding is less restrictive than similarity decoding, and Set Transformer object encoding is essential.Attention-based pooling has little effect within distribution but achieves much higher accuracy on a different dataset, indicating better generalization to novel viewpoints.
4 Related Work
Related work places SCAE among capsule, unsupervised-classification, and geometric-reasoning approaches. Its distinguishing features are unsupervised training, amortized inference, and explicit part–object geometry without transformation inputs.
- Capsule Networks: SCAE combines ideas from Transforming Autoencoders and EM Capsules while requiring only the image, unlike Transforming Autoencoders, which use an explicit transformation input.The transformation input is described as a form of supervision.
- Capsule Networks: Prior capsule networks use iterative routing from part votes to object poses, whereas SCAE uses objects to predict parts and avoids that routing procedure.The paper notes that iterative routing is inefficient and that SCAE’s OCAE encoder makes object predictions from all parts.
- Capsule Networks: SCAE performs unsupervised learning, whereas previous capsule networks generally used discriminative learning, and it further improves prior unsupervised AFFNIST generalization results.Rawlinson et al. are identified as a notable earlier exception using unsupervised capsule-conditioned reconstruction.
- Unsupervised Classification: Representation-learning approaches can discover class-informative latent variables but do not decompose objects into constituent parts or explicitly describe their shapes.The paper contrasts these approaches with SCAE’s templates and poses.
- Unsupervised Classification: Mutual-information methods directly optimize class-related probability representations, while geometric neural methods improve efficiency for rotations or reflections but not additional affine degrees of freedom such as scale.The paper distinguishes these approaches from SCAE’s explicit geometric part–object representation.
5 Discussion
SCAE is presented as a structured representation-learning method whose decoders train encoders to segment images into parts and poses and compose them into coherent wholes. The discussion notes competitive unsupervised classification without mutual information, weaker CIFAR10 performance, and possible extensions for more complex images.
- SCAE uses structured decoder networks to train separate encoders for part segmentation and whole-object composition.
- SCAE achieves competitive unsupervised object classification without mutual-information objectives or their sophisticated data augmentation.
- SCAE under-performs on CIFAR10, possibly because fixed templates are insufficiently expressive for real data.
- Deeper capsule-autoencoder hierarchies and input-dependent shape functions are proposed as possible ways to model more complex scenes.
- SCAE could serve as a figure component alongside a versatile ground component, with sequential attention analyzing one figure at a time.
A.1 Constellation Experiments
The constellation experiment uses a Set Transformer encoder for object capsules and separate capsule-specific neural networks to predict candidate part parameters.
- A four-layer Set Transformer encodes three object capsules as 32-dimensional vectors.Each layer has four attention heads with 128 hidden units per head and is followed by layer normalization.
- Separate neural networks for each object capsule predict four candidate part predictions per capsule, yielding 12 candidates overall.
A.2 Image Experiments
The image experiments use convolutional and Set Transformer encoders, MLP decoders, learned affine-transformed templates, and normalized Sobel-filtered reconstruction targets. Architecture, optimization, sparsity settings, and reconstruction examples are documented in the cited materials.
- Image experiments use a convolutional encoder for part capsules, a Set Transformer for object capsules, MLP decoders, and affine-transformed learned templates.
- Architecture details and loss-weight settings are provided in Tables 3 and 4.Table 4 distinguishes within and between quantifiers in the sparsity losses.
- The models use ReLU nonlinearities, sigmoid presence probabilities, 3 × 3 kernels, and 16 special features per part capsule in most experiments.
- Normalized Sobel-filtered images are used as SVHN and CIFAR10 reconstruction targets to emphasize shape importance.The procedure applies Sobel filtering, subtracts median color, takes absolute values, and normalizes values to [0, 1].
- Training uses RMSProp with momentum .9; batch size is 64 for constellations and 128 for other datasets.MNIST and constellation learning rates are 10^-5, while SVHN and CIFAR10 use searched rates from 5 × 10^-5 to 5 × 10^-4 with decay settings.
- Figure 7 compares Sobel-filtered targets with reconstructions from the Part Capsule Layer and object-predicted part poses.The learned templates converge to black-and-white patterns, so reconstructions lack color diversity.
C Constellation Capsule Sparsity
The constellation experiment adds a sparsity loss encouraging each active object capsule to explain at least two parts, using winning-part assignments to create binary capsule labels.
- The additional sparsity loss requires every active object capsule to explain at least two parts.
- An object capsule wins a part when it has the highest posterior mixing probability for that part.
- Binary labels mark capsules as 1 when they win at least two parts and 0 otherwise.
D Modelling Deformable Objects
The model represents object-part relationships with static mean-shape and data-dependent deformation components, while penalizing large deformations to encourage object capsules to specialize.
- D Modelling Deformable Objects: Object-part transformations combine a static mean-shape component with a dynamic component that depends on the input.The dynamic component is computed as OPdynamic_k,n = MLP(c_k) and models shape deformations.
- D Modelling Deformable Objects: Large deformations are discouraged so an object capsule learns a specific mean shape instead of modeling several objects at once.The deformation penalty is a weighted Frobenius norm added to the loss.
- D Modelling Deformable Objects: The deformation penalty uses weight α, typically set to 10 in the experiments.
E Part Capsule Encoder with Attention-based Pooling
The part-capsule encoder uses attention-based pooling to predict each part’s parameters by weighting spatial feature-map locations according to learned attention confidence.
- E Part Capsule Encoder with Attention-based Pooling: The encoder predicts a feature map and an attention mask for each part capsule, then pools spatial information using softmax attention.This mechanism supports predicting pose and existence when a part can occur at most once in an image.
- E Part Capsule Encoder with Attention-based Pooling: Each part capsule predicts six pose parameters, one presence parameter, and special features across the spatial feature map.The feature map has 6 pose + 1 presence + c_z special-feature dimensions.
- E Part Capsule Encoder with Attention-based Pooling: Softmax weights are applied along spatial dimensions to compute the final capsule parameters.