Functorial Language Models

We introduce functorial language models: a principled way to compute probability distributions over word sequences given a monoidal functor from grammar to meaning. This yields a method for training categorical compositional distributional (DisCoCat) models on raw text data. We provide a proof-of-concept implementation in DisCoPy, the Python toolbox for monoidal categories.


Introduction
We introduce functorial language models: a principled way to compute probability distributions over word sequences given a monoidal functor from grammar to meaning. This yields a method for training categorical compositional distributional (DisCoCat) models on raw text data. We provide a proof-of-concept implementation in Dis-CoPy , see appendix A.
Language models, i.e. probability distributions over word sequences, are a cornerstone of natural language processing and information retrieval (Ponte and Croft, 1998). Neural models (Bengio et al., 2003), where these distributions are learnt by a neural network, are now built in everyday tools such as virtual assistants, automatic translation, etc. However effective they may be, neural models are not compositional in the sense that they do not take the grammatical structure of sentences into account, at least not explicitly. This severely limits our ability to interpret them: if we open the black box, we only see matrices of weights.
On the other hand, categorical compositional distributional (DisCoCat) models (Clark et al., 2008(Clark et al., , 2010 use grammatical structure to compose the distributional meaning of words together into a meaning for the sentence. Grammar is explicitly represented as string diagrams, which allow formal reasoning about natural language semantics, for example analysing ambiguity (Kartsaklis et al., 2013(Kartsaklis et al., , 2014Piedeleu et al., 2015) and entailment (Sadrzadeh et al., 2018a;de Felice et al., 2019).
However, despite emprical validation on small-scale examples (Grefenstette and Sadrzadeh, 2011), applying DisCoCat models at a large scale to real-world text data appears like a far-away goal. The challenge is at least two-fold: 1) we need a robust way to predict the grammatical structure of a given word sequence, 2) we need an efficient way to learn the meaning of each word, given the grammatical structures in which it occurs. The first point requires a probabilistic grammar: it is not enough to know whether a sequence is grammatical, we need a probability distribution over its possible parsings.
It has also been argued that a robust model ought to be incremental (Sadrzadeh et al., 2018b): scanning through the given word sequence left-to-right and updating its prediction for what comes next. This plays a crucial role in the context of informal language and dialogues (Purver et al., 2006). Previous work by the first author and collaborators (Shiebler et al., 2020) characterise these incremental, probabilistic parsers in terms of a functor from formal grammars to automata.
The second challenge, i.e. how to use grammar for learning the meaning of words, is essential if we want DisCoCat to work as a standalone model for natural language. Indeed, the first experiments (Grefenstette and Sadrzadeh, 2011) assumed the vector embedding for nouns to be given by some other means, such as co-occurence counts or neural language models. On the other hand, in more recent experiments on quantum hardware (Meichanetzidis et al., 2020a,b;Coecke et al., 2020) the meaning of words is learnt directly from the training data of a supervised questionanswering task. We coin the term functorial learning for this approach, as it can be understood as learning a functor from data.
We give a formal definition of this functorial approach and define functorial language models by composition with a probabilistic grammar. Concretely, we train a DisCoCat model to predict which word is missing in a sentence with a hole. We argue that this captures precisely the idea of distributional compositionality: extending Firth's principle, we shall know a word by the company it keeps and by the grammatical structure in which it occurs.
DisCoCat models have a one-sentence definition: they are rigid monoidal functors from the category generated by a pregroup grammar to vector spaces and linear maps. Let's unpack this definition and repack it in a format suitable for computation.
Fix a set of words , called the vocabulary, and a set of basic types . A rule is a pair of sequences of words and types ∈ ( + ) ★ × ( + ) ★ , where ×, + and ★ denote respectively the Cartesian product, the disjoint union and the free monoid ★ = ∈N with unit 1 the empty string. A grammar = ( , , , ) is given by a set of such rules together with a distinguished ∈ called the sentence type. This defines a signature in the sense of (Selinger, 2010) which generates a free monoidal category G. The objects are sequences of words and types, the arrows are progressive planar string diagrams with rules as boxes. The language of is defined as the sequences ∈ ★ such that there is a diagram : → in G: its grammatical structure.
This definition merely reformulates the notion of semi-Thue system (Thue, 1914) in the language of monoidal category theory. (Post, 1947) and (Markov, 1947) independently proved the undecidability of the parsing problem, i.e. given and ∈ ★ decide whether there exists : → in G. (Chomsky, 1957) then put the equivalent notion of unrestricted grammar at the bottom of his well-known hierarchy. In parallel, (Lambek, 1958) defined his calculus in terms of closed monoidal categories, taking inspiration from the categorial grammars of (Adjukiewicz, 1935) and (Bar-Hillel, 1953). Half a century later, pregroup grammars (Lambek, 1999(Lambek, , 2001(Lambek, , 2008 simplified this calculus by going from closed to rigid monoidal categories. Pregroup grammars are at the basis of the original DisCoCat model of (Clark et al., 2008). Definition 1.1. A pregroup grammar has types × Z where for ∈ N the types ( , − ) and ( , + ) are written ... and ... respectively and are called iterated left and right adjoints. The rules are of two kinds: cups and triangles. Cups have the shape ( , ) ( , + 1) → 1 for ∈ Z, i.e. they cancel a basic type with its right adjoint. Triangles have the shape → for a word ∈ and a type We only consider finite dimensional vector spaces. A diagram is progressive when its wires go monotonously top to bottom, i.e. they do not bend up or down.
A diagram is planar when its wires do not cross.
∈ ( × Z) ★ . Note that we draw words as the labels of triangle boxes rather than as input wires.
A pregroup grammar = ( , , , ) defines a rigid signature: the generating objects are given by words and types + , the generating arrows are given by dictionary entries → ∈ . Thus, it generates a free rigid monoidal category G where the object are sequences of words and types with iterated adjoints, the arrows are planar string diagrams with triangles as boxes. Cups are given by the rigid structure of G, they are drawn as bent wires. Again, the language of is defined as The main distinction between Chomsky's phrase structure grammars and the categorial grammar tradition of Adjukiewicz, Bar-Hillel and Lambek is that in the latter "all the grammar is in the lexicon". The only language-dependent rules are dictionary entries of the form : → for a word ∈ and a type ∈ Ob(G). They are drawn as triangle boxes. All the grammatical rules come from the structure of the category G, e.g. the cups of pregroups from the rigid structure. This has both conceptual and computational advantages. Conceptually, it makes categorial grammars universal: the same rules apply to all languages, only the dictionaries change. Computationally, these universal rules have a canonical semantics, thus it is enough to define the meaning of each dictionary entry.

Definition 1.2. A DisCoCat model is a rigid monoidal functor
: G → Vect from the category generated by a pregroup grammar. On objects is defined by a mapping 0 : → N which sends each basic type to the dimension of a vector space. The image of word types ∈ is the monoidal unit ( ) = 1, the image of complex types = ( 1 , 1 ). . .( , ) is the product of the dimensions of their basic types ( ) = 0 ( 1 ). . . 0 ( ). On arrows is defined by a mapping 1 ∈ ( , ) ∈ R ( ) that sends each dictionary entry ( , ) ∈ to a vector ( , ) ∈ R for = ( ) the dimension of its type. Cups are sent to the rigid monoidal structure of Vect. The meaning of a sentence : → is given by the vector ( ) ∈ R for = ( ) the dimension of the sentence type. When ( ) = 1, the meaning of sentences is given by real-valued scalars that can encode either truth value or likelihood.

Remark 1.3. DisCoCat models have been generalised both in their domain and codomain:
one can vary the grammar category as in (Coecke et al., 2013) or the semantic category as in (Piedeleu et al., 2015;Coecke et al., 2018b;Delpeuch, 2019). While we stick to the original definition, our proposal can apply to any model, so long as the meaning of words lives in some vector space, regardless of how these word vectors are composed together.
We fix the map 0 : → N and consider it as the hyper-parameters of the model. All the data defining the DisCoCat model is contained in the finite set of vectors 1 ∈ ( , ) ∈ R ( ) . Thus, we may consider these = ( , ) ∈ ( ) parameters as a machine learning landscape. The idea of functorial learning is to take a training set of pairs ⊆ (G) × (Vect) of diagrams ∈ (G) and vectors ∈ (Vect), and learn a functor such that ( ) = . In practice, this exact functor may not exist thus we fix a loss : ∈N R × R → R + and a regularisation : R → R + then using gradient-based methods we approximate: arg min :G→Vect

Example 1.4. Fix a type
∈ for yes-no questions with ( ) = 1, i.e. the meaning of a question : 1 ... → is a scalar ( ) ∈ R. Take a training set of such questions together with their answer in {0, 1} ⊆ R. Then the functor may be seen as a DisCoCat model for question-answering, see (de Felice et al., 2019). This model has been deployed on quantum hardware, see (Meichanetzidis et al., 2020a,b;Coecke et al., 2020).

Functorial language models
The definition of DisCoCat models can be recast in terms of encoding matrices, first introduced in (Toumi, 2018;Coecke et al., 2018a). The set of meaning vectors 1 ∈ ( , ) ∈ R ( ) can be packed into a set of matrices : | | → 0 ( ) indexed by grammatical types ∈ ( × Z) ★ , where ⊆ is the set of words ∈ with ( , ) ∈ . If we compose the matrix with the one-hot vector for a word : 1 → | |, we get its image under the functor ( , ) = o 9 . In the other direction, if we compose a meaning (co)vector : 0 ( ) → 1 with , we get a (co)vector o 9 : | | → 1 which gives its inner product with the words in . The key insight of our functorial language models is to take softmax( o 9 ) as the probability distribution over words in given a meaning : 0 ( ) → 1.
Remark 2.1. softmax is not a linear map, hence technically we cannot draw it as a box in a Vectvalued diagram. We can however draw it as a bubble, see (Toumi et al., 2021). We can then consider functors which send bubbles to softmax.
Fix a corpus of grammatical sentences ⊆ ∈ ★ G( , ) and assume ( ) = 1, i.e. the meaning of sentences are scalars. We generate a training set from this by taking each sentence : 1 ... → , removing one dictionary entry → by replacing it with the identity → . Thus, we get a diagram : → and take the one-hot vector for the word as target label. Intuitively, given the diagram of a sentence with a hole, we want to predict what word is missing.
This does not yet define a language model in the usual sense, indeed we are assuming that text data comes annotated with grammatical structure. That is, we are computing the conditional distribution ( | ) over words ∈ of type ∈ Ob(G) given the grammatical structure : → in which they occur. In order to get a language model, we need to compose this conditional with a distribution ( | 1 ... ) over diagrams given word sequences: a probabilistic grammar. Learning this distribution from data is called probabilistic grammar induction, which (Shiebler et al., 2020) formulate in terms of monoidal categories. In future work, we plan to train our functorial language model end-to-end on raw text data, i.e. learning both the probabilistic grammar and the DisCoCat model at once.

Implementation
We implemented a proof-of-concept in Dis-CoPy (de Felice et al., 2020) and used jax (Bradbury et al., 2020) for automatic differentiation and just-in-time compilation on GPU. We initialise the encoding matrices at random and use singular value decomposition to obtain word vectors that are close to orthogonal. We used crossentropy loss, a weighted sum of l1 and l2 (1e−1 and 5e−2 resp.) regularisation and Adam optimization (Kingma and Ba, 2017) (learning rate = 5e−2). We set the hyper-parameters to ( ) = 1 and ( ) = 7. After training on a subset of handcrafted data (~100 sentences with three distinct grammatical structures), the model was able to infer the missing word in previously unseen sentences with accuracy > .75, see appendix A. The notebook can be found on DisCoPy's GitHub.

A Experiment
We reproduce our hand-crafted dataset together with some sample predictions from our model. In the following cherrypicked examples, the prediction is correct in the first three cases and incorrect in the last one. This failure can be explained by the fact the word "krill" only appears once in the training set.