DiagrammaticLearning: A Graphical Language for Compositional Training Regimes

TL;DR

This paper introduces learning diagrams, a graphical language for representing and composing complex training regimes in deep learning, enabling better understanding and manipulation of multi-component models.

Contribution

It presents a novel graphical language for modeling deep learning training setups, with a formal category theoretic semantics and practical implementation in PyTorch and Flux.jl.

Findings

Learning diagrams unify various training regimes like multi-task and multi-modal learning.

The library facilitates building and manipulating complex models as compositions of simpler components.

Experiments demonstrate improved understanding and flexibility in model training workflows.

Abstract

Motivated by deep learning regimes with multiple interacting yet distinct model components, we introduce learning diagrams, graphical depictions of training setups that capture parameterized learning as data rather than code. A learning diagram compiles to a unique loss function on which component models are trained. The result of training on this loss is a collection of models whose predictions ``agree" with one another. We show that a number of popular learning setups such as few-shot multi-task learning, knowledge distillation, and multi-modal learning can be depicted as learning diagrams. We further implement learning diagrams in a library that allows users to build diagrams of PyTorch and Flux.jl models. By implementing some classic machine learning use cases, we demonstrate how learning diagrams allow practitioners to build complicated models as compositions of smaller components, identify relationships between workflows, and manipulate models during or after training. Leveraging a category theoretic framework, we introduce a rigorous semantics for learning diagrams that puts such operations on a firm mathematical foundation.