Paper: Injecting a Structural Inductive Bias into a Seq2Seq Model by Simulation (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Novelty):
The scope feels limited. We already have prior works showing the viability of synthetic pre-training and knowledge transfer from natural language tasks. It is not as clear what the motivation for exploring particularly FST-related tasks is. Transformers have been shown to underperform in OOD generalization on logical inference [1,2], ListOps [2], Flip-Flop languages [3], parity tasks/sensitive tasks [4], automata tasks [5], and others [6]. It would have been good to contrast the approach with some of such works, reconcile with them, and see if the synthetic pre-training proposed here can be used.

Author response:
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.

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.

Reviewer comment (Experiments):
What are the average and maximum sequence lengths in the pre-training data, training data, and iteration generalization data?

Author response:
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:
|  Num States | Data  | Min | Max | Average |
| ----------- | ----- | --- | --- | ------- |
| 4           | train | 2   | 11  | 4.66    |
| 4           | test  | 4   | 30  | 18.97   |
| 5           | train | 2   | 14  | 5.39    |
| 5           | test  | 4   | 30  | 19.53   |
| 7           | train | 2   | 20  | 6.12    |
| 7           | test  | 4   | 30  | 20.13   |
| 10          | train | 2   | 25  | 7.31    |
| 10          | test  | 4   | 30  | 20.62   |

Reviewer comment (Experiments):
Would it be possible to explore generalizations to higher lengths e.g. 100 or more?

Author response:
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:


|  Model      | Max pre-training length | Test length | Acc  | Edit distance |
| ----------- | ----------------------- | ----------- | ---- | ------------- |
| ByT5        | 1024                    | 40 to 70    | 29.3 | 15.6          |
| SIP-d4      | 35                      | 40 to 70    | **69.4** | **2.61**          |
| ByT5        | 1024                    | 90 to 110   | 1.4  | 55.37         |
| SIP-d4      | 35                      | 90 to 110   | 3.4  | 34.5          |
| SIP-d4-long | 110                     | 90 to 110   | **81.5** | **1.09**          |

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.

Reviewer comment (Experiments):
The author needs to provide more ablation experiments and analysis to understand the changes brought about by the FST on the results and why it can bring consistent improvements.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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](https://github.com/liujch1998/memo-trap)). For example, consider the following simple text editing problem (similar to Section 7.2)
|Input | Output|
| ------ | ------- |
| Howard Phillips Lovecraft | H.P. Lovecraft |
| John Ronald Reuel Tolkien | J.R.R. Tolkien |
| Thomas Stearns Eliot | T.S. Eliot |

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

Reviewer comment (Theory):
The proposed framework would be limited when the number of states or transitions becomes very large.

Author response:
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.

Reviewer comment (Theory):
Positional embedding from the transformer is used in the usual way, which causes identical FSTs with different state ordering to be encoded as different things.

Author response:
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).

Reviewer comment (Theory):
The need to encode the FST as a sequence of prefix encoding limits the design of the FST topology.

Author response:
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.

Reviewer comment (Theory):
Limitation of FST which has to be deterministic.

Author response:
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.

Reviewer comment (Experiments):
The setup looks up state and transition from the embedding table, which is not scalable.

Author response:
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.

Reviewer comment (Theory):
Positional information of the transformer is used as usual, meaning two identical FSTs with different state ordering have different representations in the transformer, which could result in unwanted behavior.

Author response:
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).

Reviewer comment (Evaluation):
Question — What exactly is learned from the tunable embedding? As the author assumes no FST information is available from the downstream task, such embedding should be tuned toward certain encoding from the pre-training stage. Have the authors performed any analysis on such embedding?

Author response:
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.

Reviewer comment (Experiments):
Question — The nature of FST accepts compositional design (i.e., addition or composition of different FSTs can be easily combined). Would the proposed approach, when trained on different FST encodings, generalize well to a task solvable by their composition?

Author response:
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).