Neural Circuit Diagrams: Robust Diagrams for the Communication, Implementation, and Analysis of Deep Learning Architectures
Authors: Vincent Abbott
Venue: Accepted by TMLR ()

Abstract: Diagrams matter. Unfortunately, the deep learning community has no standard method for diagramming architectures. The current combination of linear algebra notation and ad-hoc diagrams fails to offer the necessary precision to understand architectures in all their detail. However, this detail is critical for faithful implementation, mathematical analysis, further innovation, and ethical assurances. I present neural circuit diagrams, a graphical language tailored to the needs of communicating deep learning architectures. Neural circuit diagrams naturally keep track of the changing arrangement of data, precisely show how operations are broadcast over axes, and display the critical parallel behavior of linear operations. A lingering issue with existing diagramming methods is the inability to simultaneously express the detail of axes and the free arrangement of data, which neural circuit diagrams solve. Their compositional structure is analogous to code, creating a close correspondence between diagrams and implementation.

In this work, I introduce neural circuit diagrams for an audience of machine learning researchers. After introducing neural circuit diagrams, I cover a host of architectures to show their utility and breed familiarity. This includes the transformer architecture, convolution (and its difficult-to-explain extensions), residual networks, the U-Net, and the vision transformer. I include a Jupyter notebook that provides evidence for the close correspondence between diagrams and code. Finally, I examine backpropagation using neural circuit diagrams. I show their utility in providing mathematical insight and analyzing algorithms' time and space complexities.

Published in Transactions on Machine Learning Research (12/2023)
Neural Circuit Diagrams: Robust Diagrams for the Communication, Implementation, and Analysis of Deep Learning Architectures
Vincent Abbott Australian National University
Vincent. Abbott edu
Reviewed on OpenReview:
Abstract
Diagrams matter. Unfortunately, the deep learning community has no standard method for diagramming architectures. The current combination of linear algebra notation and ad-hoc diagrams fails to offer the necessary precision to understand architectures in all their detail. However, this detail is critical for faithful implementation, mathematical analysis, further innovation, and ethical assurances. present neural circuit diagrams, graphical language tailored to the needs of communicating deep learning architectures. Neural circuit diagrams naturally keep track of the changing arrangement of data, precisely show how operations are broadcast over axes, and display the critical parallel behavior of linear operations. A lingering issuc with existing diagramming methods is the inability to simultancously express the detail of axes and the free arrangement of data, which neural circuit diagrams solve. Their compositional structure is analogous to code, creating close correspondence between diagrams and implementation.
In this work, 1 introduce neural circuit diagrams for audience of machine learning searchers. After introducing neural circuit diagrams, I cover host of architectures to show their utility and breed familiarity. This includes the transformer architecture, convolution (and its diffieult-oo-expain extensions), residual networks, the U-Net, and the vision trans- former. I include a Jupyter notebook that provides evidence for the close correspondence between diagrams and code. Finally, I examine backpropagation using neural circuit di- agrams. 1 show their utility in providing mathematical insight and analyzing algorithms' time and space complexities.
1 Introduction
1.1 Necessity of Improved Communication in Deep Learning
Deep learning models are immense statistical engines. They rely on components connected in intricate ways to slowly nudge input data toward some target. Deep learning models convert big data into usable predictions, forming the core of many AT systems. The design of model its architecture can significantly impact performance (Krizhovsky et al., 2012), case of training (He ct al., 2015; Srivastava ct al., 2015), genoralization (Toffe & Szegedy, 2015; Ba al., 2016), and ability to efficiently tackle certain classes of data (Vaswani al., 2017; Ho et al., 2020).
Architectures can have subtle impacts, such as different image models recognizing patterns at various scales (Ronneberger et al., 2015; Luo et al., 2017). Many significant innovations in deep learning have resulted from architecture design, often from frighteningly simple modifications (He et al., 2015). Furthermore, architecture design is constant flux. New developments constantly improve on statc-of-the-art methods (He al., 2016; Leo, 2023), often showing that the most common designs are just one of many approaches worth investigating (Liu et al., 2021; Sun et al., 2023).

---

Published in Transactions on Machine Learning Research (12/2023)
However, these critical innovations are presented using ad-hoc diagrams and linear algebra notation (Vaswani et al., 2017; Goodfellow et al., 2016). These methods are ill-equipped for the non-linear operations and actions on multi-axis tensors that constitute deep learning models (Xu et al., 2023; Chiang et al., 2023). Furthermore, these tools are insufficient for papers to present their models full detail. Subtle details such as the order of normalization or activation components can be missing, despite their impact on performance (He et al., 2016).
Works with immense theoretical contributions can fail to communicate equally insightful architectural de- velopments (Rombach et al., 2022; Nichol & Dhariwal, 2021). Many papers cannot be reproduced without reference to the accompanying code. This was quantified by Raff (2019), where only 63.5% of 255 machine learning papers from 1984 to 2017 could be independently reproduced without reference to the author's code. Interestingly, the number of equations present was negatively correlated with reproduction, further highlighting the deficits of how models are currently communicated. The year that papers were published had no corrciation to reproducibility, indicating that this problem is not resolving on its own.
Relying code raises many issues. The reader must understand specific programming framework, and there is burden to dissect and reimplement the code if frameworks mismatch. Without reference to a blueprint, mistakes in code cannot be cross-checked. The overall structure of algorithms is obfuscated, raising ethical risks about how data is managed (Kapoor & Narayanan, 2022).
Furthermore, papers that clearly explain their models without resorting to code provide stronger scientific insight. As argued by Drummond (2009), replicating the code associated with experiments leads to weaker scientific results than reproducing procedure. After all, replicating an experiment perfectly controls all variables, including irrelevant oncs, making it difficult to link any independent variable to the observed outcome.
However, in machine learning, papers often cannot be independently reproduced without referencing their accompanying code. As result, the machine learning community misses out on experiments that provide general insight independent of specific implememnatiins. Improved communication of architectures, therefore, will offer clear scientific value.
1.2 Case Study: Shortfalls of Attention is All You Need
To highlight the problem of insufficient communication of deep learning architectures, I present a case study of Attention is All You Need, the paper that introduced transformer models (Vaswani et al., 2017). Introduced in 2017, transformer models have revolutionized machine learning, finding applications in natural language processing, image processing, and generative tasks (Phuong & Hutter, 2022; Lin et al., 2021).
Transformers' effectiveness stems partly from their ability to inject external data of arbitrary width into base data. I refer to axes representing the number of items in data as a width, and axes indicating information per item as depth.
An attention head gives weighted sum of the injected data's value vectors, V. The weights depend on the attention score the base data's query vectors, Q, assign to each key vector, K, of the injected data. Value and key vectors come in pairs. Fully connected layers, consisting of learned matrix multiplication, generate Q, K, and V vectors from the original base and injected data. Multi-head attention uses multiple attention heads in parallel, enabling efficient. parallel operations and the simultaneous learning of distinct attributes.
Attention All You Need, which I refer to as the original transformer paper, explains these algorithms using diagrams (see Figure 1) and equations (see Equation 1,2,3) that hinder understandability (Chiang et al., 2023; Phuong & Hutter, 2022).
2

---

