Interchangeable Token Embeddings for Extendable Vocabulary and Alpha-Equivalence
This work addresses a foundational limitation in language models for formal reasoning domains, offering incremental improvements in handling interchangeable tokens.
The paper tackled the problem of language models lacking interchangeable token representations, which limits generalization to larger vocabularies and recognition of alpha-equivalence, by proposing a dual-part token embedding strategy that improved generalization in tasks like linear temporal logic solving and propositional logic assignment prediction.
Language models lack the notion of interchangeable tokens: symbols that are semantically equivalent yet distinct, such as bound variables in formal logic. This limitation prevents generalization to larger vocabularies and hinders the model's ability to recognize alpha-equivalence, where renaming bound variables preserves meaning. We formalize this machine learning problem and introduce alpha-covariance, a metric for evaluating robustness to such transformations. To tackle this task, we propose a dual-part token embedding strategy: a shared component ensures semantic consistency, while a randomized component maintains token distinguishability. Compared to a baseline that relies on alpha-renaming for data augmentation, our approach demonstrates improved generalization to unseen tokens in linear temporal logic solving, propositional logic assignment prediction, and copying with an extendable vocabulary, while introducing a favorable inductive bias for alpha-equivalence. Our findings establish a foundation for designing language models that can learn interchangeable token representations, a crucial step toward more flexible and systematic reasoning in formal domains. Our code and project page are available at https://necrashter.github.io/interchangeable-token-embeddings