Learning Diagrams: A Graphical Language for Compositional Training Regimes
Authors: Mason Lary
Richard Samuelson
Alexander Wilentz
Alina Zare
Matthew Klawonn
James Fairbanks
Venue: ICLR 2025 Poster (2025)
Keywords: ML Libraries
Training
Foundation Models
Multi-Task Learning
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We rigorously model training setups diagrammatically and provide a library for powerful manipulations thereof.

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.

Published as a conference paper at ICLR 2025
LEARNING DIAGRAMS:
A GRAPHICAL LANGUAGE
FOR COMPOSITIONAL TRAINING REGIMES
Mason Lary*1 University at Buffalo
Richard Samuelson" University of Florida
Alexander Wilentz
Harvard University
Alina Zare University of Florida
Matthew Klawonn Air Force Research Lab Information Directorate
James P. Fairbanks University of Florida
ABSTRACT
Motivated by deep learning regimes with multiple interacting yet distinct model components, we introduce learning diagrams, graphical depictions of training se- tups that capture parameterized learning as data rather than code. A learning dia- gram compiles to a unique loss function on which component models are trained. The result of training on this loss collection of models whose predictions "agree" with one another. We show that a number of popular learning setups such learning
as few-shot multi-task learning, knowledge distillation, and can be depicted as learning diagrams. We further implement learning diagrams in library that allows users to build diagrams of PyTooreh and Flux.) jI models. By im- plementing some classic machine learning use cases, we demonstrate how learn- ing 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 rigorous semanties for learning diagrams that puts such operations a firm mathematical foundation.
INTRODUCTION
The deep learning literature is rife with training regimes that exhibit non-trivial interactions between distinct models. Examples include molti- architectures with vision and language components (e.g) Vinyals et al. (2015), Ramesh et al. (2022), knowledge distillation schemes Hinton et al. (2015), multi-task learning setups, and on. The practitioner who manages such collections often juggles in- teracting models that at times are training or frozen, sometimes classifiers and sometimes feature extractors, sometimes trained on a single task and sometimes on many. Motivated by such practical settings we introduce formalism for building models that treats data sets, models, and their inter- actions as structured data rather than as code. Our primary contribution, the learning diagram, is graphical in nature and has rigorous mathematical semanties, meaning that learning diagrams be interpreted unambiguously. manipulated with confidence, and related to one another. Learning diagrams are also compositional, meaning that complex training setups can be built from easier to understand pieces.
At the heart of our formalism is the observation that machine learning problems can often be framed as search for commuting diagrams. In what is arguably the simplest case, we begin with a collec- tion of patterns T; TT and labels yı E. We select architecture f : I8 m 12n and learn a parameter 0 E R such that holds for all pairs (I, ys) E R x IR.
f 0) R y
(1)
If we model the patterns, labels, and parameters functions
"Equal contributions
1 "orresponning Authors: mason ar@buffal edu

---

Published as a conference paper at ICLR 2025
X:I-R", Y:I+R, 0:1 IR"
where is finite set that indexes patterns and labels, then the equality in means that the following diagram (approximately) commutes, i.e. that the pair of paths starting at 7 and ending at R produce
the same real valued output for any choice of element
T
R
X x
Xx0.f=Y
(3)
Anatomically, Diagram 3 is graph with nodes that represent spaces and edges that represent maps between them. The goal of parameterized learning is find l such that the diagram comes as close possible to commuting. We observe that the diagram itself can contribute to this goal if the common codomain of labels and function f, role played by R the diagram above, imbued with loss function. In that case and ground truth labels and train the component models to minimize said difference. As the loss minimized, the diagram gets closer commuting.
measure the difference between predictions of
The process of generating losses is not limited to diagrams with only one pair of paths that must commute. In fact, will see that many training regimes designed make multiple pairs of paths commute all at once. Further, the implementation of training setups as learning diagrams is an academic exercise. Rather, doing so allows for the construction and manipulation of training setups in way that would be more difficult without a structured representation. We believe such a shift from learning pipelines code learning pipelines as data needed address open problems identified by prior work, such as the need for "programming interface to specify that various adapted models are derived from the same pre-trained model" Boomasani et al.] (2021), or "tools that will allow us to build pre-trained models in the same way that we build open-source software" Raffel 202I). Our specific contributions are as follows.
We formalize the process of producing diagrams of models and data and compiling them to loss functions. We imbue learning diagrams with a rigorous semantics by way of category theoretic machinery.
We introduce software library, Diagrammaticl_aamee .eaning.l. that realizes the theory of learn- ing diagrams to supply convenient operations for building and manipulating training setups, both before models are trained and after, when they may be used as components of other training setups.
We provide several examples of common training paradigms that can be captured as learn- ing diagrams, and show how their implementation as such affords convenient functionality for common ML tasks.
The plan of the paper is to present these contributions reverse order. To motivate the imple- mentation we will first show in Section 2 how certain popular training setups can be realized using our approach. To do so, we have selected small sample of classic machine learning papers and re-implemented them using our framework, along the way recovering the results of the original pa- pers and demonstrating the utility of our diagrammatic formalism for building and manipulating models. In Section 3 we describe the beta version of our library, highlighting available features and the engines that make them work. Finally, in Section 4 we cover the mathematical underpin- ning learning diagrams that enables the operations realized in our library. We conclude with an examination of related work and directions for future research.
LEARNING DIAGRAM DEMONSTRATIONS
Our goal in this section is to show that our approach to developing training loops is expedient and feature rich without negatively affecting the final performance of the produced models. We do not
2

---

