Timezone: »

Equivariance with Learned Canonicalization Functions
Sékou-Oumar Kaba · Arnab Kumar Mondal · Yan Zhang · Yoshua Bengio · Siamak Ravanbakhsh

Thu Jul 27 04:30 PM -- 06:00 PM (PDT) @ Exhibit Hall 1 #127
Symmetry-based neural networks often constrain the architecture in order to achieve invariance or equivariance to a group of transformations. In this paper, we propose an alternative that avoids this architectural constraint by learning to produce canonical representations of the data. These canonicalization functions can readily be plugged into non-equivariant backbone architectures. We offer explicit ways to implement them for some groups of interest. We show that this approach enjoys universality while providing interpretable insights. Our main hypothesis, supported by our empirical results, is that learning a small neural network to perform canonicalization is better than using predefined heuristics. Our experiments show that learning the canonicalization function is competitive with existing techniques for learning equivariant functions across many tasks, including image classification, $N$-body dynamics prediction, point cloud classification and part segmentation, while being faster across the board.

Author Information

Sékou-Oumar Kaba (Mila, McGill University)
Arnab Kumar Mondal (McGill University / Mila)
Yan Zhang (Samsung - SAIT AI Lab Montreal)
Yoshua Bengio (Mila - Quebec AI Institute)
Siamak Ravanbakhsh (McGill - Mila)

More from the Same Authors