Published in Transactions on Machine Learning Research (12/2023)
The Q, K, and V inputs to each attention head come from the learned linear projections of multi head attention.
The dimension of matrix multiplication and the SoftMax is ambiguous.
Scaled Dot- Product Attention
Probabltes
MarAAl
SoftMa
Loptz
Encoder
Docoder
Feod Forword
local
Mattful
7AC Npr
&.Nc Feud Forwrr
2
Multi-Head Ateention
No
ARS Norm MLB- Musiion
Foc Nor
Lindar
Mali Attecioon
Gonozt
Pociloonll Fricziing
Eroodiing
Nitention
KCL
Emoediing
Embeesting
Linoor
Inpuis
Outputs (et ged right)
These V, K, and Q values cupies in situation (A);
while Q is seperate in situation (B)
Figure 1: My annotations of the diagrams of the original transformer model. Critical information is missing regarding the origin of Q, K, and V values (red and blue), and the axes over which operations act (green).
QK"
Attention(Q,KK, V) Soft
V (d& is the key depth)
(1)
Jdly
MultiHead K, V) Concat(heady, ...,head,))) W where head, Attention (QW? KW," vw.)
(2)
(3)
The original transformer paper obscures dimension sizes and their interactions. The dimensions over which SoftMax and matrix multiplication operates is ambiguuus (Figure 1.1, green; Equation 1, 3).
Determining the initial and final matrix dimensions is left to the reader. This obscures key facts required to understand transformers. For instance, K and can have a different width to Q, allowing them to inject external information of arbitrary width. This fact is not made clear in the original diagrams or equations. Yet, is necessary understand why transformers are effective tasks with variable input widths, such language processing.
The original transformer paper also has uncertainty regarding Q, K, and V. In Figure 1.1 and Equation 1, they represent separate values fed to each attention head. In Figure 1.2 and Equation 2 and 3, they are all copies of each other at location (A) of the overall model in Figure 1.3, while Q is separate in situation (B).
Annotating makeshift diagrams does not resolve the issue of low interpretability. As they are constructed for a specific purpose by their author, they carry the author's curse of knowlodge (Pinker, 2014; Hayes & Bajzek, 2008; Ross et al., 1977). In Figure low interpretability arises from missing critical information, not from insufficiently annotating the information present. The information about which axes are matrix multiplied or are operated on with the Max is simply not present.
Using and index over data, we have Soft Max(v)[[] exp(v(+))/4 exp(v[k])
3

---

Published in Transactions on Machine Learning Research (12/2023)
Therefore, we need to develop a framework for diagramming architectures that ensures key information, such as the axes over which operations occur, is automatically shown. Taking full advantage of annolating the critical information already present in neural circuit diagrams, I present alternative diagrams in Figures 20, 21, and 22.
1.3 Current Approaches and Related Works
These issues with the current ad-hoc approaches to communicating architectures have been identified prior works, which have proposed their own solutions (Phuong & Hutter, 2022; Chiang al, 2023; Xu et al., 2023; Xu & Maruyama, 2022). This shows that this is known issue of interest to the deep learning community. Non-graphical approaches focus on enumerating all the variables and operations explicitly, whether by extending linear algebra notation (Chiang et al., 2023) or explicitly describing every step with pseudocode (Phuong & Hutter, 2022). Visualization, however, is essential to human comprehension (Pinker, 2014; Borkin et al., 2016; Sadoski, 1993). Standard non-graphical methods are essential to pursue, and the community will benefit. significantly from their adoption; however, standardized graphical language is still nceded.
The inclination towards visualizing complex systems has led to many tools being developed for industrial applications. Labview, MATLAB's Simulink, and Modelica are used in academia and industry to model various systems. For deep learning, TensorBoard and Torchview have become convenient ways to graph architectures. These tools, however, do not ofler sufficient detail to implement architectures. They are often dedicated to one programming language or framework, meaning they cannot serve as a general means of communicating new developments. Besides, a rigorously developed language for deep learning architectures would help to improve these tools. This requires diagrams equipped with mathematical framework that captures the changing structure of data, along with key operations such broadcasting and linear transformations.
innepeeenn graphical
Many mathematically rigorous graphical methods exist for a variety of fields. This includes Petri nets, which have been used to model several processes in research and industry (Murata, 1989). Tensor networks were developed for quantum physics and have been successfully extended to deep learning (Biamonte Bergholm, 2017: Xu et al., 2023; Xu Maruyama, 2022). Xu al. (2023) showed that re-implementing models after making them graphically explicit can improve performance by letting parallelized tensor algorithms be employed. Robust diagrams, therefore, can benefit both the communication and performance of architectures. Formal graphical methods have also been developed in physics, logic, and topology (Baez & Stay, 2010; Awodey, 2010).
All these graphical methods have boen found to represent an undorlying category, a mathematical space with well-defined composition rules (Meseguer & Montanari, 1990; Baez & Stay, 2010). category theory approach allows a common structure, monoidal products, to define an intuitive graphical language (Selinger, 2009; Fong & Spivak, 2019). Category theory, therefore, provides a robust framework to understand and develop new graphical methods.
However, a noted issue (Chiang et al., 2023) of previous graphical approaches is they have difficulty expressing non-linear operations. This arises from a tensor approach to monoidal products. Data brought together cannot necessarily be copied or deleted. This represents, for instance, axes brought together to form a matrix and this approach makes linear operations elegantly manageable. Tt, however, makes expressing copying and delction impossible. The alternative Cartesian approach allows copying and deletion, reflecting the mechanics of classical computing.
The Cartesian approach has been used to develop mathematical understanding of deep learning (Shiebler et al., 2021; Fong et al., 2019; Wilson & Zanasi, 2022; Cruttwell et al., 2022). However, Cartesian monoidal products do not automatically keep track of dimensionality and cannot easily represent broadcasting or linear operations. These works often rely on the most rudimentary model of deep-learning networks as soquential linear layers and activation functions, despite residual networks having become the norm (He et al., 2015; 2016). The graphical language generated by a pure Cartesian approach fails to show the details of architectures, limiting its ability to consider models as they appear in practice.
4

---

Published in Transactions on Machine Learning Research (12/2023)
The issue of only looking at rudimentary, linear layer-activation layer models is pervasive in deep learning research (Zhang et al., 2017; Saxe et al., 2019; Li et al., 2022). There are uncountably many ways of relating inputs to outputs. Every theory or hypothesis about deep learning algorithms has to assumee that we are working with some subset of all possible functions. However, specifying this subset means theoretical insights can only apply to that subset. This precludes us from using such theories to compare disparate archilectures and make design choices.
The problem of only considering rudimentary models is partially a consequence of us not having the tools to robustly represent more complex models, never mind the tools to confidently analyze them. Category theory-based diagrams can serve as models of intricate systems. Structure-preserving maps allow analyses to scale over entire models. Therefore, developing comprehensive diagrams that correspond to mathematical expressions can be the first step in a rigorous theory of deep learning architectures with clear practical applications.
The literature reveals combination of problems that need to be solved. Deep learning suffers from poor communccttion and needs a graphical language to understand and analyze architectures. Category theory can provide a rigorous graphical language but typically forces choice between tensor Cartesian approaches. The elegance of tensor products and the ficxibility of Cartesian products must both be available to properly represent architectures. A category arises when a system has sufficient. compositional structure, meaning a non-category theory approach to diagramming architectures will likely yield a category anyway. The challenge of reconciling Cartesian and tensor approaches, therefore, remains.
1.4 The Philosophy of My Approach
As I am introducing these diagrams, I have burden to explain how I think they should be used and to address criticisms of creating diagramming standard in the first place. I will take brief aside to address these points, which believe will aid in the adoption of neural circuit diagrams.
These diagrams are intended to express sequential-tensor deep learning models. This is in contrast to machine learning or artificial intelligence systems more generally. Deep learning models are machine learning models with sequential data processing through neural network layers. do not cover recursive or branching models in this work. Furthermore, assume data is always in the form of tuples of tensors. Generalizing diagrams to further contexts is an exciting avenue for future rescarch.
By making these assumptions, I develop diagrams specialized for some of the most essential but difficult-to- explain systems in artificial intelligence research. Researchers outside the narrow scope of requential-tensoe deep learning models often rely on these tools. By more clearly communicating them, rescarchers who may not be up to date on the latest innovations or aware of their options stand to benefit an immense deal.
do not expect two independent teams to diagram architectures the exact same way. Indeed, I do not believe the appropriate diagramming framework would have this property. Diagrams should have the flexibility to allow for innovations and to appeal to the audience's e's level of knowledge. Instead, the benefit of my framework is to have comprehensive, robust diagrams with clear correspondence to implementation and analysis, in contrast to ad-hoc diagrams, which often fail to include critical information.
Neural circuit diagrams can be decomposed into sections that allow for layered abstraction. The exact details of code can be abstracted into single-symbol components. Scctions of diagrams can be highlighted for the reader's clarity, and repeated patterns can be defined components. Diagrams have an immense compo- sitional structure. The horizontal axis represents sequential composition, and the vertical axis represents parallel compasition. Sections and components can be joined like Lego bricks to construct models.
This sectioning allows for a close correspondence between diagrams and implementation. Every highlighted section becomes a module in code. Diagrams, therefore, provide a cross plassorm blueprint for architectures. This allows implementations to be cross-checked to a reference, increasing reliability. Furthermore, which components are abstracted and the level of abstraction can vary depending on the audience, leading clearer, specialized communication.
5

---

Published in Transactions on Machine Learning Research (12/2023)
A common criticism is that introducing new standard simply increases the number of standards, worsening the issue trying to be solved (below). T do not believe this is a relevant crit ique for deep learning diagrams. Currently, there are no standard diagramming methods. Every paper, in a sense, has its own ad-hoc diagramming scheme. Compared to this, neural circuit diagrams only need to be learned once, after which architectures can be clearly and explicitly explained. Furthermore, they build on existing research on robust. monoidal string diagrams, which have been found be a universal standard for various fields (Bacz & Stay, 2010).
1.5 Contributions
To address the need for more robust communication and analysis of deep learning architectures, I introduce neural circuit diagrams. Neural circuit diagrams solve the lingering challenge of accommodating both the details of axes (the tensor approach) and the free arrangement of data (the Cartesian approach) in diagrams. They are specialized for sequential algorithms on memory states consisting of tuples of tensors.
Diagramming the details of axes means the shape of data is clear throughout a model. They easily show broadcasting and provide a graphical calculus to rearrange lincar functions into equivalent forms. At the same time, they clearly represent luples, copying, and deletion, processes that lypical graphical methods struggle with. This makes them uniquely capable of accurately representing deep learning models.
Inspired by category theory and especially monoidal string diagrams (Selinger, 2009; Bacz Stay, 2010), this work builds on a literature of robust diagramming methods. However, the category theory details are omitted to maximize impact among machine learning researchers.
The benefits of neural circuit diagrams are many. They allow for clearer communication of new developments, making ideas more rapidly disseminated and understood. They offer robust blueprints for designing and implementing models, accelerating innovation and streamlining productivity. Furthermore, they allow for rigorous mathematical analysis of architectures, bringing us closer to a theoretical understanding of deep learning.
These points are evidenced by diagramming a host of architectures. I cover basic multi-layer perceptron, the transformer architecture, convolution (and its difficult-to-expliin permutations), the identity ResNet, the U-Net, and the vision transformer. provide a Jupyter notebook that implements these diagrams, which provides further evidence for the close relationship between diagrams and implementation. Finally, offer a novel analysis of backpropagation, which shows the utility of noural circuit diagrams for rigorous analysis of architectures.
2 Reading Neural Circuit Diagrams
2.1 Commutative Diagrams
We aim to craft diagrams that precisely represent deep learning algorithms. While these diagrams will eventually be generalized, we will initially concentrate on common models. Specifically, we will explore models that successively process data of predictable types. To facilitate understanding, we will introduce diagrams of gradually increasing complexity. To begin, lot's delve into an intuitive diagram, where symbols represent data types, and arrows signify the functions connecting them.
Note, 1 use forward composition with ",", meaning f str - int composes with g int float by (fig): str + float.
Figure 2: We have two functions: f f: str - int and : int - float. These functions can be composed into single function (fig) str float. In commuting diagrams, we represent data types, such as str, int, and float, with floating symbols, while functions are denoted by arrows connecting them.
int
g
(Function)
fig
float
(Data type) str
6

---

Published in Transactions on Machine Learning Research (12/2023)
2.1.1 Tuples and Memory
Algorithus are rarely composed of operations on a single variable. Instead, their steps involve operations on memory states composed of multiple variables. The data type of a memory state a tuple of the variables which compose So, a state containing an int and str would have a type int str.
Consider single algorithmie step acting on compound memory state Ax BxC. A function f BxC D acting this memory state would give overall step with shape Id[A] x A (x - Ax D. Note that [d[A] is the identity. We need to indicate A, even though f does not act on it, so that the initial and final memory states are properly shown. In Figure 3, I diagram f along another function g A D E.
AxD
Figure 3: Here, diagram two functions, f: : C D and g: AxD - E, acting together. To represent the full memory states, we are required to amend into [d[A]xf= Ax [ C) A D. The composed function (Id[A) x fig Ax (B E.
|d| A X f
g
AxBx C
fig
2.2
String Diagrams
These commuting diagrams fall short, however. As algorithms scale, operations and memory states get more complex. Usually, functions only act some variables. However, it is not clear how to these targeted functions. Compound data types and compound functions are better suited by reorienting diagrams as in Figure 4. We will have horizontal wires represent types, and symbols represent functions. Diagrams are forced to horizontally go left to right.
AxBxC [d[A] X
x
AXXD
AXD
E
AXBXC h E
g
Figure 4:
We reorient diagrams to go left to right. Wires represent data types, and symbols represent
functions. This expression defines
This reorientation allows us to represent compound lypes and functions easily. We can diagram tupled types A x B a wire for A and a wire for B vertically stacked, but separated by a dashed line. For increased clarity, we can draw boxes around functions. In Figure 5, we see a clear reexpression of Figure 4. Here, we have the unchanged A variable untouched by which acts only on BxC.
vertical section represents memory
Diagrams also be split along horizontal dashed lines.
state of type A x B C.
A
A
E
E
D
g
g
C
f
f
This vertical section represents the function
f:BxC-D, which stacked into,
(Bx D.
Figure 3: Tupled data types diagrammed with wires separated by dashed lines. This clearly shows when functions act on only some variables.
Every vertical section of a diagram represents something. Either, shows which data type present memory, or which function is applied at this step. Diagrams can always be decomposed into vertical sections, each of which must compose with adjacent sections ensure algorithms well-defined. Diagrams also

---

Published in Transactions on Machine Learning Research (12/2023)
be split along dashed lines. Diagrams are built from these vertically and horizontally composed sections, with wires acting like jigsaw indents.
2.3
Tensors
We will specialize diagrams for deep learning models whose memory states tuples of tensors. Tensors are numbers arranged along axes. So, scalar R is rank 0 tensor, vector R' is rank tensor, table R is rank 2 tensor, and so on. If our diagram takes tensor data types, we get something like Figure 6.
R3
R
R 1x2
Figure 6: Similar to Figure 5, but with data types being tensors.
R
R33x3
R
However, benefit from diagramming the details of axes. Instead of diagramming wire labeled IR axb diagram a wire labeled and wire labeled b, without dashed line separating them. This lets us diagram Figure 6 into the clear form of Figure 7.
We express tensor data types by representing axes with wires.
R 4x2
4
2
This us reexpress diagrams by drawing the details axes.
RR
R3
R lx2
3
1
2
R
R62
3333
g
g
3
R
Figure 7: We can diagram types gabb as two wires labeled and b, without a dashed line separating them. (See cell 2, Jupyter notebook.)
2.3.1 Indexes
Values in tensors are accessed by indexes. A tensor A E R for example, has constituent values A[4,, ja] e IR where E (0...3) and ja (0...2). Indexes also be used to access subtensors, expressions A[4,: :] a This subtensor extraction is therefore an operation IR4x R². We diagram by having indexes act the relevant axis. Indexes are diagrammed with pointed pentagons, or kets L). This
have
type of subtensor extraction is diagrammed according to Figure 8.
We express extractions, grabbing A|2, :], by an index applied the
2
3
Figure 8: We diagram indexes with pointed pentagons labeled with the index being extracted. (See cell 3, Jupyter notebook.)
8

---

Published in Transactions on Machine Learning Research (12/2023)
Here, the symbols ig iterate over {0 ... a 1/. and Ja over (0... b 1). This diagram covers all the indexes.
A
jb
b
jb
Figure 9: These subtensors are defined such that A[44, :][/a] A[i.. This expression is the same in the reverse order. (See cell 4, Jupyler nolebook.)
2.3.2 Broadcasting
Broadcasting critical to understanding deep learning models. It lifts operation to act in parallel over
additional axes. Here, we show operation G R Ig² lifted to operation G' R 1x3 - R gx2 We
diagram this broadcasting by having the 4-length wire pass over G, adding 4-length axis to its input and output shapes. Formally, we define ((()) = :))[ This shown in Figure 10.
An operation single 3- length is applied in parallel separate rows broadcasting.
This lifts an operation G: R² R² to R 1x3
R Rxx
x 4
Broadcasting
3
G
2
G
G
Broadcasting be formally defined of indexes. The left and right hand side below
equal. Graphically, we see that broadcasting allows indexes pass over expressions. Note how the
expressions holds for 2 replaced by any valid index ig [0, 1, 3J.
4
2
def
2
3
G
2
2
3
3
G
Figure 10: An operation is lifted over a 4-length axis by broadcasting. This applies G over corresponding subtensors. Broadcasting can be formally defined by equating indexes before and after an operation. (See cell 5, Jupyter notebook.)
Inner broadcasting acts within tuple segments. A 1x3
x R collection of data can be reduced to IR x RR in different ways. Therefore, there are 4 ways of applying an operation
IR x R IR² to it. This gives a function lifted by "inner broadcasting", which has a shape IR 4x*3 RR* IR 4x2 We diagram this by drawing wire from the source tuple segment over the function, as shown in Figure 15. This adds an axis of equal length to the target tuple segment and to the output, reflecting the shape of the lifted operation.
Broadcasting naturally represents dementtwise operations. A function values f R - R when broad- cast, gives an operation Xx R Ixa One length axes do not change the shape of data, and can be freely amended or removed from pre-existing shapes by arrows. This means we diagram element-wist functions by drawing incoming and outgoing arrows, which represent the amendment and removal of a 1-length axis. This is shown Figure 12.
9

---

Published in Transactions on Machine Learning Research (12/2023)
IR 4x3 x collection data can be reduced to x R 4 different ways. Therefore, an operation
R R R2 can be lifted an operation R 5x3 R
R 5x2
inner broadcasting.
X4 Inner Broadcasting
2
4
2
2
H
H
4
H
Figure 11: Lifting operation within a tuple segment gives inner broadcasting, We diagram it by having wire from the target tuple segment over the function, reflecting the shape of the lifted function. (See cell 6, Jupyter notebook.)
Broadcasting with 1-length axis leaves data types unchanged, 1-length axes can freely introduced and removed with arrows; >
Operation on f:
Element-wise f
Figure 12: Element wise operations can be naturally shown with broadcasting. (See cell 7, Jupyter notebook.)
f
4
4
x4x3 Broadcasting
3
Starting with function values, we apply it onto every value in data by broadcasting.
2.4 Linearity
Linear functions are an important class of operations for deep learning. Linear functions can be highly parallelized, especially with GPUs. Previous works have shown how graphically modeling linear functions, and reimplementing algorithms can improve performance (Xu et al., 2023). Linear functions have immense regularity. Standard monoidal string diagrams rely on these properties to provide elegant graphical languages for various fields (Baez & Stay, 2010).
Linear functions are required to obey additivity and homogeneity, as shown in Figure These operations are closed under composition, so applying lincar maps onto each other gives another lincar map. Importantly to us, they are natural with respect to broadcasting. This means for any two linear functions and g, the equality in Figure 14 holds. This means they can be simultaneously broadcast. This lets series of linear functions be efficiently parallelized and flexibly rearranged.
Elementwise
Elementwise
Addition
Addition
A
A
A
b
Figure 13: A subset of functions between R to R² are linear, obeying additivity and homogeneity. This class of functions are closed under composition and has many im- portant composition properties.
A
Scalar Multiplication
Scalar
Multiplication
A
b
8
A
10

---

Published in Transactions on Machine Learning Research (12/2023)
Linear operations natural with respect broadcasting. meaning they slide past each other.
a
F
b
a
F
b
a
F
C
C
G
d
G
d
d
c
G
d
Figure 14: Linear functions are natural with respect each other and broadcasting. This means the above equality holds, letting expressions be flexibly rearranged.
However, a pure monoidal string diagram has difficulty representing non-linear operations, a noted issue (Chiang et al., 2023). Neural circuit diagrams have Cartesian products and broadcasting, which are not generally analogous to how monoidal string diagrams combine linear functions. we know functions are linear, we can use diagrams to efficiently reason about algorithms. By focusing on lincar functions, we can take advantage of their parallelization properties.
2.4.1 Multilinearity
There is an important distinction between linear and multilincar operations. Inner products, for example, are multilinear. The inner product v(x, = x = Σ, a[i] w[i] is linear with respect to each input. So, u(x + y. u(x,y) + u(z,y), and similarly for the second input. However, is not linear with respect to dise addition over its entire input and output, u(x, X2 + y2) # v1)1,y1) + u(x2. Y2). Compare this to copying A, which can show is linear.
A: A:R²
R axa and x,y E R², A
A(x) == (x,x)
A(x+y) = (x (x,x) yy
=((x) A(y)
A( = (1.x,) xx = (x,x)
=xx((x)
To simultaneously broadcast multilinear functions, we note that every multilinear operation equals an outer product followed by a lincar function. The outer product is the ur-multilincar operation, taking a tuple input and returning a Lensor, which takes the product over one element from each tuple segment. given by @ IR x R2 - R axb All muple-mutillinarr functions M : R x H² - R have an associated tensor-linear form MA paxb R such that Ø; M, = M. We diagram the outer product by simply having tuple line ending, which will often occur before host of linear operations are simultancously applied.
2.4.2 Implementing Linearity and Common Operations
Key lincar and multilinear operations be implemented by the einops package (Rogozhnikov, 2021), leading to elegant implementations of algorithms. Some key linear operations are inner products, which sum over an axis, transposing, which swaps axes, views, which rearranges axes, and diagonalization, which makes axes take the same index.
With neural circuit diagrams, we can clearly show these operations. We show inner products with cups, transposing by crossing wires, views by solid lines consuming and producing their respective shapes, and diagonalization by wires merging. As these operations are linear, they can be simultaneously applied. The interaction of wires shows how incoming axes coordinate to produce outgoing axes. The einops package symbolically implements these operations by having incoming and outgoing axes correspond to symbols.
An example Ihat combines many of these operations is a section of multi-head allention shown in 16. It employs an outer product, a transpose, a diagonalization, an inner product, and an element-wise operation. The input to this algorithm a tuple of tensors. Axes with an overline are width, representing the amount of rather than detail per thing. Though a complex expression, we can break this figure up as in Figure 7 and implement the interaction of wires using einops, shown in Figure 17.
11

---

Published in Transactions on Machine Learning Research (12/2023)
Inner Product
Matrix Multiplication
We have 4 lists
Inner Broadcasting
choose from segment, from the second.
This gives us 4 x x2 possible det products for the output.
x Inner Broadcasting
The inner product takes two equal-length lists, mutliplies corresponding values, and results.
Multiplication
X4 Inner Broadcasting
Inner Praduct
drawn as a cup, 5
3
X Inner Broadcasting
Figure 15: An in-depth example of matrix multiplication, key multilinear inner broadcast operation. Inner products are defined on vectors R" x R=" - R Then, we inner broadcast them to act over matrices. The new Tx R - RRx operation is matrix multiplication. Therefore, we see that matrix multiplication is an instance an inner broadcast operation.
Equation Expression
1
X
k
Q a y x shaped tensor.
Implicit product
We can diagram a portion of multi-head at- tention, a sophisticated al- gorithm, with clarity using neural circuit diagrams.
Figure 16:
V
a lsaixke k
Q
k
Scale 1/2
1/2
shaped tensor
ht
*k
Inner Product
K
X
K
k
k
h
h
Transposing
Diagonalize
Implementation using einsum
Local menory contains, x kk # Transpose K,
V
# Q:
k
k
Figure 17: This section multi-hoad attention can be implemented using the einsum operation. Note the close relationship be- tween diagrams and im- plementation and how di- agrams reflect the memory states and operations of al- gorithms. (See cell 8, Jupyter notebook.)
h
h
Q, einops. KRR
ht)
1/22
# Implicit outer product and diagonalize,
k
X einops.einsun(Q, K, 'y k1 h, k2 x H
K
k
X
->
h
h
# Inner product,
X winops..ensunte
'y Rxh->yxh))
Scale,
x math. sqrt(()
Implementation using einsum
k
(with simultaneous broadcasting of linear functions)
12
# Local menory contains,
Q: y x
x einsun(,,
K
x x mgth.
h
12

---

Published in Transactions on Machine Learning Research (12/2023)
2.4.3 Linear Algebra
All linear functions I IR" - R2" have associated Raxb tensor that uniquely identifies them. This hints at the ability to transpose this associated tensor get new linear function, f :Rk" KK To extract these associated transposes, use the unit. The unit for shape a, given by 1j R R gaxa is linear map which returns r times the R=xa identity matrix, for IR.
Note that the associated transpose, which sends a linear function f : R" - R to f : - R" by transposing the associated R"x" tensor, different to transpose operation which sends R=xm to R=*" Associated transposes are used for mathematical rearrangement and are not usually directly implemented in code, though provide code examples in cell 9 of the Jupyter notcbook.
The unit and the inner product be arranged to give the identity map R R" as in Figure 18. This identity map can be freely introduced, split into unit and the identity matrix, and then used to rearrange operations. For example, this allows us convert the linear map F R R ga into F" R oxe R.. These associated tensors and transposes be used to better understand convolution (Section 3.3) and backpropagation (Section 3.6).
x ld
Associated Transpose
b
a
F
b
b
b a
FT.
C
Identity
The unit and inner product arranged in this manner give the identity.
G
H
Naturality among linear functions with respect broadcasting allows them to be horizontally shifted.
H
a
G
a
G
G
b
H
T
H
H
Simultaneously broadcasting linear operations is product, meaning praducts can shifted.
G H
G H
K =
a
Figure 18: Linear operations have a flexible algebra. Simultaneous operations may increase efficiency (Xu al., 2023). As the height of diagrams related to the amount of data stored in independent segments, it gives rough idea of memory usage. This is further explored in Section 3.6. (See cell 9, Jupyter notebook.)
These rearrangementss can transpose specific axes. linear operation Rab - R has an associated aaxxx tensor. This tensor can be associated with various lincar operations, such I bxa - R.. These different forms are often of interest to us, as they can efficiently implement the reverse of operations (see Figure 25, 30). To extract these rearrangements, we can selectively apply units and the inner product to reorient the direction of wires for linear operations.
13

---

Published in Transactions on Machine Learning Research (12/2023)
3 Results: Key Applied Cases
3.1 Basic Multi-Layer Perceptron
Diagramming a basic multi-layer perceptron will help consolidate knowledge of neural circuit diagrams and show their value as a leaching and implementation lool, as shown in Figure 19. We use pictograms to represent components analogous to traditional circuit diagrams and to create more memorable diagrams (Borkin et al., 2016).
isport torch.nn Baste Image secogniser This close copy an introductory Pylorch tutorial;
wwt.ttchhh..gtttttaaaatt... hts.h//ggnnnor///ss//uuuaaaaaaa wtw.d/dddeeetttaallleeema hhtel
class htt.:rggeeecoppiieemmmoobbbl
def init (self):
super()- in1t init__)
Flatten nn.Flatten()
self. 1ilfar relu. stack m.Sequential(
L28-ea 512), mn.ReLU(),
Basic Image Recogniser for Digits
(512 512), mn.ReLU(), (512, 18),
28
R
>
SoftMax
1
28
512
512
k= 10
10
X 28
L
L
L
Forward(self, x):
Linear
self- (x)-eelatte(()
Y 10
self.]inear: relu stack(x)
Xrpedd mn.50ftnxx(Y)
return
Y_ared
Figure 19: PoTorch code and neural cireuit diagram for basic MNIST (digit recognition) neural network taken from introductory PyTorch tutorial. Note the close correspondence between neural circuit diagrams and PyTorch code. (See cell 10, Jupyter notebook.)
Fully connected layers are shown boldface L, with boldface indicating component with internal learned weights. Their input and output sizes are inferred from the diagrams. If fully connected layer is biased, we add a "+" in the bottom right. Traditional presentations easily miss this detail. For example, many implementations of the transformer, including those from PyTorch and Harvard NLP, have a bias in the query, key, and value fully-connected layers despite Atlention All You Need (Vaswani et al., 2017) not indicating the presence of bias.
Activation functions just clement- wise operations. Though traditionally ReLU (Krizhevsky al., 2012), other choices may yield superior performance (Lee, 2023). With neural circuit diagrams,, the activation function employed can be checked at a glance. SoftMax is a common operation that converts scores into probabilitics, and we represent it with a left-facing triangle (4)) indicating values being "spread" to sum to 1.
As mentioned in Section 1.2, how operations such Sof are broadcast can be ambiguous in traditional presentations. This especially worrisome SoftMax can be applied to shapes of arbitrary size. On the other hand, the neural circuit diagram method of displaying broadcasting makes it clear how SoftMax is applied.
14

---

Published in Transactions on Machine Learning Research (12/2023)
3.2 Neural Circuit Diagrams for the Transformer Architecture
In Section 1.2, covered shortfalls in Allention is All You Need. We now have the lools to address these shorteomings using neural circuit diagrams. Figure 20 shows scaled-dot product attention. Unlike the approach from Attention All You Need, the size variables and the axes over which matrix multiplication and broadcasting occur is clearly shown. Figure 21 shows multi-head attention. The origin of queries, keys, and values are clear, and concatenating the separate attention heads using einsum naturally follows. Finally, we show the full transformer model in Figure 22 using neural circuit diagrams. Introducing such a large architecture requires an unavoidable level of description, and take some artistic license and notate all the additional details.
The input is tuple of data, with shapes clearly indicated. Overlines indicate associated
Scaled Dot- Praduct Attention
The attention scores are used for weighted sum over the value vectors.
amount of things, rather than their details.
Equation from the original paper.
Q
MatMul
Scale k
The output data has the same shape as the queries.
QKI
SoftMax
MatMui
SoftMax
V
K
Vdl
X
N
V
k-axis and K. The result is scaled, an element-wise operation. The output shape y relating each queries to each X injected items.
Matrix multiplication
A SoftMax proportions the attention scores over the axis. For each y queries, we now have proportioned X-doop vector.
Figure 20: The original equation for attention against a neural circuit diagram. The descriptions are un- necessary but clarify what is happening. Corresponds to Equation and Figure 1.1. (See cell 11, Jupyter notebook.)
Multi Head Attention
Learned fully connected layers Q. K. and V/ from the inputs.
These generate /-shaped data, k-deep vector for each attention head.
Linear
Scaled Dot-Prtduct Attention
m
k
L
h
1/2
Figure 21: Neural circuit diagram for multi-head attention. Implement- ing matrix multiplication is clear with the cross-platform the cinops package (Rogozhnikov, 2021). Corresponds Equation 2 and 3 and Figure 1.2. (See cell 12, Jupyter notebook.)
Copy
X
m
m
LK
k
h
y
k
k
We the injected data
L
V
Concat (Diagonalize)
Linear
generate K and V.
During matrix multiplication, diagonalize the axes, avoiding cross-awe axis ensuring parallel operations.
final fully connected layer combines the outputs of the separate heads into
15

---

Published in Transactions on Machine Learning Research (12/2023)
Neural Circuit Diagram for Transformers
framcwork representing deep learning Transfamer and comprehensve diagram for the orginal architecuure from Attention is All Vos Need We descrthe all nocessary companerts, enabling trchnicaly proficent novices who have read understand the transfommer architecture.
discrams
changed ward, we
Encoded Input
A
an ingut sentence, the sentence length, and nnt the number input model depth.
middep embodding, Multi-Head Self-Attention
Dropouts simply returiing their input like identity function, bus other simes O, deltting informationn They roduce ovtrrdlinnce any one neuron, nabilsnn the model
Impet Embeeddig
N(
6)
Capy
Linear
Scaled Dot Product Actention
E
Encoder Stack
21
hhe encoder processes the input with attention heads to extract necessary information for the decoder, which will it output.
12
farward layer applies connected layers to malleably adjust data
Positional Encoding positional encoding identifies the location of inputs by adding a sinosoidal pattern.
Add & Narm
Foed Forward
Add & Norm
R
>>
11
We have 6 sequential encoder stacks. Each learns its own wiights
Concat (Diagonalize)
I!
deep learning, we almost always use residuial networks. Instead of input, learn add small change.
completely change
activation
non-linear allowing
Output V HH
zotterss be discovered, Typically RELU.
The target Y sceds the generator by sither giving sentence to mimic or our best estimate so far, 8i' the number of words the target language. During training, the decoder's cupput attempts
d is
size.
For each of words, LayerNorm adds and scales the ands so that has mean of and standard deviation Each of the mn indeses has gain, This stabilizes network during
by
Target Embedding
mimic a targat. mask musurs sach of satimated wards only have
gresious words.
Y Pr
N( 6)
Masked Multi-Head Self-Attention
Decoder Stack
E
Capy
The decader uses the encoded input and generate output. The encoded input provides external information, for example, English sentence. The target provides the translation mimic the best estimate far.
(Opt.)
>4
Positional Encoding
Add & Norm
As data through each of the decoder stacks, information from the encoded
Feed Forvard
Add & Norm
>R
m
M
input A is repeatedly injected.
Multi-Hess Cross Attention
1
m
L
generated me-deap vectors mapped in' rn'-vectors containing the probability of each word in target language.
Encoded Imput
x
A
Adid Norm
Feed Forward
Add & Norm
R-4
>>
Copy
L
Linear
SoftMax
L
111
Figure 22: The fully diagrammed architecture from Attention is All You Need (Vaswani et al., 2017).
16

---

Published in Transactions on Machine Learning Research (12/2023)
3.3 Convolution
Convolutions are critical to underslanding computer vision archilectures. Different archilectures exlend and use convolution in various ways, so implementing and understanding these architectures requires convolution and its variations to be accurately expressed. However, these extensions are often hard to explain. For exam- ple, Py Torch concedes that dilation is "harder to describe". Transposed convolution is similarly challenging communicate (Zeiler al., 2010). A standardized means notating convolution and variations would aid in communicating the ideas already developed by the machine learning community and encourage more innovation of sophisticated architectures such vision transformers (Dosovitskiy et al., 2021; Khan et al., 2022).
In deep learning, convolutions alter a tensor by taking weighted sums over nearby values. With standard bracket notation to access values, a convolution over vector of length I by a kernel w of length given by, (Note: we subscript indezes by the aris over which they act.)
(iny( Σ Lold + f.ll-
The maximum in value is such that it does not exceed the maximum index for u[iy jx]. Starting indexing at 0, we get I 1 imax + Jmax = k 2, so the length of the output is therefore ÿ = H k + Note how convolution is a multilinear operation; it is linear concerning each vector input and w. Therefore, it has a tensor-lincar form with an associated tensor, the convolution tensor, that uniquely identifies it.
Conv(u, w)y
ΣΣ (x) Juy ul).v[e) w[ie] t, if la bn + jk.
(x)
else.
We diagram convolution with the below diagram, Figure 23. We then transpose the linear operation into a more standard form, letting the input be to the left, and the kernel to the right.
V x
iv
k+1
in
1)
Figure 23: Convolution is multilinear operation, with an associated tensor. This tensor is transposed into a standard form.
W k
We contract over the convolution Here, the + operation that returns 1 the sum of the indexes equals the coindex, zero atherwise.
We transpose the operations arrive at the standard form.
We typically work with higher dimensional convolutions, in which case the indexes act like tuples of indexes. We diagram axes that act in this tandem manner by placing them especially close to each other and labeling their length by one bolded symbol akin to a vector. In dimensions the convolution tensor becomes;
1
if (Lao, (m) = (1 in1) + (Jea. jea).
(* 2D) igo I Jkn jaI Cog CoT
0 clsc.
Figure 24 shows what convolution does. It takes an input, uses linear operation separate it into overlapping blocks, and then broadcasts operation over each block. Using neural circuit diagrams, we
17

---

Published in Transactions on Machine Learning Research (12/2023)
easily show the extensions of convolution. A standard convolution operation tensors the input with channel depth axis, and feeds each block and the channel axis through a learned linear map.
Additionally, we can take an average, maximum, or some other operation rather than a linear map on each block. This lets us naturally display average or max pooling, among other operations. Displaying convolutions like this has further benefits for understanding. For example, 1 convolution tensors give a linear operation K R xxx1 which we recognize to be the identity. Therefore, x kernels are the same as broadcasting over the input.
Standard Convolution
Average Pool
Max Pool
X
x
*
X
++1
x-k+1 k
k
k
E
Max
Figure 24: Convolution and related operations, clearly shown using neural circuit diagrams.
L
Co
4
C1
Co
Shorthand
x
k+1
Stride and dilation scale the contribution of in or ju in the convolution tensor, increasing the speed at which the convolution scans over its inputs. This changes the convolution lensor inlo the form of Equation 4. We diagram these changes by adding the or d multiplier where the axis meets the tensor in Figure 25. These multipliers also change the size of the output, allowing for downscaling operations.
lg =
ig + d ju-
(* 8, dy ex] =
(4)
clse.
(k 1) -
T =
(5)
We often want to make slight adjustments to the output size. This is done by padding the input with zeros around its borders. We can explicitly show the padding operation, but we make it implicit when the output dimension does not match the expectation given the input dimension, kernel dimension, stride, and dilation used.
Stride make the output axis have far lower dimension than the input axis. This perfect for down- scaling. We implement upscaling by transposing strided convolution, resulting in operation with many more output blocks than actual inputs. We broadcast over these blocks to get our high digh-dimessional output.
Stride and Dilation
Padding
As the labeled output differs from the expected output, we infer the input padded.
Transposed Convolution
Transposed convolution accesses convolution tensor in different manner, meaning the output is larger than input. This be combined with stride implicit padding upscale inputs.
x
x
indexes the two axes are multiplied by the associated amount, letting us sum over instances where & s +dx
+
*
X k
t
Figure 25: Stride, dilation, padding, and transposed convolution shown with neural circuit diagrams.
Transposed convolution is challenging to intuit in the typical approach to convolutions, which focuses on visualizing the scanning action rather than the decomposition of an image's data structure into overlapping blocks. The blocks generated by transposed convolution can be broadcast with linear maps, maximum, average, other operations, all easily shown using neural circuit diagrams.
18

---

Published in Transactions on Machine Learning Research (12/2023)
3.4 Computer Vision
In computer vision, the design of deep learning architectures is critical. Computer vision tasks often have enormous inputs that are only tractable with high degree of parallelization (Krizhevsky et al., 2012). Architectures can relate information at different scales (Tuo et al., 2017), making architecture design task- dependant. Sophisticated architectures such as vision transformers combine the complexity of convolution and transformer architectures (Khan et al., 2022; Dehghani et al., 2023).
X Identity ResNet
10
b
(N 3, nM = [16, 64, 128, 256])
b
BatchNorm & Activate
AvgPool
R
Linear
SoftMax
x
X
X
3
* 3 12
Block (1, N)
112
Block (2, N)
n2
Block (2, N)
n4
n14
E
L
10
10
b
b
b
X
x x/s
Legend
rlp
Block (s, N)
71
11 n1 4
b
b
Residual Component
> R
X/S
R
X/s
>
R
X/S
+ x/5
Mg
Ho
12
11
*
TTb
1
171
11
b
b
b
b
Parantheses fed back themselves, ranging A from
X/5
1
21
minimum to maximum
value.
N-1
Residual Component
N -
X
>
R
X x/S
>
R
X/S
>
R
X/S
X x/5
11,
R7b
3
71
B1
A ==
Figure 26: Residual networks with identity mappings and full pre-activation (IdResNet) (He et al., 2016) offered improvements over the original ResNet architecture. These improvements, however, are often missing from implementations. By making the design of the improved model clear, neural circuit diagrams can motivate common packages to be updated. (See cell 13-15, Jupyter notebook:)
These cases show why clear architecture design is promising for enhancing computer vision research. Neural circuit diagrams, therefore, are in unique position to accelerate computer vision rescarch, motivating parallelization, task-appropriate architecture design, and further innovation of sophisticated architectures.
As examples neural circuit diagrams applied to computer vision architectures, have diagrammed the identity residual network architecture (He al., 2016) in Figure 26, which shows many innovations of ResNets not included in common implementations, as well as the UNet architecture (Ronneberger et al., 2015) in Figure 27, which shows how saving and loading variables may be displayed.
19

---

Published in Transactions on Machine Learning Research (12/2023)
X
Legend
UNet
x; 512 2
Cp
(with padding)
1
i= 0
The C1 values are saved to memory...
=
A
4
Double Convolution
64 2
else..
X
X
X
X
= 2
CA
Co
CA-1
*
CΛ
*
3
CA
R
>
X
x5
2
Down Scale Block
Max
C1
1
X5
X5
x
X
T
C4
*
C5
* 3
C5
22
C4
3
>
R
>R
Up Scale Block
X5-
X
*
X4
C5
2cs
* 3
C5
* 3
C5 A
22
C4 A
C1
>
R
>
R
X
C5-A
Concatenation is view over tuple. It is linear, can be simultaneously broadcast with other linear operations. In this case, diagonalization the axes.
C5-1
the C2 values are loaded from memory.
A
X
X
*
V
Figure 27: The UNet architecture (Ronneberger et al., 2015) forms the basis of probabilistic diffusion models, state-of the-art image generation tools (Rombach et al., 2022). UNets rearrange data in intricate ways, which we can show with neural circuit diagrams. Note that in this diagram we have modified the UNet architecture to pad the input of convolution layers. To get the original UNet architecture, the X, values can be further distinguished as X the sizes of which can be added to the legend. (See cell 16, Jupyter notebook.)
Architectures often comprise sub-components, which we show as blocks that accept configurations. This is analogous to classes or functions that may appear code. The code associated with this work implements these algorithms guided by the blocks from the diagrams.
20

---

Published in Transactions on Machine Learning Research (12/2023)
3.5 Vision Transformer
Neural circuit diagrams reveal the degrees of freedom of architectures, motivating experimentation and innovation. A case study that reveals this is the vision transformer, which brings together many of the cases we have already covered. Its explanations (Khan et al., 2022, See Figure 2) suffer from the same issues as explanations of the original transformer (sce Section 1.2), made worse by even more axes being present.
With neural circuit diagrams, visual altention mechanisms are as simple as replacing the and F axes in Figure 21 with tandem y and axes and setting h 1. As 1 x convolutions are simply the identity map, Conv(u, [1]) v, broadcasting linear map R - Rx for each of pixels 1 l-convolution. This leaves with Figure 28 for visual attention mechanism.
Visual Attention
1 x 1 convolution tensors the identity, be freely added.
Scaled Dot- -roduct
Attention
2
C
Q
k
*
-1/2
X
X
2
c
y
1
V
Figure 28:
Using neural circuit diagrams, visual attention (Dosovitskly et al., 2021) is shown to be simple modification of multi-head attention (See Figure 21, Figure 16, cell 17, Jupyter notebook.)
This highly suggestive diagram calls us experiment with the convolutions' stride, dilation, and kernel sizes, potentially streamlining models. The diagram clarifies how to implement multi-head visual attention with h 1, especially using cinsum similar to Figure 16. Additionally, y does not need to match X. We could have y be image data, and T be textual data without convolutions.
This case study shows how neural circuit diagrams reveal the degrees of freedom architectures and, therefore, motivate innovation while being precise in how algorithms should be implemented.
3.6 Differentiation: A Clear Improvement over Prior Methods
leave the most mathematically dense part of this work for last. Neural circuit diagrams intend to be used for the communication, implementation, tinkering, and analysis of architectures. These aims appeal to distinct audiences, and each should conceptualize neural circuit diagrams differently. The theoretical study of deep lcarning models requires understanding how individual components are composod into modcls and how properties scale during composition. Neural circuit diagrams are highly composed systems (see Figure 7) and thus provide framework for studying composition. They have underlying category, which is not the focus of this work.
Differentiation is an example of a property that is agreeable under composition. Differentiation is key to understanding information flows through architectures (He et al., 2016). The chain rule relates the derivative of composed functions to the composition of their derivatives and, therefore, provides case study of how studying composition allows models be understood. This analysis, however, is hampered by the fact that symbolically expressing the chain rule has quadratic length complexity relative to the number of composed
21

---

Published in Transactions on Machine Learning Research (12/2023)
functions.
n'(x) = h'()) h)())) = ( (f (fogo1)(s)) (( ()) o))())
This issue of symbolic methods proliferating symbols keep track of relationships between objects was noted in the introduction. To understand how differentiation composed and encourage more innovations like that of identity ResNets, which used differentiation to understand data flows (He et al., 2016), need graphical differentiation method.
Some graphical methods have been developed and applied to understanding differentiation in the context of deep learning, drawing on monoidal string diagrams from category theory (Shicbler et al., 2021; Cockett et al., 2019). As linearity cannot be completely ensured, these graphical methods are Cartesian, not expressing the details of axes. Other graphical approaches to neural networks could not incorporate differentiation, showing the signíficance of neural circuit diagrams being able to incorporate differentiation (Xu & Maruyama, 2022).
Differentiation, however, has key linear properties. Transposing disserentiation is very important. These prior graphical methods require redefining differentiation for each transpose, making the relationships between these forms unclcar. By detailing tensors and Cartesian products, our graphical presentation can show these linear relationships clearly. While drawing on their many theoretical contributions (Shicbler et al., 2021; Cockett et al., 2019), this work provides a significant. advantage over these previous works.
In addition theoretical understanding, clearly expressing differentiation is key to efficient implementations. Mathematically equivalent algorithms may have different time memory complexities. The rules of linear algebra we have developed (see Figure 18) allow mathematically equivalent. algorithms to be rearranged into more time or memory-efficiet forms. To show the potential of neural circuit diagrams, we focus on backpropagation and analyze its time and memory complexity with neural circuit diagrams.
3.6.1 Modeling Differentiation
To model differentiation, consider a once different function F R - RR.. It has a Jacobian which assigns to every point in R" axb tensor that describes its derivative, JF R R questions, and JF answers how much function responds to an infinitesimal change. The questions we ask JF are where is the change happening (R" input), how much is it changing by (IR output axis), and which direction are we moving in (R" output axis). Inner products over the output axes "ask" these questions. The chain rule can be defined with respect to the Jacobian and is diagrammed in Figure 29.
bxa Functions answer
a
(
=
Σ
a
G
a
Fb
dx
oxb
dx
x
F(x)
Figure 29: The chain rule expressed symbolically with index notation, and with neural circuit diagrams.
F
b
][F; G| a =
JG
a
b
JF
4
This expression convoluted, and will struggle scale. Instead, we transpose JF into the forward derivative as per Cockett et (2019)'s definition This form is more agrecable for the chain rule, and the first transpose we employ.
22

---

Published in Transactions on Machine Learning Research (12/2023)
a
a b
JF
DF b
Figure 30: Definition of the forward derivative, and how functions com- pose under it.
a
A
F b
D[F; G| c
8
OF b dG
This naturally scales with depth. Furthermore, can define (, D_) functor, composition preserving map, from once differentiable functions F IR" Rº to (F, DF) R RR" H x g² (Fong et al., 2019; Cruttwell et al., 2021; Cockett al., 2019). Per the chain rule, (, [DD D This is shown in Figure 31.
A diagram composed of vertical sections,
Fbbel b F - = a F
We have composition preserving map, the D_) functor, that maps vertical sections to vertical sections, implementing the chain rule.
Figure 31:
A functor map from vertical sections to vertical sections that preserves composition. We define derivative chain functor we can apply to any diagram once-differentiblle functions.
F
G
F
DF
DG
D[F; GJ
This composes clegantly. However, when optimizing an algorithm, we are not interested in how much known infinitesimal change will alter an output. Rather, given some target change in outpul, we are interested in which direction will best achieve it. We can do this by calculating the change in the output for each element of in parallel, effectively running the algorithm multiple times. This done by applying the unit and broadcasting. Furthermore, we sum the infinitesimal change over some target R value. The inner product does this. For algorithm F; L,, where L is loss function to I can do optimization according to Figure 32.
The inputs the algorithm the desired change L (-1 to minimize L). and the location of (the current parameters.)
Figure 32: We turn a small chain rule expression an optimization fune- tion by applying the inner product over the target direction and deriva- tive output. An inner product over an axis of length 1 is just multiplica- tion. Using the unit, we run this al- gorithm for every input degree of free- dom, broadcasting over the axis.
The output is the change X required to achieve the change in f.
= 8X
DL
However, the forward derivative has large time complexity. A linear function gives matrix multiplication. Therefore, a linear map f : R -R R applied onto ge will require a x operations. In general, broadcasting multiplics time and memory complexity. The memory usage of an algorithm is related to the number of elements it stores at any step in the algorithm. We use these tricks to analyze the order of the time and space complexity for the above process.
23

---

Published in Transactions on Machine Learning Research (12/2023)
Space Complexity
1 ++a+ +a++++0 aaxa
1 x
1+b+bxa
JC
Figure 33: An analysis of the space and time complexity of the naive op- timization algorithm.
X
a
F
DF
DL
Time Complexity (a x b)
(b
We observe that this has a high time complexity, quadratic with respect to the size of X. In practice, we avoid the forward derivative, also called the Jacobian-vector product or JVP, in favor of the reverse derivative, or VJP, which more directly implements the above process. We define it relation to the Jacobian and forward derivative in Figure 34. In Figure 36, algorithm in terms of the forward derivative and show the far lower memory and time complexity required.
rules of linear algebra to re-express the optimization
Definition of the forward derivative,
b JF
DF
b
JF
a
DF
a
Naturality,
Unit Unit-inner product identity,
a
JF
JF
RF
JF
Figure 34: The definition of the forward and reverse derivative with respect to the Jacobian. This aligus with the Jacobian-vector product and the vector- Jacobian product, respectively.
Figure 35: A full expression of how the unit and the forward derivative give the Jacobian. This demonstrates how linear algebra principles il- lustrate the relationships between different forms of the derivative.
f
gf
Figure 36: Using the previously developed linear algebra rules (see Figure 18, we rearrange our optimization algorithm to use the reverse instead of the for- ward derivative. The diagrams then reveal the computational advantages of the new algorithm, called backpropagation.
DFb
DL
b
DL
b
+ b +
Space Complexity
g.
8
F
RL
b
R
b
RF à
DFb
Time Complexity
1xb
bxa
4 Conclusion
Neural circuit diagrams are method of addressing the lingering problem of unclear communication in the deep learning community. My introduction showed why this is concern and argued why a system of axis wires and dashed lines needed, solving the challenge of reconciling the detail of tensor axes and the
24

---

Published in Transactions on Machine Learning Research (12/2023)
flexibility of tuples. This work covered a host of architectures to breed familiarity, encourage adoption, and evidence the utility of neural circuit diagrams for understanding and implementing models.
Neural circuit diagrams are appealing to diverse users, from students first learning neural networks to specialized theoretical researchers investigating their mathematical foundations. This leads to immense future potential. Future work can see neural circuit diagrams explained concise and accessible manner for lay audience. More diagrams can be modelled and formal standards developed. Finally, their mathematical foundation can be more fully expressed. The underlying category theory structure can be fully investigated (Abbott, 2023, Chapter 3), allowing models to incorporate probabilistic functions (Perrone, 2022; Fritz et al., 2023), additional data types, or even quantum circuits.
Acknowledgements
Mathcha was used to write equations and draw diagrams. The Harvard NLP annotated transformer was invaluable for drawing Figure 22. I thank the anonymous TMLR reviewers for providing useful feedback the paper throughout its many rewrites, and my supervisor Dr Yoshihiro Maruyama for his support, and pointing me towards applied category theory for machine learning. This work was supported by JST (JPMJMS2033-002; JPMJFR206P).
References
Vincent Abbott. Robust Diagrams for Deep Learning Architectures: Applications and Theory. Thesis, The Australian National University, Canberra, October 2023.
Honours
Steve Awodey. Category Theory. Oxford University Press, Inc., USA, 2nd odition, July 2010. ISBN 978-0- 19-923718-0.
Lei Jimmy Ba, Jamie Ryan Kiros, and GeofTrey E. Hinton. Layer Normalization. CoRR, abs/ 1607.06450, 2016. URL http://arxiv ht////rxx...o/aab//...////////// 1e//111677..44.00 arXiv: 1607.06450.
John C. Baez and Mike Stay. Physics, Topology, Logic and Computation: A Rosetta Stone. volume 813, pp. 95 172. 2010. doi: 10.1007/978-3-6222 3-18--33-----11 12821-9 2. URL http://arxivvoor haagg//bs//// 03400 #eXiiv0903.0340 (quant-ph].
Jacob Biamonte and Ville Bergholm. Tensor Networks a Nutshell, July 2017. URL http://arxiv..o org/abs) 1708.00006. 3XXi:11708000000 (cond-mat, physics-gr-qe, physieschep-th, physicsmatt--hh, physies:quant- ph].
Michelle A. Borkin, Zoya Bylinskii, Nam Wook Kim, Constance May Bainbridge, Chelsea Yeh, Daniel Borkin, Hanspeter Pfister, and Aude Oliva. Beyond Memorability: Visualization Recognition and Recall. IEEE Transactions on Visualization and Computer Graphics, 22(1):519-528 January 2016. ISSN 1941- 0506. doi: (1110//VVG.2205.226 Computer Graphics.
Conference Name: TEEE Transactions on Visualization and
David Chiang, Alexander M. Rush, and Boaz Barak. Named Tensor Notation, January 2023. URL http:
hhrxiv.org/ab///1022222 hbs/a122..21191/ 1rX1v:2102.13196 [cs).
Robin Cockett, Geoffrey Cruttwell, Jonathan Gallagher, Jean-Simon Pacaud Lemay, Benjamin MacAdam, Gordon Plotkin, and Dorette Pronk. Reverse derivative categories, October 2019. URL http://arxiv.org
abs/1910.07005.. 1rXiv:1910.07065 [cs, math].
G. H. Cwll, Bruno Gavranović, Neil Ghani, Paul Wilson, and Fabio Zanasi. Categorical Foundations of Gradient-Basd Learning, July 2021. URL http:///rxiv. wwwpssivvooo//aaa abs///b////a waabb//210300090 arXiv:2103.0193 [cs, math].
Geoffrey S. H. Cruttwell, Bruno Gavranovic, Neil Ghani, Paul W. Wilson, and Fabio Zanasi. Categorical Foundations Gradient-Baan Learning. In Ilya Sergey (ed.), Programming Languages and Systems 31st European Symposium Programming, ESOP 2022, Held Part of the Euroopaan Joint Conferences
25

---

Published in Transactions on Machine Learning Research (12/2023)
Theory and Practice of Software, ETAPS 2022, Munich, Germany, April 2-7, 2022, Proceedings, volume 13240 of Lecture Notes in Computer Science, pp. 1-28. Springer, 2022. doi: 10.1007 1110077/778-30000-93333---- 1. URL https://doi.orr/III 101-10077778--0000-000
1.
Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim Alabdulmohsin, Avital Oliver, Piotr Padlewski, Alexey Gritsenko, Mario Lučié, and Neil Houlsby. Patch n' Pack: NaViT, Vision Transformer for any 22000/06300...0 arXiv:2307.0630 [cs].
Aspect Ratio and Resolution, July 2023. CRL
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In 9th Inter- national Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 8-7, 2021. OpenReview.net 2021. URL https:/ http:////pprreeeeereeeee wetpriiwwweeettoommddd-i/i-- Dieerummmdd-YYbaaaTTT.
Chris Drummond. Replicability is not reproducibility: Nor is it good science. Proceedings of the Evaluation Methods for Machine Learning Workshop at the 26th ICML, Jamuary 2009.
Brendan Fong and David I. Spivak. An Invitation to Applied Category Theory: Seven Sketches in Com- positionality. Cambridge University Press, 1 edition, July 2019. ISBN 978-1-108-66880- 978-1-108- 48229-5 978-1-108-71182-1 doi: 10.1017 00.1017/9811008680044. URL https:/ https://www.coomm www.c/ebbbrdgg.ogg/ooor////////e product/ identifier 10000000046////////000000
Brendan Fong, David I. Spivak, and Rémy Tuyéras. Backprop Functor: A compositional perspective supervised learning. In 34th Annual ACM//EER Symposium on Logic in Computer Science LICS 2019, Vancouver, BC, Canada, June 24-27, 2019, pp. 1 13. IEEE 2019. doi: (L.100//LCS..201.8775577
Tobias Fritz, Tomás Gonda, Paolo Perrone, and Eigil Fjeldgren Rischel. Representable Markov Categories and Comparison of Statistical Experiments in Categorical Probability. Theoretical Computer Science, 961: 113896, June 2023. ISSN 03043975. doi: 10.1016/j.tess 1039.1 URL http://arxiv.org/abbb abb/2000..7 ar arXiv:2010.07416 [cs, math, stat].
Tan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
John Hayes and Diana Bajzek. Understanding and Reducing the Knowledge Effect:: Implications for Writers. Written Communication, 25:104 118, January 2008.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. CoRR, abs/1512.03333 2015. URL http://arxiv.orr
haabb///112..3335.. arXiv: 1512.03385.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity Mappings in Deep Residual Networks. In Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling (eds.), Computer Vision ECCV 2016 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV, volume 9908 of Lecture Notes in Computer Science, pp. 330-645, Springer, 2016. doi: 10.00779976-33319--449---0 38.
38. URL httpss// wtwgg/..100/9988----0
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising Diffusion Probabilistic Models. In Hugo Larochelle, Aurelio Ranzato, Raia Hadsell, Maria-Florina Balean, and Hsuan-Tien Lin (eds.), Advances in Newral Information Processing Systems 33: Arnual Conference on Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020. URL https://proceedingsneeeipss Abstract.html
(paper/2020)
Sergey loffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. CoRR, abs/11002.05 2015. URL http://wwwwworgbbbs/////aa 1502.03167.
arXiv:
Sayash Kapoor and Arvind Narayanan. Leakage and the Reproducibility Crisis in ML-based Science, July arXiv:2207.07048 Jos, stat).
2022. URL http://arriv.org/aa//////
26

---

Published in Transactions on Machine Learning Research (12/2023)
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in Vision: A Survey. ACM Compuling Surveys, 54(105):1-41, January 2022. ISSN 0360-0300, 1557-7341. doi: 10.1145/350524. URL
Alex Krizhovsky, Πya Sutskever, and Gcoffrey E Hinton. ImageNct classification with deep convolutional neural networks. Pereira, C.J. Burges, L. Bottou, and K.Q. Weinberger (eds.), Advances in neural information processing systems, volume 25. Curran Associates, Inc., 2012. URL https: httpss//proceedings. perripsscc/paperr
Minhyeok Lee. GELU Activation Function in Deep Learning: A Comprehensive Mathematical Analysis and Performance, May 2023. URL http:/ http//aaxxvvvraababbe abt.abb//222222222 arXiv:2305.1207 [es].
Yuqing Li, Tao Luo, and Nung Kwan Yip. Towards an Understanding of Residual Networks Using Neu- ral Tangent Hierarchy (NTH). CSTAM Transactions on Applied Mathemalics, 3(4):692-70) June 2022. ISSN 2708-0560, 2708-0579. doi: 10.4208//coa (csiammmmm.S-002002 arXiv:2007.03714 (cs, math, stat].
URL http://aariiii /aa//aaxiiiorrr... wwt..r////b//000031....4
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A Survey of Transformers, June 2021. URL 220.b////00.05... arXiv:2106.04554 [cs].
Hanxiao Liu, Zihang Dai, David R. So, and Quoc V. Le. Pay Attention to MLPs, June 2021. URL
http://arxiv.org. .ow.oorg//bbbb aataaas/...../ 22/222110000000 arXiv:2105.05050 [cs].
Wenjie Luo, Yujia Li, Raquel Urtasun, and Richard Zemel. Understanding the Effective Receptive Field in Deep Convolutional Neural Networks, January 2017. URL https://arxiv.or//abbss htgg/abb///70.04122222
José Meseguer and Ugo Montanari. Petri nets are monoids. Information and Computation, 88(2):105 155, October 1990. ISSN 0890-5401. doi: 11101//090055400(0000000000
URL https://www.sccencdddee
T. Murata. Petri nets: Properties, analysis and applications. Proceedings of the IEEE, 77(4):541 580, April 1989. ISSN 1558-2256 doi: 10.1109/5.2411
Conference Name: Proceedings of the IEEE.
Alex Nichol and Prafulla Dhariwal. Improved Denoising Diffusion Probabilistic Models, February 2021. URL http://axxiv.or/aaa www.o////////o/oo//er/e 221.2222222222222 arXiv:2102.09672 [es, stat].
Paolo Perrone. Markov Categories and Entropy. CoRR, abs/ 22t//2222.11799, 2022. doi: 10.48550 005550//RXXV22211 11719. URL https://doi.org http:/ddi..rr///0005//////////// 40550//ariiv.22222.1777 arXiv: 2212.
Mary Phuong and Marcus Hutter. Formal Algorithms for Transformers. CoRR, abs/2207.09233 2022. doi: .9928.9 arXiv: 2207.09238.
10.48550/ ARXIV www.2227.092233 URL https::/doo..or//1...
S. Pinker. The sense of style: The thinking person's guide to writing in the 21st century. Penguin Publishing Group, 2014. ISBN 978-0-698-177 17098-8 URL https:/ hases//boooll books. http:///bbooosgogl..ooooo www..o/l.........ooooooo///dd-aa
Edward RafT. A Step Toward Quantifying Independently Reproducible Machine Learning Research. In Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https:// proceedings. ppoceedings.ncurips.comm http:paaapppppa/aaaaaaaaaaaa fflcs/paper//0000
Alex Rogozhnikov. Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation. October 2021. URL /ee wttp://wwww...tooom///oooooo///o
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-Resolution Image Synthesis with Latent Diffusion Models, April 2022. URL http://arxiv.org/ wwwpsiiwwwrg/b/s////////oo 2212211222 arX2v:2112.10752 [cs].
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation. CoRR, abs/1505.04597, 2015. URL http://xxiiorrr/ee
hatp///abb/11555555555 1509.04597. arXiv: 1505.04597.
27

---

Published in Transactions on Machine Learning Research (12/2023)
Lee Ross, David Greene, and Pamela House. The "false consensus effect": An egocentric bias in social perception and attribution processes. Journal of Paperimental Social Psychology, 13(3):279- -01799--0,,,0 1977.
Sadoski. Impact of concreteness on comprehensibility, interest. Journal of Educational Psychology, 85(2): 291-304, 304, 1993.
Andrew M Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan D Tracey, and David D Cox. On the information bottleneck theory of deep learning. Journal of Statistical Mechamics: Theory and Experiment, 1014(12)), December 2019. ISSN 1742-5468. doi: (0.1088/1742-5468/
ab3985. URL hutps://iapsciince..op..pp
0ttriic///1..00//777222222 ab398..
Peter Selinger. A survey of graphical languages for monoidal categories, August 2009. URL https: /arxiv. arg/abs/0908.334T7
Dan Shiebler, Bruno Gavranovie, and Paul W. Wilson. Category Theory in Machine Learning. CoRR, abs/2106.07032, 2021. URL https://arxivvorg// wwrggabb/2106.000 2206.00002.. arXiv: 2106.07032.
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidluber. Highway Networks. CoRR, abs/1150500038/ 2015. URL http:://axiv.org/// 1505.00057.... arXiv: 1505.00387.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive Network: A Successor to Transformer for Large Language Models, August 2023. URL
/tw.b///220.oomm .86211. arXiv:2307.cm 200::222770821111 [cs].
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advonces in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-609999 2017. URL https:/ httpss//ppooeddigg.eeerrrgi 3355ee223557deeee lfbd00533 1aaa:3aa115555aaa-aaaaaaaa Abstract.tttutt
haper/2017)
Paul Wilson and Fabio Zanasi. Categories of Differentiable Polynomial Circuits for Machine Learning, May 2022. URL http:// http://aaxiv.oraa////////s/ httaab///2000000 /t////223..3333 06430. arXiv:2203.06430 |cs, math].
Tom Xu and Yoshihiro Maruyama. Neural String Diagrams: A Universal Modelling Language for Categorical Deep Learning. In Ben Goertzel, Matthew Iklé, and Alexey Potapov (eds.), Artificial General Intelligence, Lecture Notes in Computer Science, pp. 306 315, Cham, 2022. Springer International Publishing. ISBN 078-3-030-93355-44 doi: 10.1007/978-3-000 071077788---00-777----------- 32.
Yao Lei Xu, Kriton Konstantinidis, and Danilo P. Mandic. Graph Tensor Networks: An Intuitive Framework for Designing Largo-Scale Ncural Learning Systems on Multiple Domains. CoRR, 0t4//233313355 2023. arXiv.2303, 13565.
doi: 10.48550/ 118550/ARXIVV ARXIV 1306.. URL
arXiv: 1356.1
Matthew D. Zeiler, Dilip Krishnan, Graham W. Taylor, and Rob Fergus. Deconvolutional networks. In 2010 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pp. 2528 22338, San Francisco, CA, USA, June 2010. IEFE. ISBN 978-1-4244-6084-0. doi: 11100//CVPR.00005533300 URL
http://ieeexplor.eeee.oorg htw.eee.rrggooccmnn/////////1 55339577
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization, February 2017. URL arXiv:1611.03555 [cs].
A Jupyter Notebook (see: weabott/Neura--Ciruut-DDaaagga
28

---

Published in Transactions on Machine Learning Research (12/2023)
tensor([6, 2, 8])
Inport torch
typlig import functorch
Figure 9: Subtensors
inport itertools
symbod
2.3 Tensors
We diagrams tensors, which can be vertically and horizantally decomposed.
Different orders of access give the sare result.
We data repruenting with
R
up randon (5 7)
2
daagiam drowing
Xab randhLr [b])
Show that all pairs of indexes give the same result in seertools.roouuct(eee assert Xab[la, jb] Kab[ia, :][]b]
R460
R203
8
3
range(b)):
assert Xah[1a, jb] -- Xab[:, [b][ta]
2.3.2 Broadcasting
diagran Function 3, constructed out of f: 6 3 and 8: 3
Figure 10: Broadcasting
outputs represent generic functions, and how diagrans relate to code. Corch.Tensor
Ar operation on over digll lenght applic partls oue exparate nows beaadcamting The lis operation G: R R R Rt
+(x0 T, xi T::
... f: 3 ...
assert stzelze torch (tech,51ze[[4,2]]
assert si.size() sir)))
return (33,)))
def g(x0
2: 3, 3 +> assert x8.size() tir))))ize(())
equil Graphicaly, we see that broddcattin allows indeees poss over exgnssions, Note how the Auprectine holds for 2 reploned by any valid incec (0,12,,3
Graadmating can formaly
ideass nignt uide
si.size() Sizeh, 31)
return (1, def h(xe T, xt : T, T::
3
... h: 2, 2
x8.size()
corch.Size((33)
assert xl.size() =* torch. 2])
assert x2.size() toe)
[3), [2], [4], [3]
g(x9, f(x1,x2))
T torch. Tensor
h(t))) rar 2}), tor)h
function
bj
def G(Xa: T) >> T:
6: return sun(Xa**2) tarch.anes(b)
tonsar[[[8.68]) B.6853]])
Ne could bootstrup definition of broadcasting, Note that we are using spaces to indicate tensoring. tupling, which is 1ine with standard notation while writing code. deff Ge(Xac: T) >> .en Ybc torch.zeros($
2,3.1 Indexes
Figure Indeexes
We express aultersur octrctions grabbing Al2, 1, by indox applted on the relevans ods
for i in range(c[e]):
Ybc[=.je] G(Xac[:.je])
3
Or use PyTorch command torch.vnap(6, -1,
Extracting process are faniliar with. Consider, It A (4 3) tensor
We feed randon input, and see whothor applying an or after
the result.
table torch. (0,)
Xac torch.rand(a c) for Je range(c(@]):
table[2,
row
alt.ccloooeeeG//a((tt.))) Gc(Xac)[: assort torch, torch..aliccooee@(aa((...... (:xXe)
29

---

Published in Transactions on Machine Learning Research (12/2023)
This shows how definition broadcasting lines with that used by Pytorch vnap.
Elenentwise operations implenented usual ie def f(x):
Figure 11: Inner Broadcasting
"F 1 1" **
gett XII of action of can on indueed to R x dfteore soye Theefon comration (R Kt
Wo broadcast an elonentwise operation,
R R gi fanee
fs torch.vnap(f))
x loner Broafccaiing
Xa tarch.
!I
range(a(e]):
And that it aligns with the index before Index after franework.
assert torchaallcloee((((a(t),). fs(x)]
elementwise operations Implied,
Inser Geppicclll, beopfeasting
special implomentation is nendad.
tow oppresiont replaned by any idec (0, 1,233)
exprescons
assert torch.allclose(((((())) f(Xa)[1))
2.4 Linearity
4
Hom, the iider a sccompan by idebbiy on the stappdd vonsblle
2.4.2 Implementing Linearity and Common Operations
3, b, c, d [3], [2], [4], Corch.Tenser
Figure 17: Multi-head Attention Einsum
Impiomentation using
it We have sone Function which can be Inner broadcast,
Local menery contains,
H(Xa: T, Xd: T)
..* www.stttrcc.gggt((aaar.....
...
K. siness sinsur 2.k KIR Implicit outar product and diagonalkns,
return tonch.
torch.sun(torch. (Xr ** 2)) + torch.ones(b)
Troee product.
strops.elnniitt
bootstrap inner broadcasting, def Hc0(Xca: T, xd T) T:
<0 H: c a,a, c.d
Impiomentation ssing eirsun (with kmuhanecin braddrattngg finction))
defined a, c, []] arrays. Yeb zerosc for ic in range(c(#)):
Locad contains,
D: kin x strops sinsunt K,
Yeb[te, :] H(Xca(ie, :], xd)
return Yeb
But veap offers clear way doing it, Hse torch.wnap(H, (8, None),
Import nath
Import einops
y, 5,
can show this satisfles Definition 2.14 by, sancch.randcc.
tonch h]) k torch. k, h])
xd torch. rand(())
for ic in range(<(o))):
Local menory contains,
assert torchaallcdosss(ieeek.. xd)(is, :]. H(Xca[ic, assert forch.allccesee(ts@@ke.. xd)(ie, :].
$ Q: y k a K: k.
Outer products, transposes, inner products, and diagonalization reduce
expressions.
H(Xcalic, tI, Xd))
Transpose K,
K einops.einsunnk, 'x k k. b.)
product and diagonalize, cinops.cinsun(Q,
h')
Figure Elementwise operations
Inner product,
einops.einsun(X, "
Broadcasting with I-length ads leaves data unchanged, 1-length axes can be freely introduced and removad with arrows;
# Scale,
X X (ot
Operation Value
Element-wise
f:R R
torch.rand([V,
K tonch k, h])
menory contains,
Q: k k.h
x4x3 3 Broadcasting
3
x einops. K, kb, sort(k)
Starting function values, apply every value in data by broadcasting.
2.4.3 Linear Algebra
30

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 18: Graphical Linear Algebra
identity.
x
The first step is an outer product with the unit, outerUnit lanbda Xh: es(Xb, torch.eye(b(0]), 'be, bi b2 -> b1 b2")
Ausotieted
The next is inner product over the first two axes, dotouter larbda Xbob: es 'b8 b1 -> bi') Applying both these *should* be the identity, and hence leave any input unchanged. Xb ranc(h assert torch.aliclose
Mdrniity
The inc and the ins product amanged
G
H
this noreer gan taa Hienct
H
2 G
G
dottutee((uuternnn((()))))
H
H
H
Therefore, confidently the expressions in Figure 18 to manipulate expressions.
Srrultsseutll brndeading live spe332rci1a1 sater woded, maring wue prolats car le skites.
G
G
K
K
3.1 Basic Multi-Layer Perceptron
H
Figure 19: Implementing Basic Multi-Layer Perceptron
will exercise implementing of these equivalences.
peoc pape eecogrtser
# the reader can follow this exercise to get better sense linear functions can implenented, and how different forms equivalent.
(1ab dest Dnagrhegg Luarge unaucityy
Lofe Rilie Flilincs
BROC deppettr
Basic Image Recagniser for Digits
a, b, c, d [3], [4], [5), [3]
eiviloel 1ita, 241.
$12
912
10
# we will be using this function ta lott es = einops.einsun
Lines
Y
with www.c//uuaaaaa..
b.c
F F_natrix tarch.rand(a b
Smport tarch.nn
exericse we show
linear nap
Basic Inage Recogniser
<> b can be transposed in two ways.
This is close copy of an introductory PyTorch
Either, we can broadcast, or take an outer product.
tutorial:
these are the sane.
https: pttps://ppt http://yytttcoorteeeee/ htwp://www.uuuuosseee//ee/e//sea
# Transposing by broadcasting
1. tutoriialll.... html
BesicInagsRecconiseerr (nduMedule):
[_aunc(Xa:: T):
def init__seeff;
..* F: b.c
super()._init.
return ex.a, 'a,a b c->b ets)
self.flatten nn.Flatten)
b.c
@e1f.1inear_rell_ce Sequentnti
__broadcast torch. func, 8)
nn.Linear(28*28, 512),
nn.RelU(),
# We then reduce it, as in the diagran, b
nn. (512, 512), ReLU(), nn. Linein (512, 18),
def b.broadcast__ee txans 1):
(b () (.b c): b ->. Xbbc F_broddcast(Xbb) ct)
forward(self, solf.flatton(s hacprreeeeeuusaaeeeeeeeeeeeaaeee
(Xboc,
x self. selff_ineerrrelu_tteee y_pred nn.Softnaxxx) y_pred
Transpoing by Hinerrity
take outer product of Id(b) and follow with inner product.
ny BessicInageRecognisen Rasicdinegstecognisse)) my sassccimgeeeecognissee..... aasssmmmeeeeggsseeerrrrrrrrrrrte htwpreeeeeeeeeeetcoocooonrrreeae
This gives us, ..uuerrroduuct es(torch.eye(b(e)), f_natrix," b1, b2 c->b8 b1 b2 c²,)
it Think of this as Id(b) Ft b8 a -> b1 b2 arranged into an associated
b2 e tensor,
Saftmax(
then take the inner product. This gives (b natrix, which can be used for a (b +> c) map.
dim=tensor([ 0.0150, 03,038,, 0.1395, -8.0558, 6.8824, -B.8613, -8.0163, 0.0134, 0.8577, -8.8624]], grad_Fn- Fn-
F. linear transsoseeee outerprorrrr,,, B c->b
caddmetcckwrdee.)
We contend that these are the sane,
torch. rand(b
assert tlrchsellclose(
3.2 Neural Circuit Diagrams for the Transformer Architecture
F_booadcast transaassssssaaaaaa transpose, b c())
es F. 2inear_transposse
Furthermore, lets prove the unit-inner product
31

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 20: Scaled Dot-Product Attention
k h')
return
Sclle PrePect Amartion
Siediion fram
We ixplement component
network
nodel.
This is necessary when there are hold, learned components that need to initialized. class multimeadtttaanionni Nodule))
M X,
S.0Maa
Multi-Heed attention has various settings, which become variables
for initializer.
Note, need acconodate batches, the capture additional axes.
def _init _seif_ h):
super()._
self.n, self.k,
k, h
do the algorithn step by step, def ScaleebocProdiccttteeeeeooa kt
boldface, learned corpanents # Note how they bind axes un want to split, which we do later with einops.
yk, xk, xk -> yk
klength [--1
self.ig mn.Linear(m, kth, False)
Transpese
self. Ln.Linear kth, False) self.Lv in. kth, False) self. mn.Linear(k"h, i, False)
simpss,eissun(k, "...
x')
# Matrix Multiply / Inner Product
x einops.einsun(G, k, .... y K, ... k ....
x*)
# Scale
# We have endogenous data (Fyn) and external injacted data (Xxm)
x nath.sqptt(klmntt)
SoftMax
def forward(self,, Eyn, Xxn):
n, +> # We first generate query, key, and value
Matrix Hultiply/ / Inner Product
x Sinops.einsumf/ v, "... y X,
vectors.
Linear layers autonetically broadcast.
return
However, the and axes are bound. deFine unbinder handle outputs,
Alternatively, we can sirultaneously broadcast linear functions.
unbind lanbda x: rearassne '... (k
def scllebortroddccttttntiiooig k: v: T) yk, xk, xk >> yk
h)->... k h', b-self.h)
q Lmbind(ae2f.L((Eym) umbind(self.uk(xxm))
klength s.si
Product and Scale
unbind(self.Lw((xa)))
x einops.einsunce, k, "... y k, x ->>>.....
x*)
k, to Nulti-Head
and
innor product Attention
torch.nn, wow.an.otaaxx111 / math. (preennggg))))
MuitHHesdoorrddcccttteettnd k, v)
# Final Inner Product
shops.eeisuunc V, .... y X, ...
Rebind feed the final learned layer, rearrs.na.,,, "... k h-> (k h)',
k*)
return
heself.h) return self.Lo(o)
Figure 21: Multi- Head Attention
Now we can run it on fake data;
y, X, 9, jG, heads [28], [22], [128], [16], 4 Intornal Data Eyn torch.rand(y n) External Data
Lownee By connecte legts geneue
Thes gemaat soadd dais
kess ncter steer head
Liner
Dos-rodut
torch.rand(x
nha (tt.tt/eeettetttttoorrll..0eeeed
Copy
assert
sim).size()) == y
copy
(Bigoraloo)
3.4 Computer Vision
During nutri rubipiiaiae, dapoolossi
fell cve layr
Here, really understand why splitting diagrams into "fenced off blocks aids implementation. addition making diagrams
We be implementing this algorithm. This shows us how we go from diagrams to implementations, and begins to give an idea of how organized diagrams leads to organized code.
understand patterns more blocks indicate structured and organized.
Figure 26: Identity Residual Network
def wuuttteeeedtrrddutttttttttooooo
T)
T:
ykh, xkh, xkh +> ykh
Identity ReaNet
klongth [.s]
3. [16.64 25K0
einsps.ei "... k. x.k h')
R
Lince
SofMae
X wtt.mmn.oottaxx22666 / soth-sgrt((1eengt))
Block (,M)
22
Block
(2
18
einops.einsun(x, '... y h, ...
32

---

Published in Transactions on Machine Learning Research (12/2023)
(ob, nb, padding=1), NorsActivate(nb), nn. (nnv n1, 1)
For Figure 26, every fenced off region is its own nodule.
# Batch norn and then activate is ropeated notif, class NorrActivate( (ww.eeeeenn..euetiialll.. def _niit_(eeff nf,, Batth, Activation:an. RelU): 1t____ntt_____(((______ Activation())
Residual connections are a repeated pattern in the diagran. So, we are notivated to encapsulate then as seperate nodule.
def size to_striing(str)) return (aan(mmp(str,iist((zzz)))))
class Residuallomeccionnmm.... (
def __nni___seff, nainline nn.Nodule, connection
The Identity ResNet nanageable sequence of components.
down
nn.Module
None) None:
super()._inft
class IdentityResMee(nn.B (nw.Seuunntiil))
self.main mainline
def _nnit_ N+3, (,,,[16664,128,235
self.secondary mn.Identity() if connection
y=10):
None connection
(upor():_intt___
def forward(self, return self.nein(x) seconaeco)))
cn.conv2d(3, n_mu[0], ,, 3, padding=1),
Black(1, N, n_mu[0], n_nu(1)),
Block(2, N, n_nu[1], n_nu(2]),
Block(2, N, n_nu[2], n_nu[3]), kornActivate(r_m[[]] m..ddpptivsseooool2((..
standard inago processing algorithn inputs shaped b
Fn.Ftatten),,
b, c, hM [3], [3], [16, 16]
(iinear(n_nn[3))) y), mm.ootrrax-11)
idresnet IdentityResNet() Xbchw (onch hw)
see the size naintained, 11st(1dresnet. (xchar .izee() [19]
The Block can be defined in a seperate madel, keeping the code manageable and closely connected to the diagram.
Lopord
The UNee is a more complicated algorithm than residual networks, The "fenced off" sections help keep our code organized. Diagrams streamline implementation, and helps keep code organized.
(.M)
Becshel Conserest
>++
31x
Figure 27: The UNet architecture
Puzzthere are red to
Ris
LiNet (with padding)
Legend
dlaatton sanding fom.
= +
adar
to stterne to manzeam
M-1-
Rendaal Component
>x
Double Conneiation
canol
The Cj valoos
1
=1=0
dbc.
61.
C2
>
1
Dl Mack
Scale
# We then fallow how diagrans define each "black" class Block((n.Sqquntiaal)
iintt (self, N, ne, ni):
and n1 inputs the initializen
implicit fron having then in the donain and codorain in the dlagran. ...
>R>
>R
// (uper()__n1t)
txple. linezz
sinallanosale beandess with opecstione disgsrallarion ove & aore
Ube Ci Wilre amo lusded bott tenscess
korsactivatorne), Residuualcameeciinn
Sequuentiaa((
mn.Conv2d(ne, nb, 1, s),
Nornactivatennb),
Cn.Con22d(bb, nb, 3, padding=1), (ornActivateenb), nn.Conv2d(nb, n1, 1),
Me that double convolution where numbers of channels change nepeated notif,
$ Me denate the input with c0 and output with cl. This can also be done for subsequent nenbers of an iteration.
(onv n1, s),
When go down Iteration have the input be cl and the output co. class DoubleConoolutionmm (tt.ionnnmmmmeetiiaa) def iiit (self, cB, c1, Activationann ReLU)) (oper()__ini____
etc. nay
Residdualcnnnection
nn. Sequentiit
Nornactivatenn2),
mm.Conv2d(n1, nb, (ornactivataaab),
nn. cn.Convdd(c8, c1, 3, padding-1),
33

---

Published in Transactions on Machine Learning Research (12/2023)
Activation(),
adapt our This is good exercise how neural circuit diagrams allow code to be easily adapted for new modalities.
apply case.
conv2deee,, cl, 3, padding-1), Activation(),
Figure 28: Visual Attention
The nodel specified for very specific number
layers,
will nake very flexible.
class Uket(nn.Nodulee
1x1 comvolation taners are the ideniiyy ftnely
def _init__sseff, (self, y-2):
DottProduct
super init
Super()._nnit___
Attension
the channel sizes;
c [1 1f == 0 else 64 ** 1 for In
ra))
Saving and loading fron nenory means нe can not use single, sequential
Y
V
D.C
Set and initialize the components; self. DownScaleBlocks (oonscaeelooo((()))(([]]]
for in range(8,4)
class www.itteentionnnnModuuea
# Noto how this initates the lambda openators
init [sel_ c, k, 1, stride 2):
diagran.
(nirt)
boublecomwolutionn(((s)] c[5])
self self.niedlelpscaae nn..covtransposeedc[[]]
gives kernel size, which we
c[4], 2,
adjustable.
self upscceeeeesss
self.c, self.k, self.h, self.u G k, heads,
upscalefloc<(((5-1))..E)
kernel
for in range(1,4)
all holdface, learned corpanents
# Note hou standard components nay not have axes bound in
self Fialloootooooiionnn co.onnvdde[1)), y)
way diagrans. requires
forward(self,, x):
rearrange
cLanbdas
# using the einops package.
for dsb in Saf.DoonSScaleiockss
X, cLanbda dsb(x)
learned layers form convolutions self.Cq Conv2dc,, heads, kennel,
cLanbdas apppndaas.ppend(Laattbbb
self.
stride)
self. niddlelpsscalee
self.ck cn.oond, heads, kernel, stride)
usb uolfaasscals:
clanbda clanhdas.
self.Cv cn.oonzd heads, kennel,
x usb(x, cLambda)
stride)
self. Finnlc(ooolllliion(())
self.Co Co.CTrrraansp
k heads, kenell,
class DownScaleBlock (ww.eeiioooccoMMMMooeeeee Module):
stride)
def init (self, c0, c1) None:
super(). super()._init__ init_
Defined previously, closely follows diagras. def multseeaddtrrraccttttttttnooooe q: T,
self. donvouutooooootiiionn Douh noutletonvoluttonnce,,,
c1)
Vt 1) >>
self, downScaler mm.xxxoo12d(2,, 2, 1) forlard(self,, x):
... ykh, xkh, ykh
klength k.size()[-2]
cLanbda self. doublecoowolurion(x)
x einops.einsun(e, k, '... y h, x.k.
2olF.dowSSaler(cL..n
y x h')
return x, cLambda
Son.h.an.oottaax((22 nath. sgtte(kggg)))))
class UpScaleelock(mm Module): c1, c0) >> None:
x Sanops.einsuntg V, '... y h, ... xkk
def
... y h')
super(). init
salf dol.lllloooooooitiioon
boubleeamoolutionn22aa.. c1)
# have endogenous data (EYc) and external injected (XXc)
self.upScalen ww..on/Tranppssssatteee2222
forward(self, cLanbda):
def forward(self, EcY, XcX):
Concatenation
channel axis
cY, cX -> CY
(din=1)
algorithn. Injects
torch.concat(y, clanbda, 1)
information
into ...
do..ubbbloooollttoon())) self.upScalorrx) return
# query, key, and value vectors.
# Me unbind the h axes which were produced by convolutions, in the namnal nuirsheeddeeProccccttteeeettt
unbind lanbda x: eenops. 'N h) (H H) h', h=self.h)
3.5 Vision Transformer
# Save size to recoven it later q = wwwf. Ca(EcY) s.s
34

