Injecting a Structural Inductive Bias into a Seq2Seq Model by Simulation

Thank you for your review!

If By-T5 is already pre-trained in natural data before the synthetic pre-training, I wonder how much of an influence there is from the "pre-pre-training" in enabling OOD generalization and such.

We tried to train a ByT5-style model from scratch. However, none of the hyperparameter configurations we tried resulted in a model that would converge well. We hypothesize that a model has to be in a reasonable space to begin with to make learning easier.

While we cannot exactly attribute what part of the generalization ability comes from the combination of the original ByT5 weights and our pre-training, ByT5 on its own does not have a strong inductive bias for FST-like tasks and performs worse in our experiments.

We already have prior works showing the viability of synthetic pre-training and knowledge transfer from natural language tasks.

We briefly review some methods for pre-training with synthetic data in our Related Work section. In contrast to these, our methodology relies on simulating a computational device (FSTs, specifically).

We view our contribution not only as a pretraining procedure but it also includes a specific fine-tuning method (prefix of tunable embeddings) mirroring the pre-training. We have now included an ablation of the tunable prefix during fine-tuning (see Tables 1,3,4). The results show that using a tunable prefix is important for achieving the best performance.

The scope feels limited [...] It is not as clear what the motivation for exploring particularly FST-related tasks is.

FSTs are an important computational device for NLP on the level of phonology and morphology. Our methodology of simulating a computational device is also quite general, and we view this as a stepping stone to more expressive computational devices with broader applications, such as Pushdown Transducers which could help with processing hierarchical structure. We hope that the results of this work can inform that exploration.

What are the average and maximum sequence lengths in the pre-training data, training data, and iteration generalization data?

In the pretraining data: the average input sequence length is 15.57 and the maximum is 35. The distribution of the training/test length for iteration generalization depends on the number of states in the FST:

Would it be possible to explore generalizations to higher lengths e.g. 100 or more?

We ran an experiment on FSTs with 4 states (i.e. from the pre-training distribution), where models are trained on inputs of length up to 15, and tested on inputs of length 40 - 70, and on lengths 90 to 110. Note that this is beyond the lengths seen during pre-training. The results are below (averaged over 5 tasks), and show that our pretraining still has a strong positive impact:

We hypothesize that the accuracy of SIP-d4 drops on lengths 90 to 110 because the relevant positional embeddings were not pre-trained. SIP-d4-long is a version of SIP-d4 that has been further pre-trained on strings of length up to 110. We refer to the new Appendix I for more details.

Thank you for your review!

[...] more ablation experiments and analysis to understand the changes brought about by the FST on the results and why it can bring consistent improvements.

We have included an analysis in Appendix G that gives some insights into what the learned prefix can represent. In a nutshell, when finetuning on tasks from the pre-training distribution and adapting only the prefix (i.e. keeping the rest of the model frozen), the learned prefix resembles an encoding of the ground truth FST.

We have also included an ablation where we don’t use a prefix of tunable embeddings during fine-tuning (see Tables 1,3,4). The results show that a tuneable prefix is important for achieving the best performance.

Despite the many experiments conducted in this paper, I still hope that the authors can apply the method to large language models. It is very important to determine whether the proposed method is still effective in LLMs and whether it can solve some problems existing in larger models, such as hallucinations.

In future work, we are planning to integrate pre-training to simulate computational devices (e.g. FST or grammars) into LLMs. This will likely require training on a mixture of simulated data and other pre-training data, so will be computationally expensive.

We think that simulating computational devices could potentially help with some problems of LLMs, in particular systematic generalization, but also specific forms of hallucinations, where the model ignores a rule and resorts to default behaviour (see also the MemoTrap dataset). For example, consider the following simple text editing problem (similar to Section 7.2)

The current version of ChatGPT outputs “J.K. Rowling” for the name “John Edward Rowling”, hallucinating the K.

Please feel free to reach out if you have any additional questions or would like any clarification. We'll do our best to respond in the remaining 40 hours before the discussion period ends.

Thank you for your review!

1. Limitation of FST which has to be deterministic. / The need to encode FST as a sequence of prefix encoding also makes the design of FST topology limited

Any FST can be encoded with our method (as you mention though, the encoding is not unique, e.g. because of different orderings of transitions, see also our reply to your third point). In addition to our experiments with deterministic FSTs, we pretrain on non-deterministic FSTs in section 6.4, which improves OOD accuracy on tasks with underlying non-deterministic FSTs in comparison to pretraining on deterministic FSTs.

Our framework could also easily be extended to stochastic/weighted FSTs by including the weights in the prefix as well.

2.The setup looks up state and transition from the embedding table which is not scalable

the current proposed framework would be limited when #states/transition explodes

While encoding very large FSTs into a prefix can make pre-training expensive (due to the quadratic runtime complexity of the transformer), our results in section 6.3 show that improvements can already be achieved by pre-training on smaller FSTs, which mitigates this issue to a certain extent.

We would also like to mention that other works in this area use black box meta-learning, where the prefix contains a small synthetic dataset of input/output pairs [1, 2]. Instead of encoding a synthetic dataset, we encode the underlying rules of the task, which is much more scalable.

3.Positional information of transformer is used as usual, which means two identical FST with different state ordering have different representation in the transformer. This could result in unwanted behavior.

We ran an experiment to test how sensitive our pre-trained model is to different orderings of the transitions in the prefix: when given an encoding of an FST (from the same distribution as the training data), and a version of the encoding where the order of transitions is randomly permuted, the model predicts the same output in 98.3% of the cases. In the cases where the predictions are not exactly the same, they differ by 1.65 characters on average. We hypothesize that the transformer learns to encode them in a relatively position-agnostic way (e.g., by suppressing positional encodings); however, we did not perform a 'mechanistic' analysis of how it has been achieved.

To a lesser extent, this robustness towards permuting the prefix holds also in the fine-tuned case, when only the prefix is tuned and the other model parameters are frozen during finetuning: randomly permuting the order of the vectors in the tuned prefix reduces accuracy by 7.0 percentage points on average, with a median reduction of only 1.3 points (see the new Appendix G for details).

1.What exactly is learned from the tunable embedding? [...] such embedding should be tuned toward certain encoding from pre-training stage right?

We performed an analysis that computes the cosine similarity between the tuned prefix of a model (when only the prefix is tuned and the rest of the model is frozen) and an encoding of the “ground truth” FST. This analysis revealed that the learned prefix is reliably more similar to the ground truth FST than to 10,000 different FSTs with similar properties (see the new Appendix G for details). This shows that information is represented in the prefix in a somewhat similar way to encodings of symbolic FSTs.

2. The nature of FST accepts compositional design (that is, addition or composition of different FSTs can be easily combined). I wonder if the proposed approach, when trained on different FSTs encoding, would generalize well to the task that is solvable by their composition?

Thank you for suggesting this interesting research direction! It is not clear what the composition function could be for the encoding learned through fine-tuning. However, we think it might be possible to adapt the model to the composition of two FSTs when their encodings are available using a small amount of training data from the composition of the FSTs. There may also be a way to meta-learn the model to accept a certain composition function (e.g., involving the concatenation of prefixes).

[1] Compositional generalization through meta sequence-to-sequence learning - Lake. NeurIPS 2019.

[2] Human-like systematic generalization through a meta-learning neural network - Lake and Baroni, Nature 2023.

Please feel free to reach out if you have any additional questions or would like any clarification. We'll do our best to respond in the remaining 40 hours before the discussion period ends.