Published as a conference paper at ICLR 2025
aim outperform our baselines, but rather to recreate them. In doing, we will motivate the in- creasingly abstract constructions of later sections. To facilitate intuition before rigorous definitions, will leave some terms undefined for now.
We have identified training setups from classic papers that are both representative of widely used paradigms, and are sufficiently complex to highlight aspects of approach that practically useful. A particular paper's omission from our experiments does not necessarily indicate that we cannot model it as a learning diagram, nor that doing so would be a fruitless endeavor. Among the diagrams we have worked out but not included are norm-based regularization techniques and autoencoders.
2.1 LEARNING DIAGRAM NOTATION In each example that follows define learning di- agram first by considering domains of data corresponding to datasets, their components (such as images and labels), and the representation spaces induced by architectural choices. Each space has an associated Lawvere metric Lawvere (1973), i.e. is a generalization of metric space relaxing the symmetry and separation axioms, where distances between points may be infinite. For spaces where elements ought to be comparable and generate loss, we assign one and denote such a space Y as (Y, C), otherwise we assume that distinct points are infinitely far from one another and denote the space (Y,o).. The edges in a diagram that connect one domain to another depict models and data. We will see in Section 4 that such spaces and maps between them form an object of math- ematical significance called category. Our formal theory requires that parameter spaces also be represented explicitly, but for the sake of exposition we will omit or include them in our illustrations as convenient.
With this recipe for specifying learning diagrams in mind, we move on to our first example of learning diagrams in action: Neural Image Captioning Vinyals et al. (2015). It is the oldest of the group of examples and increases the complexity over standard predictions only slightly, but allows us introduce some useful features of approach and library.
2.2 IMAGE CAPTIONING Vinyals et al. (0155) was the first paper to demonstrate an end to end trainable model for image captioning. Named the "Neural Image Captioner" (NIC), uses a fixed visual encoder to extract features from an image, and a trainable recurrent component to produce caption. We can depict the model as in Diagram 4
(0) (CNNNY)
LSTM
(ZI, 20)
Prediction
(4)
Label
Cec)
In replicating the NIC diagrammatically, we aim to show that diagrams enable easy swapping of architectures for different components, and easy manipulation of model properties such as whether model updates is frozen. The components of the NIC diagram include graph that captures the "syntax" of our training set-up. as well as "semantic" data associated with the graph that includes the particular choice of architectures and spaces. Such separation can be exploited; whereas the original paper has a single choice of CNN encoder [Szegedy et al. (2015), we can easily vary the encoder backbone by simply assigning different architecture the CNN edge of the diagram.
In other words, there is no need to manually chain together the CNN and LSTM, as our backend handles it for us. The code required for such switch is depicted in Listing T and is only slightly edited from our implementation. We trained three semantic variations of Diagram 4 corresponding to three choices of CNN encoder; GoogLeNet (Szegedy al. (2013). Resnet-50 He al.) (2016), and an image transformer| Dosovitskiy (2020). VVnyaals et al. (2015) used Googl eeNt in the original paper. By using more powerful architectures and replicating other portions of the model and training procedures, were able to outperform the original paper's BLEU score Papineni al. (2002) on the Flickr8k Hodosh et al. (2013) and Flickr30k Plummer et al. 2015) image captioning data sets. The results of our experiments
captured in Table 1
The mix of trainable and fixed models offers a chance to highlight some model manipulation. In particular,
utilize model homomorphisms to identify sub-components and apply functions to them. To specify a homomorphism. we first specify the shape of sub-diagram. then describe mapping from the sub-diagram to the larger model with components want to manipulate. In
3

---

Published as a conference paper at ICLR 2025
Data Set
GeogLeNet
Resner-50
ViT-B-16
Original
Flickr8k BLEU 66.85 70.06
66.93 76.72
71.39
63
Flickr30k BLEU
71.78
66
Table 1: Results of image captioning models trained as learning diagrams, named after the convolu- tional component.
Listing we show code that allows us to set the vision encoder to "eval" mode. While this is easy enough to do with standard PyTorrch for such simple model, for more complex diagrams with more models becomes especially advantageous to leverage homomorphisms which can pick out subgraphs based on some structural pattern matching and apply different transformations to the different pieces of the pattern Ehrig et al. (1973).
2.3 KNOWLEDGE DISTILLATION Much like the simple Diagram the image captioning case contains only one pair of paths that induce loss; we want the captions produced by the CNN+LSTM team to be the same as the ground truth. Learning diagrams are capable of handling much more complicated setups where multiple pairs of parallel paths contribute terms to the loss. To demonstrate this capability, we now describe how to implement knowledge distillation learning diagram.
Knowledge distillation can refer to many different techniques (Buciluă et al. (2006); Hinton et al. 2013)) that share the common goal of taking information from a powerful model or ensemble called "teacher" and recovering that information in a simpler, easier to use model called the "student." A particularly popular approach to distilling knowledge is to train the student on the output of the teacher, whether such output be logits or a softmax with a temperature. Further, the student trained on labeled data that can be the same as that used in training the teacher, and so receives supervisory signals from two sources; the outputs of the teacher model and labeled data.
Figure recovers a knowledge distillation setup; notice there two pairs of parallel paths. One terminates at Y. the label space, and the other in G. the soft target space. The learning diagram therefore induces loss with two terms, relative to the ground truth labels, and the other corresponding to the softmax outputs of the stu- dent model relative the soft outputs of the teacher. The official PyTorch knowledge distillation introduction [Chariton (2024) has a more detailed specification of the experiments than the original paper Hinton al. (2015), 5), so we compare those benchmarks. We see in Table 2 diagram produces results comparable to those of the official implementation.
corresponding the predictions of the student model
that the learning
v. v_model load (cnn_ _ame) train_de ds = get_dataset ()
diagram @LD begin
N + [train_ds]
1_1_x_y x_Y = (Any, Vector (FLoatt614)
1_1. + Any
+ "Vector(Float66
cross entropy) x
CNN x.
(v (v_model (x(0)), x(1))
LSTM: Q + L = (x, y) -> (x)
Label: f.Ixx_ - -Y Y (x, y) ->
Prediction: I + x -> predict
CNN_x_Y LSTM Prediction YNN_X_Y Label
end
Listing 1: Key functionality for the NIC implementation. Users specify diagrams as Python or Julia data structures (Julia shown) similar an edge list representation of the underlying graph and assign models to components of the diagram. Because the parts of the model are represented as data instead of code, users can specify and manipulate sub-diagrams and components programmatically. without the complexity of manipulating unstructured code. We also demonstrate that mapping to data spaces provides type annotations different levels of granularity for models the diagram.
4

---

Published as a conference paper at ICLR 2025
(T,
8:. n1 t
(N, (X, 00)
(S, 00) (G,Lat)
8 72
Con)
A diagrammatic representation of Knowledge Distillation (left), the red and blue paral- lel paths encode equations that the optimization procedure attempts to solve (right). Components models used in both pairs of parallel paths shown in violet.
Figure 1:
Teacher
Student
Distilled Student
Py Torch Implementation 74.72
70.50 67.67
70.81
Learning Diagram
76.24
69.21
Table 2: Performance comparison of official Py Torch knowledge distillation vs DiagrammaticLeam- ing.jl CIFAR-10.
2.4 FEW SHOT LEARNING We now turn to example that will highlight the utility of building more complicated models in a compositional way. The multi-task and few-shot learning literature contains many examples that would benefit from compositional perspective, thanks to their perva- sive use of shared backbones conjunction with task specific modules. We select Tian et al. (2020) our replication target demonstrate the utility of compositional perspective, since it exemplifies the idea of training a feature extractor on as much available data as one can, and subsequently train- ing small task specific modules that exploit the learned features. More convoluted training pipelines have had limited success in improving upon this approach. In the case Tian et al. (2020). the data used comes from Tiered Imagenet (Ren al.) (018)) Mini-Imagenet Vinyals et al. (2016), FC100 Oreshkin et al. (2018), and CifarFS Bertinetto et al. (2018). A ResNet-12 backhone trained on all data, and subsequenty few shot tasks are sampled at random from the respective test sets. Simple affine classifiers are then trained on the feature-label pairs extracted by the pre-trained backbone.
The meta-train and meta-test steps of Tian et al. 2020) have very similar syntax, in that they both extract features and train a classifier to make predictions using those features. The difference is that during meta-training the feature extractor is updated on all tasks. Intuitively, we can track the task specific interactions of data and models and use them to define the composite interaction seen during mnea-training, then we can easily re-use them during meta-testing and avoid implementing new meta-testing loop from scratch. Learning diagrams provide such functionality by allowing us to define task specific classification setups, and subsequently glue those setups along the shared axis for composition, colimit, is depicted and
of common encoder. The exact construction described in Figure
Approach
MI 1-Shot
MI 5-Shot
TI 1-Shot
TI 5-Shot
Tian et al. (200)) Simple
62.02
79.64
69.74
84.41
Learning Diagram
58.13
78.82
66.48
80.15
Approach
CFS 1-Shot
CFS 5-Shot
FC100 1-Shot
FC100 5-Shot
Tian et al. (2020) Simple
71.5
86.0
42.6
59.1
Learning Diagram
70.5
85.6
43.1
59.3
Table 3:
Performance comparison of Tian et al. 20020) vs Diagrammatill.amiingg. on Mini-
Imagenet (MI) Tiered-Imagenet (TI), Cifar-FS (CFS), and FC100.
The compositional perspective not only allows us to build large training pipelines quickly, but also to track the presence notable sub-diagrams via the inclusion morphisms computed as part of the colimit computation. As result, upon training the foundation model, we can efficiently proceed to the meta-test phase. Using the same "edata" syntax as in Listing [ we first assign a new classifier head and N-way K-shot data set to each of the classifier sub-diagrams, the encoder to "eval
5

---

Published as a conference paper at ICLR 2025
N
X
X
N
X
colim
7
m
Y
Y
F
h
Ni
N2
N
N2
21
22.2
81
22
X
X
X
v1
colim
Vi
2/2
Y2
F
F
F
Figure 2: Learning diagrams can be built using the categorical construction of colimit, which generalizes the concept of set union quotienting by equivalence relation. (Top) A data set and classifier/encoder pair are composed by identifying the common image and label domains and (dashed arrows). The result of the colimit computation is square whose induced loss trains the classifier and encoder on the labeling task of the associated data set. (Bottom) Two classifier squares are merged along a common feature extraction backbones m. The resulting colimit is a learning diagram whose loss trains classifier heads on their respective tasks and the common backbone model on all tasks. For this colimit be well specified, both the image space X and the feature space F have to be identical. Additional classification heads can be attached by further colimits.
mode", build the model, and then train on a fewshot data set. Repeating this procedure per randomly sampled few-shot data set gives the performance noted in Table
IMPLEMENTATION
Having motivated the functionality DiagrammatieLeammingn/ Wiith with some common ML tasks, we now turn to describing the implementation that provides such functionality.
3.1 SPECIFYING THE DIAGRAM DATA To generate composite loss, users must specify dia- gram. Listing shows an example of the structures required by the library. The required components
Vertices with their labels, mapping to data space, and an optional distance metric. Data spaces are captured by Julia types, which allows the user to annotate their diagram with type information for computations. This can be circumvented by supplying Any type within the diagram.
Edges with their labels, their source and target vertices, and mapping the computation the edge performs. These computations are simply Python Julia functions.
Relationships between edges, represented with S. This requirement specifies preorder on the edges of graph, concept specified further in Section 4 It can also be seen specification of which edges in the graph should take part in loss calculations.
3.2 FINDING PARALLEL PATHS From a learning diagram, we aim to calculate single loss all the parallel paths, first need compute all
function. Since the loss function pairs of parallel paths from the graph.
To accomplish this, we apply an algebraic perspective and use the fact that powers of the adjacency matrix count paths in a graph. This approach is a generalization of the celebrated Floyd- Warshal algorithm Floyy 1962). We represent the graph as V V matrix with entries that are of edges
https:: //gi thub. com www..ooow.lllaai.oo cJul a wwwp:///aaaaiaaaaaaaaaaaa agrammatt cliear ng-