---

Published in Transactions on Machine Learning Research (12/2023)
By appropriately nanaging the axes, nininal
changes to our previous code
necessary.
unbind(q)
unbind(self (k((x)))
unbind self. (el.f.v(xx)))
self walttseeooddrrdddcttttttnnnnnii k,
Rebind to feed to the transpased convolution
layer.
Sanops.rearrangeeC 'N (H KQ k. (k
W',
b-self.h, Mel))
self.Co(o)
Single hatch elenent, = [1]
Y, [16, 16], [16, 16], [33], the additional configurations, heads, kemel, stride 4, 3,
Internal Data,
BYc rand(b c + v)
Data,
canchhrand(b
We can now run the algorithn, visualAttention Visualattention(((())) k, heads, kernel, stride)
Interestingly, the height/width reduces by for stride values above 1. Otherwise, It stays the sane. visualAttention. (EV..,, xxc)_size()
torch. 33, 151)
Appendix
A container the nodules, Replace module definition
self.Cq co.vodee heads, kernel, stride) with;
self.cq Tracker(nn, Convzd(e, heads, kernel, stride), "Query convolution")
And the Input / output sizes (to check diagrans) will be printed.
class Tracker (nn.Mooue)))))))
def (self, nodule: nn. Module, nane str "*):
super() super() init_()
self.nadule. nodule if name: self.nare nane
self.nare self.nodule._gee_aaree def forward(self,, x):
x_size size tece_t__strinneee sttttnngxx.siz(()
y_size size stiznng(x.sizz())
(f \t (x_size) >> (y_size)*)
return
35


Reviews:
[{"review_id": "xBmGBDlikH", "reviewer": "Reviewer_wTkA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Starting with strengths; I really appreciate the work that was put into this paper. It's no easy task to think of a (visual!) language that properly defines an extremely wide set of algorithms. I think this paper has value on that sole basis. It's also an interesting (I'm not sure what to call it?) socio-technical contribution, rather than the usual empirical paper.\n\nThat being said, I'm honestly not sure what to make of what the authors propose. This is perhaps on a more personal & opinion-level, but in some sense this is where the paper is trying to appeal to the reader. The technical contribution of the paper is obvious, but the paper is claiming to try to do more than that; it's suggesting ML as a field changes its paper-writing culture. I'm having difficulty with the latter.\n\nI have to say when I first scrolled through the paper and happened upon Figures 15 & 16 I was a bit skeptical; too high information density can be just as bad as too low density. \n\nAnother question I ask myself, of course inspired by [XKCD 927](https://xkcd.com/927/), is whether the proposed approach is constraining or liberating. Does it enhance our ability to think of new architectures and communicate them? Or does it narrow our thinking [1] to only algorithms that can be expressed (beautifully [2]?) within the language?\n\n[1] The Hardware Lottery, Sarah Hooker, 2020  \n[2] Lost in Math: How Beauty Leads Physics Astray, Sabine Hossenfelder, 2018\n\nLet me be more concrete.\n\nThe basis of what is meant to be widely adopted is already non-trivial. Students come out of a typical CS undergrad knowing linear algebra; as far as I know, category theory, not so much. I certainly only have a superficial understanding of it after 15+ years of ML. Perhaps I'm not the right reviewer for this, but in some other sense, am I not the perfect audience? I suspect the authors want to target _me_.\n\nAfter reading the paper, I must admit I'm still fairly confused and I don't know how to parse the proposed diagrams. I think part of it may be my lack of familiarity with the basics, but I think the presentation could be significantly different. To be honest, if this work is meant to appeal to a wide audience, there should be way more hand holding. Two suggestions:\n- associate each diagram with its equivalent in numpy/torch/some current standard\n- associate each diagram with the corresponding linear-algebraic operation (broadcasts and shapes can typically be expressed with proper indexing; this feels feasible)\n\nAt least for the first few diagrams, this would be helpful for the more code/linear algebra minded to understand this work. I'm also still unsure what dashed lines represent. Dotted boxes similarly seem to have an important semantic meaning, from what I've gathered elementwise application of a function, but that doesn't seem to work everywhere.\n\nThere is one other major downside to the proposed approach: software and science, especially an empirical science like ML, I think is aptly described as a stack of abstractions. The neat thing about that is that if I know what a self-attention block is, and I suspect the reader knows what it is (with an appropriate reference lack-thereof) then I can simply make a figure using self-attention as a building block. We take for granted now that a \"neural net layer\" is well represented by a matrix multiply + a non-linearity, so much so we don't even need to mention it anymore in papers; this wasn't the case in the 80s. This is great, and leaves more space for more interesting stuff. If we had to redefine everything from scratch in every paper, we'd have enormous papers. I worry that as architectures grow more complex or differ from the usual thinking, neural circuit diagrams would grow to unwieldy sizes. Already, Figure 19 is a whole page; some papers use transformer encoder-decoders as just a block in their overall architecture.\n\n\nSome other comments:\n\n> The reader must understand a specific programming framework\n\nPresumably here the reader must understand a specific diagram language.\n\n>  Improved communication of architectures, therefore, will offer clear scientific value.  \n\nI don't really disagree of course, but I feel like this paragraph presents a false dichotomy. Deep learning's impact on the world was catalyzed by the existence of well written open source software such as theano, pytorch, gym, and so on, as well as well written open source implementations of various architectures that could easily reused. At the end of the day, it's unlikely there will be such a thing as a pure description of an ML algorithm; should we also list the random number generator algorithms used to initialize the model as part of an architecture? A static artifact that we can safely rely on is the code that produced a paper and the corresponding package versions used (or better, some container like Docker cementing the environment in which some code ran); these are becoming more popular and expected from papers.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel way of presenting tensor transformations as commonly used in deep learning architectures. The so called neural circuit diagrams are effectively a well-defined graphical language that can exactly describe deep networks. This is presented as an alternative to ad-hoc figures+equations contained in most DL papers.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "My main requested change is the hand-holding I suggested above. I know I come off as mostly negative in my review but I believe in the work. I do think my inability to understand most of the paper says something about it; either this paper is meant to appeal to mathematicians analyzing deep nets (in which case I am the wrong reviewer & audience), or this paper is overestimating the reader's familiarity with some concepts and thus giving a presentation that is confusing to most.\n\nI'm of course assuming here that I'm a representative ML researcher.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "KADqlZlI3p", "reviewer": "Reviewer_XvGP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nI strongly agree with the premise of this work, which is that deep learning papers and their diagrams omit too many details, which makes it difficult to reproduce them. This limits the overall pace at which advances can be incorporated. \n\nThe proposed neural circuit diagram look great from an artistic standpoint (as far as I am able to judge this). The bright colors are particularly appealing, and I can imagine how (once a certain degree of familiarity has been obtained) these diagrams could be useful when no code is provided. \n\n**Weaknesses**\n\nGenerally, this paper is a bit unusual in the kind of contribution it makes, and I find it difficult to review. A framework is proposed for diagramming deep learning architectures, yet only few guidelines and standard notations are presented for creating these. For example, after reading the paper I am not a lot wiser about drawing these diagrams for other architectures like UNets, LSTMs, or deep CNNs like Inception Networks. Because of the lack of standardization, it feels that different people creating such diagrams would end up with different drawings, which perhaps defeats the purpose a little.\n\nTo highlight some broader issues in particular:\n\n* Before reading the paper in detail, I was not able to understand the diagrams very well. I initially took some time to study Figure 12, and I honestly found it quite difficult to make sense of all the details. Operations like tensor products, softmax axis normalization, and other symbols for drop-out are not intuitive. Similarly, the distinction between lines for coupling and collection is difficult to keep track off. Perhaps these diagrams are not meant to be intuitive at a glance, and they behave more like code or pseudo-code whose language needs to be understood first. However, in that case, having annotated code seems more useful. I note that this could also be a lack of familiarity with these kind of visualizations that can improve over time.\n* Few guidelines are provided for creating these kind of diagrams, and the level of detail used for various components of a models seems up to the user. For example, in the Transformer figure (Fig 12), the masking / dropout and normalization operations are presented in relatively little detail compared to the matrix multiplications. This is partially resolved via the text boxes, which almost act like a legend, though limited guidance is provided for structuring these. Because of this, it feels like there is no standard neural circuit diagram even for the Transformer, and different researchers may end up with different diagrams for the same thing. Indeed, after reading the paper in detail I don’t feel very confident in my own ability to create similar diagrams for other architectures (like models in my own papers). This could well be an issue with my own artistic abilities, yet it makes it difficult to judge the significance of this kind of diagramming. At the moment, I don’t see much benefit in adopting this kind of diagramming compared to sharing code, and I would not expect this approach to be well adopted in the community.\n\nI would ask the authors, who these diagrams are meant for, in what situation they are expected to be used, and who the intended creates of such diagrams are. This would help determine the impact of this contribution and whether it is reasonable to expect adoption.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes Neural Circuit Diagrams as a way of visualizing deep learning architectures in detail. It revolves around shaped data whose axes are combined via coupling or collecting. Morphisms, having fixed input and output shapes act as operations on the data.\n\nThe paper proceeds with several examples of fundamental data operations (Figures 1–6). Afterward, a Neural Circuit Diagram of an MLP is shown (Figure 8), followed by diagrams of scaled-dot-product attention, multi-head attention, and eventually a transformer (Figures 10–12). Along the way, a few standard morphisms are introduced, such as for transposing, binding, and copying (Figure 7).\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would like to see a lot more standardization and guidelines included for creating diagrams for many of the standard architectures. Examples of vision architectures like U-Nets or Inception nets, and recurrent nets like LSTMs would be particularly informative. Other examples focusing on handling multiple modalities (eg vision and text), or cases where processing proceeds differently along various input dimensions (eg. factorizing attention spatially and across time, as is commonly done in video generative models).\n\nAs I mentioned before, to be able to judge the potential impact of this work, it would be good to be more specific about the intended target audience and designer of such diagrams. It would also be good to clarify how it can be ensured that multiple diagrams of the same model end up being the same, since this would overall benefit adoption. Finally, commenting on how these diagrams complement annotated (pseudo-)code would also be useful, especially since readers may generally already be more familiar with code.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "munTkXKNva", "reviewer": "Reviewer_82Vd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I appreciate the need for improving diagrams, and even standardizing to a greater degree.  The notation and techniques suggested here offer some nice ways looking at and communicating operations, particularly the analogy between line connections and Einstein summation notation, and improved clarity of which operations apply to which dimensions.  Certainly, a reader can get some good ideas for drawing effective specifications.\n\nHowever, I think this currently lacks the empirical rigor and/or interesting enough technical contribution to meet TMLR criteria.  The current paper presents a system of diagramming, but no evidence to support its improvement over a combination of ad-hoc diagrams, equations and descriptions (user studies or questionnaires could supply this, for example).  Similarly, a script or tool that can automatically generate these figures (or a first draft that can be edited) would go a long way in encouraging adoption (or adoption of pieces), as researchers could view more different models to get a feel for how the diagrams work (and without this, I wouldn't expect authors to be able to get all the details right).\n\n\n---\n\nPutting forward these diagrams makes it all too tempting to provide nitpicks on the diagrams and styles themselves.  I don't think it's appropriate to base this review on my own response to these particular diagrams, and tried to give more substantive feedback first, both above and in the \"requested changes\" section below.  Still, I've listed many of my observations, which I hope can also be helpful:\n\n* the colored boxes are a little too busy, and group together data that is separate; dotted lines separate tensors, but the pink boxes are more apparent and look like the larger grouping.  is it possible to write these in black-and-white or grayscale, with color optional addition?\n\n* representing dimensions with different lines makes a lot of sense for some operations, particularly multi-dim transposes and multiplications, but gets in the way for others, e.g. copy operation (what if only the x dim is copied in fig 11 but not m, what does that even mean?)\n\n* These diagrams are very exacting, with small details having large differences (the dashed lines indicating separate tensors is one of the largest here).  As such, I don't expect many researchers or engineers would include all details based purely on spec, leading again to ad-hoc/freeform diagrams, though with perhaps a few new useful notations.  That strikes me as similar to UML, which is scarcely followed to the letter.\n\n* the scale by 1/sqrt(k) in fig 10 in particular isn't very intuitive to me.  The arced connection for matrix mul makes some sense after getting used to it (it combines the two dims, similar to einstein summation notation, as pointed out), but the arrow doesn't really go anywhere to connect to anything.  Scaling could just as easily be applied over any of these dimensions (y, k, x).\n\n* reshaping is a fundamental operation I didn't see covered.  It would be convenient to refer to passing multi-dim tensors with a single line, and split into multiple dims only when needed.\n\n* taking a page from circuit diagrams, what about using cross-hatch with the dimension to denote shapes?\n\n* softmax arrow notation is nice, though I think I've seen right-pointing arrow used for this in the past more often, suggesting comparison into (soft) 1-hot outputs\n\n\n* binding diagram is almost the same as copying (no dotted line) but with much different effect.  also keeping in mind addition is the grad of copy, it would make sense for the copy diagram flipped horizontally to perform a sum, but the \"bind\" operation doesn't do that (getting diagonal values of a matrix also not very common)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a diagramming method for communicating deep learning models, focusing on transformers.  Notations are described for tensor representation, combination, and function application.  After examples describing MLP and multi-head attention, the paper culminates in describing a transformer model.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Improvements in at least one of (1) and/or (2) below are critical, in my view (and preferably both).  Additionally, (3) would further strengthen the work.\n\n1. Research into the practical usefulness and benefits of these diagrams would help bring a level of rigor, for example user studies or questionnaires.  Maybe a study where a user needs to code up an implementation based on a description with vs without the new diagram, for example?  Or a questionnaire asking readers to answer questions about specifics with vs without this diagram?  Diagram vs equations?\n\n2. Likewise, pairing this with an automated tool that can generate these diagrams, perhaps based on a pytorch definition, would be useful in encouraging usage and potential adoption.\n\n3. The diagrams are applied exclusively to sequence transformers (and a small MLP), as case studies.  However, other models and domains can have different details that are important to get across.  Object detectors, for example, are principally concerned with object scale and alignment, some of which can be shown with shape dimensions (scale), but others (alignments, anchors) are not as clear.\n\n", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper itself is unconventional, but I think that is precisely the sort of thing TMLR is designed for. It underwent multiple major revisions through the review process, first adding more technical depth and then dialing back and making it more accessible. The current version strikes a reasonable balance between the two. The addition of code examples and extension to many contemporary architectures greatly clarifies things and makes it more compelling for readers in the field.

The paper has already gone through so much editing that I hesitate to condition acceptance on further changes, since I don't think it will dramatically affect the paper. However, I would encourage the authors to consider Reviewer 82Vd's advice about adding further explanation for transpose and convolution operations.
Decision: Accept (poster)