The Mystery of the Pathological Path-star Task for Language Models (2024)

Abstract: The recently introduced path-star task is a minimal task designed to exemplify limitations to the abilities of language models (Bachmann and Nagarajan, 2024). It involves a path-star graph where multiple arms radiate from a single starting node and each node is unique. Given the start node and a specified target node that ends an arm, the task is to generate the arm containing that target node. This is straightforward for a human but surprisingly difficult for language models, which did not outperform the random baseline. The authors hypothesized this is due to a deficiency in teacher-forcing and the next-token prediction paradigm. We demonstrate the task is learnable using teacher-forcing in alternative settings and that the issue is partially due to representation. We introduce a regularization method using structured samples of the same graph but with differing target nodes, improving results across a variety of model types. We provide RASP proofs showing the task is theoretically solvable. Finally, we find settings where an encoder-only model can consistently solve the task.

## introduction language models (lms) have become increasingly capable of solving a variety of complex tasks @xcite @xcite . lms can do many spectacular things, making it surprising when they perform poorly or require help on simple tasks @xcite @xcite @xcite . recently, @xcite introduced one such seemingly simple task designed to showcase pathological behaviour of causal (decoder-only) autoregressive (ar) lms trained via teacher-forcing. the task is simple by design and thus failure of ar models is both surprising and informative. we begin by describing the task in sec. 1.1, before analyzing why it is difficult for lms in sec. 2. ## methods and results we solely focus on the small lm setting under the belief that such models should be able to learn such a simple task and that the biases from the pretrained data and any emergent abilities of llms will just obfuscate the root problem. we implement our models using fairseq @xcite . our ar models have 200 dim. embeddings and 6 layers, each with 800 dim. feed-forward projections and 8 heads. this is smaller than both 12-and 36-layer gpt2 models used by @xcite . we use adam with a learning rate of 0.0005, a dropout rate of 0.1, and a weight decay of 0.01. we use 16-bit training and a batch size of 1024. each model is given a maximum of 100 epochs but stops if the validation loss drops below 0.001. each experiment is trained from scratch. following bachmann and nagarajan (2024), |v | = 100 and @xmath0. each dataset for @xmath1, 3, 4, 5} is made up of 2,000,000 training and 20,000 test samples of randomly generated (g, q) pairs without overlap. unlike them, we randomly permute g at every epoch instead of just once prior to training, in an attempt to prevent overfitting. we present our work as an investigation over a series of hypotheses and corresponding experiments to get at the heart of the path-star mystery. as such, we report intermediate results and describe new methodology as it becomes motivated. we present results in order of our findings, however, we need to give some post-hoc explanations for our methodology to inform the contents of tables 1, 2, 3, 4. first, in initial experimentation (under modified task conditions), we found that the models could solve the task seemingly at random. this motivated the use of running multiple seeded trials for each experiment. for all experiments, we consider the percentage of trials that correctly succeed in learning the task across 11 trials. second, we found it necessary to set attention dropout to zero, since the task requires routing exact node information. third, we also found that we required learned-instead of sinusoidal-positional-embeddings. we suspect that the latter results in too strong of a positional bias when permuting the edges in g. we also found no embeddings works for the decoder-only model as positional information will arise out of the asymmetry induced by causal masking @xcite @xcite ## conclusion we finally have a model that consistently solves the task, however, we still have many open questions. in this section, we consider these questions as well as the progress we have made in the task. first, we consider the role sensitivity has in the task's difficulty. in sec. 2.4 we hypothesized the model's overfitting was due to sensitivity @xcite @xcite . in particular, we argued that the solution is sensitive to a single token and that the model needs to learn this in order to solve the task. however, because of the size of the sample space of graphs to targets, z, and the fact that all nodes are unique, there will be many possible spurious correlations between g and t which the model can learn and solve the task via these shortcuts. as these shortcuts are non-generalizable, the model will fail during inference. thus by using structured samples, we are exposing the underlying sensitivity of the task to the model by removing noise that would otherwise disguise this aspect of the task. the empirical results show this accounts for some of the task's difficulty but not all of it. there is another role that sensitivity may play. when we consider the question of why t is insufficient information, one reason may be that sensitivity to t is the actual cause of the difficulty. this is a conjecture the task is difficult because it is sensitive. a function is sensitive if small changes in the input cause large changes in the output. sensitive functions are known to be hard for transformers to learn @xcite @xcite . such analyses depend on a formal description of the sensitivity of the task and we found formulizing this to be challenging. we ran an initial exp. to evaluate this conjecture. teacher-forcing should decrease sensitivity via the provided ground-truth (input) support. teacherforcing is not viable for ar models. however, iar models have a viable teacher-forcing training method while nar models foregoe teacher-forcing. thus finding a performance difference between them would empirically support this conjecture. we train a nar model in exp. 26x to contrast with exp. 24x. contrary to expectation, the performance of the nar model is nearly as good as the iar model, with @xmath2. we believe this result may not hold under increased task difficulty, especially as m grows, since more samples can offer ground-truth support while bypassing the chc. as we have not formally defined sensitivity, this remains a conjecture only, with, potentially, some primary counter-evidence. one open question is why the task's difficulty increases with d even though rasp analysis indicates the solutions do not depend on d. this may be explained by z depending on d. it may also relate to the sensitivity conjecture, as d increases the number of possible targets for a given graph. another open question is why increasing the layers helps for the encoder-only model. as there were some successful trials with the smaller model, it can not be the case that the number of layers was a constraint on solving the task, however, it seems easier to find the solution with more layers. perhaps the biggest open question is why only the encoder-only model consistently solves the task. if this is solely due to the non-causal parameterization, we would expect the encoder-encoder model to perform the same as the training method is identical @xcite @xcite @xcite . as this is not the case, there must be a significant difference between these two non-causal parameterizations. this was an unexpected finding. the encoder-only model is unique in that the source-side representation conditions and is dependent on the targetside. this is not true of the decoder-only model due to the causal constraint and not true of the encoder-decoder and encoder-encoder models as both employ a source-side encoder that isolates and prevents it from conditioning on the target-side. one explanation for this behaviour is the encoder-only model may learn to write t into the correct target position and then condition on this latent variable when forming the representation of g. such an algorithm would explain why 1-step nar and m -step iar inference both produce the same results, as the iterative process may already be applied in the latent space across layers. this also relates to exp. 26x, which indicates that conditioning on ground-truths is not a critical component of such behaviour and that t may be sufficient supervision given the appropriate parameterization. the path-star task is seemingly trivial but deceptively difficult and we make large headway into solving the mystery behind this difficulty. while we do not solve it using the decoder-only model, we expose several issues that make the task harder but are not caused by teacher-forcing. we demonstrate that it is solvable in the simplified setting using arm-wise permutation. we show that the relative position of q to g induces a representation problem for decoder-only models (exp. 4 vs. 5). we explain why this works and why edge-wise permutation increases the difficulty via rasp. we also show that overfitting is an issue (exp. 4 vs. 7). we show the task can be learnt, if inconsistently, via an ar encoder-decoder model trained using teacher-forcing (exps. 14 and 15). these successes depend on using a non-causal representation of g and structured samples. thus the difficulty of the task is reduced when accounting for these hidden factors. the arm-wise exps. show a key factor to the task's difficulty is the edge-wise permutation of g and providing more structure to the graph makes the task solvable. given this, we expect the issue with next-token prediction to apply to graph tasks where the graph needs to be reconstructed from a set of shuffled edges. we are skeptical that it will apply to general tasks without these conditions. lastly, we find the encoder-only model can consistently solve the task (exp. 24x). thus the task is not pathological to teacher-forcing per se but specifically when combined with a causal ar model.