Source-linked AI summary
BAE-NET: Branched Autoencoder for Shape Co-Segmentation
Zhiqin Chen, Kangxue Yin, Matthew Fisher, Siddhartha Chaudhuri, Hao Zhang
TL;DR
Shape co-segmentation needs consistent structural explanations across collections, but conventional segmentation commonly relies on labels. BAE-NET addresses this with a branched autoencoder that reconstructs shapes through recurring part representations, and its one-shot variant outperforms supervised methods using far less annotated data.
Problem
Shape co-segmentation seeks consistent structural understanding and correspondence across collections, while many segmentation methods formulate the task as supervised classification requiring ground-truth labels.
Method
BAE-NET uses a CNN encoder and branched implicit decoder trained with shape reconstruction, with each branch representing a recurring part from point coordinates and shape codes.
Results
Using 1–3 segmented exemplars, BAE-NET outperforms state-of-the-art supervised methods trained on 77–230 shapes, while also producing unsupervised and weakly supervised results.
Takeaways & Limitations
BAE-NET provides consistent co-segmentation over large collections and adapts the same architecture across unsupervised, weakly supervised, and one-shot settings.
Takeaways & Limitations
The method assumes consistently aligned input shapes and provides no theoretical guarantee of segmentation consistency or universal part counts.
Abstract
from arXiv · showhide
We treat shape co-segmentation as a representation learning problem and introduce BAE-NET, a branched autoencoder network, for the task. The unsupervised BAE-NET is trained with a collection of un-segmented shapes, using a shape reconstruction loss, without any ground-truth labels. Specifically, the network takes an input shape and encodes it using a convolutional neural network, whereas the decoder concatenates the resulting feature code with a point coordinate and outputs a value indicating whether the point is inside/outside the shape. Importantly, the decoder is branched: each branch learns a compact representation for one commonly recurring part of the shape collection, e.g., airplane wings. By complementing the shape reconstruction loss with a label loss, BAE-NET is easily tuned for one-shot learning. We show unsupervised, weakly supervised, and one-shot learning results by BAE-NET, demonstrating that using only a couple of exemplars, our network can generally outperform state-of-the-art supervised methods trained on hundreds of segmented shapes. Code is available at https://github.com/czq142857/BAE-NET.
1. Introduction
BAE-NET treats shape co-segmentation as representation learning, using a branched autoencoder to discover recurring parts without ground-truth labels and adapt to one-shot supervision. Its evaluations cover unsupervised, weakly supervised, and one-shot settings, with one-shot results outperforming supervised methods trained on hundreds of shapes.
- Shape co-segmentation seeks consistent structural explanations across shape collections, supporting structural correspondence and applications such as attribute transfer and mix-and-match modeling.
- BAE-NET frames co-segmentation as learning simple recurring part representations with a branched autoencoder and shape reconstruction loss, without ground-truth segmentation labels.Each branch learns a representation for one universal part of the input collection.
- Max pooling groups branch outputs into the final implicit field and permits overlapping part representations.Branches may represent parts or output nothing when other branches cover them.
- The decoder concatenates each encoded shape feature with a point coordinate and predicts inside/outside values, while L3 neurons represent individual shape parts.Point coordinates provide spatial awareness that convolutional encoder features may lose.
- One-shot BAE-NET sets branch count from one or more exemplars and adds label reconstruction to shape reconstruction, outperforming supervised methods trained on 10%, 20%, or 30% of shapes.The reported comparison includes PointNet++ and PointCNN, with BAE-NET using only one to three segmented exemplars.
2. Related work
Prior shape co-segmentation methods use unsupervised, weakly supervised, semi-supervised, template-based, or handcrafted-graph strategies. BAE-NET instead provides an end-to-end differentiable, geometry-only encode-and-reconstruct pipeline whose branches learn recurring parts and can accommodate complex geometries.
- BAE-NET performs unsupervised co-segmentation of geometric shapes without colors, using branches to localize recurring parts across examples and reassemble the original shapes.The same framework adapts to weakly and semi-supervised scenarios.
- The method’s architectural prior relies on a shallow fully connected stack that forces compact representations of recurring parts in the shape space.This design is presented as analogous in spirit to Deep Image Prior and Deep Geometric Prior.
- Unlike prior graph-based approaches requiring manual features or large-scale graph optimization, BAE-NET is an end-to-end differentiable pipeline.
- Compared with cuboid-based primitive decomposition, BAE-NET’s nonlinear neuron stack can represent complex non-convex geometries and groups of disjoint parts.The comparison cites rings and disjoint lamp parts as examples that cuboidal boxes cannot accommodate.
- Unlike template-based and other supervised alternatives, BAE-NET does not depend on supervision or a base template, though annotated examples can separate strongly correlated part pairs.
3. BAE-NET: architecture, loss, and training
BAE-NET uses a CNN encoder and branched implicit decoder to reconstruct shapes while representing recurring parts in separate branches. Its losses support unsupervised, supervised, and one-shot training, with branch outputs used to assign part labels.
- Architecture: BAE-NET encodes a shape with a CNN, combines its feature code with point coordinates, and predicts inside/outside values through separate decoder branches.Each branch represents a recurring shape part, while max pooling merges branch fields into the reconstructed shape.
- Losses and training: Unsupervised training reconstructs shapes from sampled 3D points and inside-outside statuses using a mean square loss.The reconstruction loss uses the decoder output for a point against its ground-truth inside-outside value.
- Losses and training: Supervised training adds branch-specific losses for labeled parts and combines voxel-grid samples with ShapeNet part point clouds.Because surface point-cloud inside-outside status can be ambiguous, the method assumes those points lie inside thicker reconstructed shapes and sets α to 1.
- Losses and training: One-shot training alternates unsupervised updates on all shapes with supervised updates on a few labeled exemplars.The implementation performs one supervised iteration after every four unsupervised iterations rather than explicitly using the joint loss weight β.
- Training details: A small L1 regularizer on the final decoder layer discourages unnecessary overlap between parts represented by different branches.The decoder notation records layer widths, while task-specific experiments vary feature dimensions, decoder widths, sampling, and training iterations.
- Inference and labeling: After training, points receive the label of the branch producing the highest value, with exemplar guidance assigning branch labels automatically.In unsupervised settings, branch meanings must be assigned manually; a branch can, for example, be designated as “jet engine.”
4. Experiments and results
BAE-NET experiments evaluate architecture choices and segmentation across unsupervised, weakly supervised, and one-shot settings. The results show recurring part extraction, cross-shape correspondence, and strong one-shot comparisons, while also exposing sensitivity to exemplar selection and part coverage.
- 4.1. Network design choices and insights: The 3-layer network separates overlapping synthetic patterns, assigns each branch one pattern, and produces correspondence across shapes.The experiments use the synthetic “elements” and “triple rings” datasets with four branches.
- 4.1. Network design choices and insights: Unsupervised BAE-NET produces reasonable colored part fields and can align chair seats with table tops and chair legs with table legs in joint training.Unsupervised training may instead produce coarser segments, such as combining a chair’s seat and back.
- 4.3. Comparison with Tags2Parts: Weakly supervised training separates parts that remain merged under original unsupervised training, including plane engines, and yields eight effective plane branches.The branches represent recurring structures such as body, wings, engines, wheels, stabilizers, and related parts.
- 4.4. One-shot training vs. supervised methods: 73.5% IOU is achieved by the one-exemplar model, compared with 70% for a method using 5% of the training data.The one-exemplar result sets car IOUs to zero.
- 4.4. One-shot training vs. supervised methods: BAE-NET with one, two, or three exemplars outperforms supervised methods trained on 10%, 20%, or 30% of the data, respectively.The comparison is averaged over 15 shape categories.
- 4.4. One-shot training vs. supervised methods: One-shot performance varies substantially with exemplar choice, ranging from 8.9% to 72.3% for earphones because of shape variation and normalization misalignment.Other categories, such as laptops, show more consistent results around 94%.
- 4.4. One-shot training vs. supervised methods: One-shot segmentation also depends on exemplars covering all ground-truth parts and representing the collection, which fails across lamp types with different part semantics.For ground-lamp exemplars, the evaluation adds a zero-output fake branch for the missing ceiling-lamp part.
5. Conclusion, limitations, and future work
BAE-NET supports unsupervised, one-shot, and weakly supervised shape co-segmentation, outperforming supervised methods with far less training data. Its limitations include coarser segmentations, missing theoretical guarantees, alignment assumptions, and sensitivity to shape complexity and semantics.
- Limitations: BAE-NET tends to produce coarser but correct segmentations that can serve as starting points for further refinement.This coarseness distinguishes its outputs from those of supervised methods.
- Conclusion: The network handles collections of more than 5,000 shapes, whereas prior unsupervised co-segmentation methods were tested on fewer than 50.The reported results generally use default network settings.
- Limitations: The method has no theoretical guarantee of segmentation consistency or universal part counts and assumes consistently aligned input shapes.Rotation-invariant models are identified as future work.
- Limitations and future work: Unsupervised branch assignments can vary with random initialization, while similar nearby parts may be grouped incorrectly and high-resolution models remain difficult to train.These issues motivate future work on shape semantics, deeper representations, and coarse-to-fine hierarchies.