Neural Extensions: Training Neural Networks with Set Functions

Nikolaos Karalias, Joshua David Robinson, Andreas Loukas, Stefanie Jegelka

Keywords: continuous extensions, algorithmic reasoning, set functions, machine learning, combinatorial optimization, image classification

Abstract: Integrating discrete computational steps into deep learning architectures is an important consideration when learning to reason over discrete items. However, many tasks that involve discrete choices are defined via (combinatorial) set functions, and thereby pose challenges for end-to-end training. In this work, we explore a general framework to construct continuous extensions of such discrete functions that enables training via gradient methods. Our framework includes well-known extensions such as the Lovasz extension of submodular set functions and facilitates the design of novel continuous extensions based on problem-specific considerations, including constraints. We demonstrate the versatility of our framework on tasks ranging from combinatorial optimization to image classification.

One-sentence Summary: A principled framework for continuous extensions of set functions in machine learning.