---

Published as a conference paper at ICLR 2025
in the graph. This allows to define a generalized version of matrix multiplication which computes paths between vertices Fong and Spivak (2018); Pratt (1989). Multiplication of elements represents concatenation of paths and summation represents union of sets, which leads to algorithm that computes all the paths in a graph. Powers of this set-of-paths valued matrix then enumerate the paths between each pair of vertices of a fixed length. Parallel paths can then be collected from any matrix element with more than 1 element.
Many common machine learning loss functions are not symmetric. These loss functions require parallel paths to occur in specified order. To facilitate this, we assume that all pairs of single edges are ordered (i.e. e1 S e2 and eg S e₁ for edges e1, e2), unless only a single pair is specified in the graph. The process that builds the matrix containing parallel paths additionally builds up a preorder between paths, allowing a filter to include only parallel paths that conform to the preorder.
3.3 INTEGRATION WITH EXISTING LIBRARIES Due to the generic nature of our representation, we can view a learning diagram as a sort of intermediate language, independent of a particular neural network framework. To this end, we compile either PyTorch Flux models from a given diagram. This flexibility allows for the creation of diagrams for new workflows and experiments, as well as the retrofitting of diagrams existing codebases. This paper focuses on PyTorel models due to its ubiquity in machine learning research and practice.
The correctness of the compilation process is explored rigorously in the next section, but for more familiar explanation, we refer to PyTorch components. Vertices in a graph in our framework contain both data sources (such as a PyTorch dataloader) as well as a loss function. A path in our graph can quite easily be converted into Py PyTorch Sequential model. Given parallel paths, a module can be created which samples x from the data source present at the domain of the paths, passes x through both sequential models to obtain outputs y, U. and outputs [(y, ú), the result of applying the loss function / found in the codomain of the paths. Summation of these modules represents composite loss functions. This process thus encapsulates data sampling, network computation, and loss computation into a single PyTorch module.
The requirements for compatibility with our framework are similar to those for compatibility with pure PyTorch. For data, we require iterable components, whether it be DataLoaders or just lists. For computations, require functions, either in the form of torch modules, or pure Python functions. Thus, popular frameworks within the Py PyTorc landscape, such as models from HuggingFace, are compatible with our system.
4 THEORETICAL GROUNDING
Having discussed some applications and our implementation, we now turn to the underlying category theoretic formalisms and definitions. It is the content of this section that ensures learning diagrams have rigorous semantics, and therefore that manipulations of diagrams and maps between them can be interpreted mathematically. The ultimate goal of the section is a machinery that "mathematically compiles" a learning diagram into composite loss function that drives participating models to form approximately commuting diagrams. Throughout this section we will use some category theory that we may not define due to space constraints. Everything left undefined should be easy to find in introductory references Riehl (2017): Fong and Spivak (20T8). Our story proceeds by defining what learning diagram is, explaining how parallel paths can be extracted from it, and finally how these parallel paths combine to form a composite loss function.
4.1 LEARNING GRAPHS AND LEARNING DIAGRAMS A learning graph (G, S) is directed multigraph G = (V, E, src, tgt) together with a preordering < of its edges.
We will motivate the preorder in due time. On its own, a learning graph doesn't associate loss function to each node of the graph. Rather, we consider it a syntactic presentation of training regime. As we saw in the image captioning example, this distinction can be useful when varying the exact data or models being used. In the end, however, what we want is a full learning diagram. The choice of the word diagram no accident; it has proper category theoretic definition.
Definition 4.1. (1.6.4 of Richl (2017)) A diagram in a category C is a functor D: + C whose domain, the indexing category, is small category.
7

---

Published as a conference paper at ICLR 2025
To define diagram, therefore, we must first specify domain category and codomain category C. Our domain category comes from our learning graph and identifies the syntactic structure of our training regime: how many spaces of data representations etc. The codomain category defines what kind of thing each node and edge in our learning graph is, i.e. the semantics of the training setup, and the functor D is the mapping from syntactic structure to semantic content that assigns particular spaces and models to edges. Constructing a category to serve as our domain is fairly straightforward; we can take a learning graph and turn it into a category by taking vertices of the graph to be objects of a category and paths in graph to be morphisms. Such a construction is called the free category on a graph. Defining the semanties conferring codomain requires more slightly more work. Given that we want to associate a potentially asymmetric loss function to each node learning graph, we should consider category C whose objects are some generalization of metric spaces. We choose the category Law.
considering, how many models,
Definition 4.2. Denote by Law the category that has objects Lawvere Metric Spaces Lawvere 1973) and morphisms Lipschitz continuous functions.
4.2 FINDING PATHS AND BUILDING LOSSES Let Par be the free category on the following graph.
(5)
We the category Par (5) to formalize the construction of losses from parallel paths. If (G, <) is a learning graph, and if Free(G) is the free category generated by G, then a functor P Par + Free(G) chooses a pair of paths in G with the same initial and terminal vertex. If D: Fr) + Law is learning diagram on G, then the composite functor D o P Par - Law chooses a pair of Lipschitz continuous functions f X - Y and g X Y with the same domain and codomain. We will assign loss to these functions as follows. If X is finite set, then we define ((f.g) E [0, + to be the sum
((f.g) = Σ [d((f()),s(z))
ZEX
If X is an infinite set, then we take the supremum over all finite sums. This quantity measures the distance between outputs of ff and g summed over all inputs. It also has the following contractive property: if f X + Y, g
X + Y. X' - Y', and g' : X' + y' are Lipschitz continuous functions, and if p: X + X" and q: Y + Y" surjective I-Lipschitz functions such that
f'((x)) = (x))) and g' (((()))) = ((g(x)))
for all æ E X, then ((f, () 2 ((f','') Formally, loss is enriched functor é C C Cost, where C the subcategory of the functor category Law Par whose fibers are surjective I-Liischitzz functions, and where Cost is the partially ordered set co).
Theorem 4.3. The mapping P + E(f.g) g) defines a a.functor { C + Cost.
Proof. Let f. 9. f', g', p and q be the functions described above. Functoriality is the contractive property, {(f.g) [( ,g'). To see this, let C X' be finite set. Then,
Σ (f()) s Σ dy (f((()))) Z'ES
zEp-is
Σ
(r
zEP-1S
Σ
dv())x),g(z)
zEP-
where the first line follows from the surjectivity of p and the third from the contractivity of q. The desired inequality follows from the arbitrary selection of S.
8

---

Published as a conference paper at ICLR 2025
4.3 LIMITING THE SEARCH Diagrams in Law with shape given by learning graph very nearly capture the data we require to build composite loss, but for two problems. For one, if we allowing asymmetric loss functions. then must provide the creator of learning diagram with a way to say how arguments to a loss should be ordered. Additionally, some vertices in the learning graph will be assigned finite collections for indexing data sets and other vertices will be assigned vector spaces for representing the spaces containing the data values. The loss functions should only use parallel paths that start indexing, because these lead finite sums in the loss functions.
Ordering paths The preorder on edges in learning graph naturally leads to preorder structure on the free category associated with a graph. This produces locally posetal 2-category, which is category with preorder every hom-set, such that morphism composition is monotonic.
The upshot is that our syntactic 2-category now tracks user specified ordering models, i.e. which model's predictions should be the first argument loss and which should be the second.
Indexing and Nonindexing Vertices The other problem is that we only want to build composite losses for parallel paths that start specific nodes. The machinery of categories provides succicint way to specify the relevant constraints. Let Ind be the free 2-category generated by the learning graph
By setting - S +, we may upgrade the category Par defined in section 2 to 2-category. These categories are related by a 2-functor a : Par + Ind.
We use the 2-category Ind to assign labels to each vertex of learning graph. Specifically, for all learning graphs G, a functor 3: Free(G) - Ind is labelling the vertices G as either indexing or non-indexing, and the structure of Ind ensures that indexed vertices precede non-indexed ones. A commutative diagram of the form chooses a pair parallel paths that begin at indexed nodes.
Par
Free(G)
(8)
Ind
Given learning diagram D Free(G) + Law, the composite functor D P Par Law chooses pair f X + Y and g X + Y of Lipschitz continuous functions, and we can compute the loss of these functions as in section 2 (We define the composite loss triple (G, D2 to be the sum
e D, E [0, +oo given by
= Σ E(f.9)
where and are the functions corresponding the functor D P. and where the taken over all choices of P making the the diagram in g commute. This construction allows us to build losses automatically from graphical specifications of machine learning architectures. The contrac- tive nature of loss functions ensures that composite losses are well behaved with respected to taking subsets of each dataset in the diagram.
RELATED WORK
We find ourselves aligned in vision with the field of AutoML in that we seek to make the life of the machine learning practitioner casier. More specifically, fit with work that focuses on structuring machine learning pipelines; a pipeline focused AutoML survey can be found Zöller and Huber 2021). The formulation of Zöller and Huber 2021) describes pipelines Directed Acyelic Graphs

---

Published as a conference paper at ICLR 2025
(DAGs). We differ from this formulation and other DAG-based formulations in many important ways. For one, we have a mathematically grounded definition of what data is captured by nodes and paths in our learning diagrams, whereas prior art appeals to imprecise terms such as "primitives" Lippmann et al. (0016). "algorithms", "operators". "transformations" and so on. There are some works (e.g. Drori al. (2019) that place more structure pipelines, but in this framework the grammar is still defined on an unstructured set of ad-hoc operations. Further, works Kalyuzhnaya al. (2020); Nikitin et al. (2022); Polonskaia et al. (202I) that define somewhat flexible pipelines usually limit themselves to DAGS that have a single source and sink representing the data and label space. As best we can tell, there appear to be a few exceptions to this rule, e.g. Klein et al. (2017) which is limited to Bayesian settings with a common data domain shared among tasks. Thanks to the rigor of our model we can accommodate arbitrarily many tasks interacting in arbitrary ways, so long as the nodes are Lawvere metric spaces and the paths between nodes are continuous.
Learning diagrams also provide a graphical programming language for constructing machine learn- ing models, and in that sense are related to works that study foundations of languages for machine learning models Elliott (2018), though such techniques tend to focus on probabilistic program- ming Gordon al.) (2014) or the semantics of gradient descent and less on the semanties of the objective to be optimized.
Category theory has already been used examine machine learning models compositionally, for ample in Fong et al. (2019); Cruttwell et al. 2002]; Gavranovié (002). Our approach is somewhat different from most such works in that focus on the semantics of the loss or objective function as opposed to the semantics of model updates. That said, we believe there are connections to ex- plore, particularly with Hanks et al. (2024) as it also considers how specify compositionally loss for machine learning models. More generally, our work belongs to an emerging tradition of using category theory to provide a diagrammatic yet rigorous approach describing models in domain of interest Decapodes.jl |Mori et al. (2024); Patterson et al. (2023): Libkind et al. (2022)
and A AlgebraicPetris Libkind et al. (2022).
FUTURE WORK AND CONCLUSION
We introduce DiagrammaticlLcamingg jl. a library built on the notion of learning graphs and learning diagrams allowing machine learning practitioners to specify complex interactions between multi- ple models in a way that is inherently compositional and structured. The utility of this approach is established by capturing some popular training setups as learning diagrams and demonstrating convenient functionality that came as a result. We further outline the implementation and theoreti- cal backing that make learning diagrams possible, demonstrating rich and rigorous semantics that leads naturally to principled operations for manipulating structured collections of models.
This paper has only scratched the surface of potential capabilities. Future work should investigate how model structure can be exploited to improve distributed training and manage high performance systems. Additionally, our structured representation seems like candidate for the foundation of a more complete machine learning model management system. We hope for example to elevate software centric definitions of reproducibility. such as the one given in He et al. (2021), ematical plane that is "implementation free" in the sense that any implementation faithfully repro- ducing the math will yield equivalent results. Homomorphisms between diagrams could capture re- lationships between development iterations and facilitate operations like searching for models with particular properties or component sub-diagrams. On the more theoretical side, our formalism is currently limited to building losses for finite data sets, but extensions to methods where continuous distributions can enter the diagrams are possible. Additionally, our implementation of training net- works of models uses mini-batches, but the theory does not address such aspects of practical training algorithms.
math-
Development and adoption of these tools would reduce the labor necessary to produce effective ML pipelines, accelerate iteration of ideas within and between research groups, reduce errors intro- duced in software implementations, and provide quicker turnarounds while evaluating novel ideas. We hope the learning diagrams paradigm will accelerate the overall process of machine learning research.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Luca Bertinetto, Joao F Henriques, Philip HS Torr, and Andrea Vedaldi. Meta-learning with differ- entiable closed-form solvers. arXiv preprint
2018.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, al. On the opportu- nities and risks of foundation models. arXiv preprint arXi::2108.0723 2021.
Cristian Buciluă, Rich Caruana, and Alexandru Niculescu-Mizil. Model compression. In Pro- ceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 535-541 2006.
Alexandros Chariton. Knowledge distillation tutorial - PyTorch tutorials 2.4.0+cu121 docu- mentation, 2024.
distillation tutorial.htm.
Geoffrey SH Cruttwell, Bruno Gavranović, Neil Ghani, Paul Wilson, and Fabio Zanasi. Categorical foundations of gradient-based learning. In European Symposium on Programming, pages 1-28. Springer International Publishing Cham, 2022.
Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1100000..1199... 1929, 2020.
Iddo Drori, Yamuna Krishnamurthy, Raoni Lourenco, Remi Rampin, Kyunghyun Cho. Claudio Silva, and Juliana Freire. Automatic machine learning by pipeline synthesis using model-based reinforcement learning and a grammar. arXiv preprint arXiv: 1905. 2019.
H. Ehrig, M. Pfender, and H. J. Schneider. Graph-grammars: An algebraic approach. In 14th Anmual Symposium on Switching and Automata Theory (swat 1973), pages 167-180, 1973. doi: 10.1109/SWA..193.111
Conal Elliott. The simple essence of automatic differentiation. Proceedings of the ACM on Pro- gramming Languages, 2(CCFP):1-29, 2018
Robert W Floyd. Algorithm 97: shortest path. Communications of the ACM, 5(6):345-33) 345, 1962.
Brendan Fong and David I. Spivak. Seven Sketches in Compositionality: An Invitation to Ap-
plied Category Theory, October 2018. URL
arXiv:1803.05316 [math].
Brendan Fong, David Spivak, and Rémy Tuyéras. Backprop as functor: A compositional perspec- tive on supervised learning. In 2019 34th Annual ACM/IEEE Symposium on Logic Computer Science (LICS), pages 1-13. IEEE, 2019.
Bruno Gavranović. Learning Functors using Gradient Descent. 323:230 245, 2022. ISSN 2075- 2180. doi: 10.4204/EPTTCS.233 15. URL httt::://axxi..ogg/bb//////////s 068000990683777
Andrew D Gordon, Thomas A Henzinger, Aditya V Nori, and Sriram K Rajamani. Probabilistic programming. Future of software engineering proceedings, pages 167-181. 2014.
Tyler Hanks, Matthew Klawonn, and James Fairbanks. Generalized gradient descent is hypergraph functor. arXiv preprint arXiv:2403.. 19845, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
Xin He, Kaiyong Zhao, and Xiaowen Chu. Automl: A survey of the state-of-the-art. Knowledge- based systems, 212:106622, 2021.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge a Neural Network, March 2015.
02531] arXiv:1503.0253 [cs, stat].

---

Published as a conference paper at ICLR 2025
Micah Hodosh, Peter Young, and Julia Hockenmaier. Framing image description ranking task: Data, models and evaluation metrics. Journal of Artificial Intelligence Research, 47:853-899 899, 2013.
Anna V Kalyuzhnaya, Nikolay O Nikitin, Pavel Vychuzhanin, Alexander Hvatov, and Alexander Boukhanovsky. Automatic evolutionary learning of composite models with knowledge enrich- ment. In Proceedings of the 2020 Genetic and Evolutionary Computation Conference Compan- ion, pages 43 43,44,77 2020.
Aaron Klein, Stefan Falkner, Numair Mansur, and Frank Hutter. Robo: A flexible and robust bayesian optimization framework in python. In NIPS 2017 Bayesian optimization workshop, pages 9-9. NIPS, 2017.
F William Lawvere. Metric spaces, generalized logic, and closed categories. Rendiconti del semi- nario matématico e fisico di Milano, 43:135-166, 1973.
Sophie Libkind, Andrew Baas, Micah Halter, Evan Patterson, and James P Fairbanks. An algebraic framework for structured epidemic modelling. Philosophical Transactions of the Royal Society A, 1.3802233))202103399 2022.
Richard Lippmann, William Campbell, and Joseph Campbell. An overview of the darpa data driven discovery models (d3m) program. In NIPS Workshop on Artificial Intelligence for Data Sci- ence, pages 1-2, 2016.
Luke Morris, Andrew Baas, Jesus Arias, Maia Gatlin, Evan Patterson, and James P Fairbanks. De- capodes: A diagrammatic tool for representing, composing, and computing spatialized partial differential equations. Journal of Computational Science, page 102345, 2024.
Nikolay O Nikitin, Pavel Vychuzhanin, Mikhail Sarafanov, Iana S Polonskaia, Ilia Revin, Irina V Barabanova, Gleb Maximov, Anna V Kalyuzhnaya, and Alexander Boukhanovsky. Automated evolutionary approach for the design composite machine learning pipelines. Future Generation Computer Systems, 127:109-122 2022.
Boris Oreshkin, Pau Rodríguez López, and Alexandre Lacoste. Tadam: Task dependent adaptive metric for improved few-shot learning. Advances in neural information processing systems, 31, 2018.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei jing Zhu. Bleu: method for automatic evaluation of machine translation. pages 311-318, 2002.
Evan Patterson, Andrew Baas, Timothy Hosgood, and James Fairbanks. diagrammatic view of differential equations in physics. Mathematics in Engineering, 5(2):1-59, 2023. ISSN 2640- 3501. doi: 2..3934/mine,2023000 2023036. URL/http: www.:w/wwmmpr.ss...o.m ccm. artpaa//ccccco////// 10. 39.3333/monn /inee//ine.../ 2023036
Bryan A Plummer, Liwei Wang, Chris M Cervantes, Juan C Caicedo, Julia Hockenmaier, and Svet- lana Lazebnik. Flickr30k entities: Collecting egioonto-phrase correspondences for richer image- to-sentence models. In Proceedings of the IEEE international conference on computer vision, pages 2641-2649, 2015.
Iana S. Polonskaia, Nikolay O. Nikitin, Ilia Revin, Pavel Vychuzhanin, and Anna V. Kalyuzhnaya. Multi-objective evolutionary design of composite data-driven models. In 2021 IEEE Congress Evolutionary Computation (CEC), pages 926 9369933, 2021. doi: 10. 10.1109/CEC48833.2222..
9504773.
Vaughan R. Pratt. Enriched categories and the warsh-warshall connection. pages 177-180, 1989.
Colin Raffel. A call to build models like we build open-source software, 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint 0666:220000 1(2):3, 2022.
12

---

Published as a conference paper at ICLR 2025
Mengye Ren, Eleni Triantafillou, Sachin Ravi, Jake Snell, Kevin Swersky, Joshua B Tenenbaum, Hugo Larochelle, and Richard S Zemel. Meta-leamnng for semi remi-supervised few-shot classifica- tion. arXiv preprint arXiv: (ath::883300077 2018.
Emily Riehl. Category theory in context. Courier Dover Publications, 2017.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1-9, 2015.
Yonglong Tian, Yue Wang. Dilip Krishnan, Joshua B Tenenbaum, and Phillip Isola. Rethinking few-shot image classification: good embedding is all you need? In Computer Vision ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XIV 16, pages 266 282. Springer, 2020.
Oriol Vinyals, Alexander Toshev. Samy Bengio, and Dumitru Erhan. Show and tell: A neural image caption generator. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 3156-3164, 2015.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advances in neural information processing systems, 29, 2016.
Marc-André Zöller and Marco F Huber. Benchmark and survey of automated machine learning frameworks. Journal of artificial intelligence research, 70:409-472, 2021.
13


Reviews:
[{"review_id": "8TcOshxr0T", "reviewer": "Reviewer_u1cH", "summary": "This work formalizes machine learning training using learning diagrams.\nA learning diagram is defined as a directed graph that allows the composition of learning components, models, and loss functions.\nThe work uses a colimit construction for building training pipelines as diagrams which consist of sub-diagrams. Finally, a softwware library, DiagrammaticLearning.jl, implements the provided formulation.", "questions": "What capability does this formalism provide that is unavailable in existing AutoML systems?", "limitations": "", "rating": 3, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This work presents an elegant mathematical notation and formulation of automated machine learning (AutoML) based on directed graphs consisting of machine learning primitives and losses. The work formalizes the creation of machine learning pipelines.\n\n2. A software library is presented that implements this formulation.", "weaknesses": "1. It is unclear what advantage this framework provides over previous work on automated machine learning, specifically AutoML systems based on directed graphs of machine learning components, or machine learning pipelines.\n\n2. The empirical performance improvement over trivial baselines, as described in Tables 2 and 3, is not significant: 1-3% compared to a weak baseline.\n\n3. The work lacks novelty with respect to the broad field of AutoML, which consists of machine learning primitives, automated machine learning systems, and user interfaces, developed over the past decade:\n\na. See for example derivatives of the D3M project.\n\nb. A book on AutoML: Automated Machine Learning, Hutter, Kotthoff, and Vanschoren, Springer Nature, 2019\n\nc. AutoML systems:\nAuto-sklearn, Google Cloud AutoML, H2O AutoML, AlphaD3M, Auto-Keras, AutoGluon, to name a few,\n\nd. Ignoring an entire field of work presented in AutoML conferences and workshops since 2016.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PsjYlobKV8", "reviewer": "Reviewer_VJye", "summary": "The paper proposes a new approach to managing complex ML training setups called learning diagrams. It offers a new, more structured, way of looking at compositionality and interaction between different components (like data and different models) in composite learning settings. More specifically, in the learning diagrams framework nodes represent objects within a training regime, which can be data spaces, intermediate representations or parameter spaces. Each node, representing these spaces, is also associated with a distance measure -- a loss function. The edges represent transformations between two spaces (can be a learnable function represented with a neural net). Learning amounts to finding parallel passes that connect the same start and end nodes, and using their composite loss for learning. Composite loss is the sum of all pairwise losses of the parallel paths through a learning diagram. \n\nThe proposed concept of learning diagrams is rigorously grounded in the category theory. Additionally, the paper introduces a software framework that enables the specification of learning diagrams as a Python data structure.", "questions": "- can the framework presented here be used to enable a type of modularity where at each intermediate \"layer\" of a modular learner a routing decision has to be made (similar to MoEs/MoErging methods https://www.arxiv.org/pdf/2408.07057 )? i.e. here intermediate modules do not necessarily produce outputs that are in the label space of the underlying data", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Originality**:\nThe ideas presented in this work appear to be original in several ways. The paper applies category theory to create a formal framework that generalizes to different common training setups in ML. It proposes a way to automatically find parallel passes that enable the calculation of the composite loss used for training. Additionally, the paper proposes to view a training setup as a manipulable graphical object with adjustable and replaceable components (as data rather than code). I think this idea is original and leads to a nice generalization of many existing training regimes potentially opening doors for new interesting training setups.\n\n**Quality&Clarity**: The paper appeared very well structured to me. The clarity of exposition is generally very good.\n\n**Significance**: the significance of the work lies in providing a theoretical and software framework for designing multi-component learning systems. It facilitates compositionally in ML system design, which can lead to more maintainable and efficient model development (subparts can be easily replaced with new ones). It also facilitates collaborative model development through the integration of a plethora of independently trained models available online. The idea of representing training frameworks as a data structure (rather than code) can be impactful by increasing the transparency and interpretability of different complex training setups. It can help gain insights into relationships between different training regimes.", "weaknesses": "Learning diagrams provide a unified framework that encompasses many different setups, and this is demonstrated in this work by representing several training settings as in the language of learning diagrams. \n\nHowever, I think this paper could benefit from some concrete demonstrations of the advantages of learning diagrams over conventional methods in terms of performance, interpretability, training efficiency, revealing unknown connections between existing training setups etc..", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PTx9lT1fRA", "reviewer": "Reviewer_KsYd", "summary": "The paper addresses the problem of formally describing training protocols. For this purpose, they introduce a graphical formalisation called *learning diagrams*. First, they motivate their proposal by providing a series of examples of how common training paradigms can be accurately and concisely described via learning diagrams (namely image captioning, knowledge distillation, and few-shot learning). Then, they introduce a library called DiagrammaticLearning, that provides an interface for describing and building training setups. Finally, they formalise the process of describing models and data and compiling them into loss functions using category theory. \n\nThe library, called *DiagrammaticLearning.jl*, provides an interface to describe learning diagrams as labeled directed graph. Intuitively, vertices describe data sources and losses, and edges represent models, which in turn are functions applied to the data. The library then detects and sums over parallel paths to determine a single loss function. Finally, the library can convert the learning diagram into an implementation using a deep learning framework (e.g. PyTorch). Thus, the library is an intermediate language to describe training protocols and converting them into actual implementations.", "questions": "- Could you elaborate on the role of DiagrammaticLearning as a compiler (Section 3.3)? How exactly does it compile the graph into e.g. PyTorch training code? \n- Is the framework compatible with common PyTorch frameworks, e.g. HuggingFace?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Communicating training setups used is often a significant challenge and can lead to miscommunication, especially in scenarios where multiple models interact with each other. Formally specifying training setups provides an opportunity to address this problem.\n- The series of continuously more complex examples makes it easy to follow the ideas behind learning diagrams. \n- The authors evaluate their framework against the original implementations of the selected examples and demonstrate that it achieves similar performance.", "weaknesses": "- Despite using an intuitive graphical representation, the code in listing 1 is hard to read and not particularly intuitive. \n- Section 3 contains little details about the actual implementation. \n- While the authors mention that they have removed the code for the double blind review process, I would have appreciated anonymised code in the supplementary material considering that the framework is the main contribution. This would have allowed the reviewers to interact with the framework, provide more detailed feedback regarding functionality and usability, and see whether it actually supports development in the ways described in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)