Не так давно я читал лекцию по DL в Сколтехе и затронул тему inductive bias — а что если мы будем строить архитектуру исходя из предположений о структуре данных? К примеру, картинки локальны и фичи в них инвариантны к расположению, поэтому ядра свёрток в CNN тоже локальны (ограничены своим размером) и применяются где угодно (потому что, ну, ядро одно). Так мы расплачиваемся обобщающей способностью за более лёгкое обучение модели — а значит, получаем более высокое качество на меньшем числе данных, но OOD по нам будут бить больнее.
Так как я ограничен числом параметров, применить inductive bias в дизайне моей архитектуры было логичным решением. Поэтому я решил повторить трюк с ограничением словаря из Z80-μLM и зафиксировал домен — функции на питоне. Это достаточно полезный домен, бенчмарки на нём уже есть, я смогу сравниться с бейзлайнами (бейзлайны, правда, будут из начала 20-х, но у меня игрушечный проект, мне можно) и, что самое главное, в питоне оказывается всего лишь 35 ключевых слов. Получается, чтобы я мог генерить простейшие программы на питоне, мне нужно:
- 35 токенов под ключевые слова (async, def, for, while, return, yield, etc.)
- 26 токенов под английский алфавит
- 10 токенов под цифры
- 22 токена под операторы (
:(){:|:};: и так далее)- 4 структурных токена (NEWLINE, INDENT, DEDENT, STR_QUOTE)
- 3 спецтокена (PAD, BOS, EOS)
Итого выходит 98 токенов. Разумеется, в питоне есть ещё и стандартные библиотеки, которые этими токенами не покрываются, да и на вход от юзера текст передаётся в виде текста — так что надо что-то придумать, чтобы это работало и не раздувало контекст через character based encoding.
Выход оказался прост: я могу делать AST parsing входа, доставать оттуда функции, заменять все литералы на однобуквенные и передавать так в модель. При генерации, модель генерирует однобуквенные литералы, а я потом после декодинга дополнительно детокенизирую, подставляя туда то, что пришло от юзера. Кроме того, чтобы обработать весь вход от юзера и не раздуть контекст, мы можем использовать эмбеддинг от какого-нибудь готового энкодера: он и так умеет понимать текст, он имеет длинный контекст и мы его не обучаем — так что число обучаемых эмбеддингов остаётся малым. Итого, архитектура получилась следующей: ettin-17m в качестве компрессора входа в один входной токен и крохотный декодер трансформер в качестве генератора. Кроме того, я попробовал сделать EncDec вариант — сделать кросс-аттеншн между ettin и декодером + SwiGLU + RoPE и потом учить уже такую штуку. Ну и, разумеется, я покрутил разные варианты архитектуры (число голов, хидден, глубину), растягивая и стягивая модель от 570k параметров до 1.2М.