SIP : Injecting a Structural Inductive Bias into a S eq2 S eq Model by Simulation (2024)

Abstract: Strong inductive biases enable learning from little data and help generalization outside the training distribution. Popular neural architectures such as Transformers lack strong structural inductive biases for seq2seq NLP tasks on their own. Consequently, they struggle with systematic generalization beyond the training distribution, e.g. with extrapolating to longer inputs, even when pre-trained on large amounts of text.We show how a structural inductive bias can be efficiently injected into a seq2seq model by pre-training it to simulate structural transformations on synthetic data. Specifically, we inject an inductive bias towards Finite State Transducers (FSTs) into a Transformer by pre-training it to simulate FSTs given their descriptions. Our experiments show that our method imparts the desired inductive bias, resulting in improved systematic generalization and better few-shot learning for FST-like tasks. Our analysis shows that fine-tuned models accurately capture the state dynamics of the unseen underlying FSTs, suggesting that the simulation process is internalized by the fine-tuned model.

## introduction inductive biases, i.e. the preferences and the abstract knowledge a model brings to the task, enable a model to learn from small amounts of data and generalize systematically outside of the training distribution. while seq2seq models perform very well on in-distribution data, they usually lack structural inductive biases and consequently struggle with systematic generalization. previous work has shown that this includes generalization to unseen combinations of known sub-strings @xcite , extrapolation to longer inputs @xcite and deeper recursion @xcite . integrating structural inductive biases into seq2seq models is challenging. one popular ap1 we release our code at https://github.com/namednil/sip proach is to develop specialized architectures @xcite @xcite , which makes it difficult to precisely control and adjust the nature of the inductive bias to changing demands as the architecture would need to be modified and models retrained. recently, some works instead have tried to inject inductive biases into seq2seq models by pretraining on a well-chosen synthetic task @xcite @xcite or meta-learning on a distribution of synthetic tasks @xcite using maml @xcite . here, the inductive bias can be controlled by the choice of the synthetic task. however, meta-learning with maml scales poorly because it requires expensive second-order derivatives and standard pre-training can be less effective @xcite . in this work, we present a computationally inexpensive way of injecting a structural inductive bias into a transformer. we focus specifically on introducing an inductive bias that is helpful for tasks that traditionally have been approached with finite state transducers (fsts). we choose fsts because they are formally well understood, are easy to generate automatically, and are one of the simplest computational devices that are useful in nlp applications. while we focus on fsts, the methodology is fairly general and our approach also provides a starting point for incorporating more general structural biases, provided by more expressive formalisms such as pushdown transducers. our approach (sip, for simulation-induced prior) is simple (see fig. 1 ): given a representation of an fst and an input string, a transformer is pre-trained to predict what the output of the fst is on the given input. we assume that fsts are not specified for fine-tuning on downstream tasks, so we replace the fst with tunable embeddings and fine-tune the model solely on input/output pairs. since we fine-tune all parameters, the model can deviate from fst-like behavior if needed. contributions. we show that a model pretrained with sip has an inductive bias that improves systematic generalization and few-shot learning for 'fst-like' downstream tasks. sip not only improves systematic generalization on fst tasks similar to those seen during pre-training but also on ones that are structurally more complex. the same pre-trained model also transfers well to natural data and achieves strong results on few-shot learning of text editing (e.g. jane @xmath0. doe) and grapheme-to-phoneme conversion. our probing experiments give insights into how the inductive bias is injected: sip not only leads to the imitation of the input/output behaviour of fsts, but encourages dynamics to emerge that simulate crucial aspects of fsts in the hidden representations. fine-tuning can leverage these dynamics, providing the inductive bias, and learn representations that resemble those of ground truth fsts. a range of approaches have been developed to tackle this, with many works focusing on special-ized architectures @xcite @xcite . @xcite find that the specialized architectures they consider do not transfer well to tasks beyond the context in which they were designed. this highlights the importance of being able to adjust inductive biases more easily than re-designing the architecture of a model. large-scale pre-training has also been shown to help with systematic generalization @xcite . however, challenges remain even for llms such as gpt-3 and palm @xcite . the methodology we present in this work can be used to create additional material for llm pre-training. here we focus on smaller models and leave this to future work. pre-training with synthetic tasks. pre-training a model on a synthetic task to introduce specific inductive biases has been explored by several recent works. @xcite identify useful 'skills' for news summarization and develop a pretraining task accordingly. lime @xcite targets mathematical reasoning and is pre-trained on string manipulation that resembles formal reasoning. papadimitriou and jurafsky (2023) consider pre-training with several synthetic languages to investigate which helps most for language modelling. in contrast to these works, our approach targets simulating a computational device and maintains a closer relation to the pre-training setting because of the tunable prefix. a challenge for using individually hand-crafted tasks is to cover a sufficient space of phenomena that are relevant to downstream tasks. instead of training on a single task only, @xcite ; mccoy and griffiths (2023) meta-learn on a distribution of tasks using maml @xcite . our approach also uses a distribution of tasks but it scales better than maml-based methods because maml requires computing and storing secondorder derivatives. for example, the transformer we train has a magnitude more parameters than the lstm of @xcite and is pretrained on a smaller gpu (a100 vs rtx 2080 ti). in addition, as the complexity of each individual task grows, maml requires more examples per task. we circumvent this by using a compact and unambiguous description of each task instead. ## finite state transducers we briefly review finite state transducers (fsts) which we use in our experiments. fsts are closely related to finite state automata (fsas). while an fsa describes a set of strings, an fst describes a relation between strings, i.e. a set of pairs (x, y), where x is an input y is an output. fsts can be visualized as labelled directed graphs (see fig. 2 ), where the nodes are called states and the edges are called transitions. consider the path q0 2b . this path is called an accepting path because it starts in an initial state (indicated by an arrow 'from nowhere' pointing to the state), and it ends in a final state (indicated by double circles). an accepting path shows what an input can be mapped to. in this case, the path shows that the fst transduces the input 001 into the output 111. we can read off which input an accepting path associates an output to by concatenating all the strings along the path occurring before ':'. the output can be determined by concatenating the strings after ':'. hence, each transition σ : ρ --→ can be thought of as figure 2 : examples of functional fsts. the fst in (a) deletes leading zeros. the fst in (b) replaces any 0 by a 1 if the last input symbol is a 1. conversely, if the last symbol is a 2, any 0 is replaced by a 2. the output can only be determined after the last input symbol. in general, an input can be paired with arbitrarily many different outputs. we call an fst f functional if every input x is paired with at most one output y, and use the notation f (x) to refer to y. all fsts we consider here are functional. in this work, we investigate generalization across different sub-classes of fsts, namely from the less expressive deterministic fsts to non-deterministic fsts. an fst is called deterministic if (i) it has a unique initial state, (ii) for all states q and input symbols σ there is at most one transition q σ : ρ --→ q ′ and (iii) σ ̸ = ϵ. intuitively, this means that in any state, for an input symbol σ there is at most one possible next state and one possible output, and hence for any input string there is at most one path that is compatible with it. because of this, we can always infer a prefix of the output by looking only at a prefix of the input string and ignoring the rest. for example, consider the input prefix 001. in the deterministic fst in fig. 2a , we know that the output has to start with 1 because there is only one path that is compatible with 001. in contrast, in the non-deterministic fst in fig. 2b , three paths are compatible with 001 that have different outputs. in that case, we can only determine the output once we look at the last symbol of the input. in short, non-deterministic fsts can take context to the right into account but deterministic fsts cannot. ## simulation-induced prior our approach largely follows the pre-training and fine-tuning paradigm. we first pre-train on synthetic fst tasks by giving the model a representation of an fst as a prefix and an input string (see fig. 1 ). the training objective is to predict the output of the fst on the input string. our research hypothesis is that training a model to predict the behaviour of an fst incentivizes the model to acquire reusable dynamics that internally simulate fsts. when fine-tuning the model using a tunable prefix instead of an encoding of an fst, these dynamics should be easy to leverage and provide a structural inductive bias for fst-like tasks. ## evaluating sip's inductive bias to understand the effects of our pre-training procedure on the inductive bias of the model and on the downstream performance, we first explore systematic generalization on synthetic fst tasks. this allows us to precisely control the similarity between the pre-training and the downstream task. ## analysis: sip leads to fst simulation we motivated our approach by the hypothesis that sip's pre-training encourages the model to simulate fsts internally, and that this provides the structural inductive bias. in this section, we present evidence that (i) sip models indeed approximately simulate fsts in the hidden states, and (ii) that the dynamics responsible for simulation are re-used after finedtuning all parameters on input/output pairs only. for a model to simulate fsts in its hidden representations, it must be able to track the fst state when processing a string, and it should be possible to extract the fst state with a probe. to test this, we mirror the pre-training setup and provide sip-d4 with an fst and an input string (fig. 4 , left). for each token, we extract the top-layer activations of the encoder, and learn a linear probe with a softmax layer to predict the id of the state that the given fst is in before processing that token. since state ids are largely arbitrary, the probe has to learn to relate the hidden representations to the fst presented in the input. the probe achieves 99.3% token-level accuracy on a test set with unseen fsts, and a wholesequence accuracy of 93.9%. we also evaluate a trivial heuristic that returns a random state that has an appropriate outgoing transition for each token in the input. this heuristic achieves a token-level accuracy of 68.9%, and a whole-sequence accuracy of only 17.8%. a probe trained on byt5 representations, i.e. before sip pre-training, performs even worse at 42.9% token-level accuracy and whole-sequence accuracy of only 7.1% (see appendix d.2). hence, the model has learned a non-trivial way to simulate transitions between states of the fst encoded in the prefix. this is remarkable because the pre-training procedure for sip-d4 does not provide supervision for how to process strings. while this shows that sip leads to the simulation of state transitions after pre-training, does the model leverage the simulation ability on downstream tasks? recall that we fine-tune all parameters of the model (section 4.2), so the model could figure 5 : row-normalized confusion matrices on the training and test data between ground truth and the state predicted by the frozen probe applied to fine-tuned models. we average across the 5 iteration generalization tasks (section 5.3). ## conclusion we present sip, a simple, efficient and adjustable method for introducing a structural inductive bias into a seq2seq model. we focus on an inductive bias towards fsts, one of the simplest computational devices that is useful for nlp applications. we achieve this by pre-training a transformer to simulate automatically generated fsts, i.e. to predict the output of an fst given an input string and a description of the fst. our experiments show that our method imparts the desired inductive bias, resulting in improved systematic generalization and better few-shot learning for fst-like tasks. in addition, we show with probing experiments that a model trained with sip simulates transitions between fst states in its hidden representations, and that the dynamics behind this are leveraged during fine-tuning. in future work, we plan to extend this methodology to more expressive formalisms such as pushdown transducers which can be used for a wider range of downstream nlp tasks.