Overcoming prompt sensitivity by tokenization robust self-distillation

Recently, Steger et al. showed that exploiting stochasticity in pretraining and fine-tuning results in models that are much less prone to adversarial attacks and prompt sensitivity. Besides, Sims et al. argued that it has benefits in some word game and reasoning-related domains.

However, efficiently applying it to post-training is non-trivial, as common RL approaches lack sample efficiency, which will be amplified by stochastic-induced tokenization variance. On-policy distillation (and its recent extension, self-distillation) provide much more sample-efficient approaches due to Rao-Blackwellization of the estimator. The only obstacle is that naive distillation assumes the same tokenization between teacher and student models.

By exploiting novel findings in cross-tokenizer distillation, we aim to bridge this gap. The project will explore whether stochastic tokenization at post-training and distillation could serve as a theoretically justified option to perform data augmentation without changing the underlying measure on the sequence space, and to which extent does it result in less prompt sensitive and more robust models beyond SFT regime.

Success of the project would bring us an ability to embed robust training into existing RL pipelines and combine with arbitrary novel distillation-based approaches.

Concrete research outputs will include RL training pipelines, reproducible evaluations, open-sourced code and checkpoints, and a research paper for A* conference with findings and recomendations for practitioners in the field.

Important:

The grant will fund execution, not an exploration phase. We've already written most of the theory and algorithms, and tested them on small-scale models.

We intentionally restrict the length of the project to 3-4 months to conduct detailed experiments and deliver a write-up researchers in the field can build upon.

Project members:

Sergei Kudriashov -- core technical contributor and scientific driver of the project, responsible for writing, experiment design and theory.

CAISH research fellow, MARS V; postgraduate student with both industrial and academic research experience. Researcher with multiple publications at top venue ML conferences (ICML, ECCV, EACL). Mentored multiple student projects on RL and adversarial robustness.
Ex-SBER, worked on risk assessment for coding agents and internal production LLM development, evals and benchmarking in financially critical applications.

Jamilia Erkenova -- secondary contributor, responsible for AI infrastructure, conducting experiments and general codebase, besides core algorithmic innovations.

Mentor on technical AI safety courses, data engineer with 2+ years of production experience in a large financial company. Developer of benchmarks for VLM models.