TokenMatch: сопоставление 3D-форм за долю секунды
0,16 секунды. Именно столько TokenMatch тратит на сопоставление пары 3D-форм, включая предсказание области перекрытия, восстановление функциональной карты и поиск попарных соответствий между точками. Для сравнения: оптимизационный метод GC-PPSM решает ту же задачу за 3,1 часа, а SM-COMB за 6 минут. Разница в четыре порядка, и при этом новая модель точнее.
За работой стоит простая, но неочевидная идея: меш, как и текст, можно токенизировать. Только вместо слов в роли токенов выступают перекрывающиеся патчи поверхности, а границы между ними проводит не токенизатор BPE, а кривизна.
Что такое сопоставление 3D-форм
Сопоставление 3D-форм (shape correspondence) это задача найти семантически согласованные соответствия между точками двух трёхмерных поверхностей: понять, какая точка на одном скане человека соответствует локтю на другом скане, даже если позы разные, а часть геометрии отсутствует. Это базовый примитив геометрической обработки, на котором строятся перенос текстур и деформаций, сегментация, статистическое моделирование форм.
С полными формами в близких позах (почти изометричные деформации) методы справляются уже неплохо. Настоящая боль начинается с частичных данных. Реальный сканер почти никогда не видит объект целиком: перекрытие двух сканов неизвестно, область стыка нужно найти, а само пространство решений резко разрастается. Задача становится плохо обусловленной. Добавьте сюда сильные неизометричные деформации, когда, скажем, скан собаки нужно сопоставить со сканом лошади, и большинство обучаемых методов начинают разваливаться.
Почему? Существующие подходы либо опираются на заранее заданные дескрипторы (сырые XYZ-координаты или признаки из DINOv2, обученной на обычных фотографиях), либо на генеративные модели поверх функциональных карт. Первые тащат в 3D-задачу признаки, которые ничего не знают о геометрии поверхности. Вторые дорогие на инференсе, плохо интерпретируются и теряются на частичных формах. И почти все обучаются на маленьких однокатегорийных датасетах, откуда выносят специфические смещения, которые мешают обобщению.
Ключевая идея: токенизация по кривизне
Авторы TokenMatch заметили, что трансформеры идеально подходят для частичного сопоставления по своей природе: механизм внимания умеет работать с множествами переменного размера, а у частичных сканов количество и расположение видимых регионов заранее неизвестны. Проблема одна: у меша нет естественной токенизации. Слова в тексте разделены пробелами, пиксели картинки режутся на квадратную сетку, а треугольная сетка имеет нерегулярную структуру и произвольную плотность семплирования.
Решение: токенизировать поверхность адаптивно, глядя на геометрию. Для каждой вершины вычисляется сигнал, который на 60% состоит из абсолютной средней кривизны и на 40% из спектральной энергии, то есть суммы квадратов собственных функций Лапласа-Бельтрами в среднечастотной полосе (с 5-й по 16-ю). Низкие частоты отброшены, потому что кодируют глобальную форму, высокие отброшены, потому что чувствительны к ремешингу и шуму. Средняя полоса оказалась самой стабильной.
Дальше центры токенов выбираются дальнейшим точечным семплированием (farthest point sampling), но расстояние взвешивается этим сигналом: геометрически информативные регионы (суставы, складки, выступы) получают более плотное покрытие токенами, а плоские участки вроде спины или бедра обходятся меньшим числом. На каждую вершину затем действует мягкое назначение: принадлежность к токену затухает по геодезическому гауссиану с шириной σ=0,15, так что патчи перекрываются. Эксперименты показали, что умеренное перекрытие критично для точности: оно делает представление устойчивым к нерегулярной тесселяции и шуму дискретизации, потому что признак токена получается взвешенной агрегацией, а не жёстким разбиением. Слишком маленькое σ даёт хрупкие назначения, чувствительные к неровностям сетки, слишком большое размывает геометрию и снижает пространственную различимость. Число токенов тоже баланс: чем их больше, тем тоньше покрытие деталей, но отдача быстро затухает.
Почему именно кривизна
Авторы честно перебрали альтернативы, прежде чем остановиться на кривизне, и эта абляция объясняет дизайн лучше любых формул. Спектральная кластеризация по собственным векторам Лапласа-Бельтрами хорошо ловит глобальную структуру формы, но производит пространственно разрозненные регионы: один «патч» может состоять из кусков на разных концах меша. Кластеризация по Heat Kernel Signature устойчивее к деформациям, но границы регионов получаются размытыми и плохо прилипают к реальным изломам поверхности. Иерархическое разбиение граней даёт локально согласованные патчи, но полностью слепо к глобальной форме. Геодезическое разбиение Вороного равномерно покрывает поверхность, но не знает ни о кривизне, ни о семантике: сустав пальца и ровный участок спины получают одинаковую плотность токенов.
Токенизация по кривизне объединяет локальную детальность и глобальную согласованность: токены концентрируются вокруг высококривизненных и подвижных регионов, сохраняя при этом покрытие всей формы. Показательно и сравнение с облаками точек: без связности меша размещение токенов теряет структуру и хуже концентрируется на информативных регионах, то есть топология треугольной сетки сама по себе несёт полезный сигнал.
Архитектура: маски, внимание и функциональные карты
Поверх токенов работает трансформерный энкодер на базе ViT-Base. Self-attention моделирует внутриформенные геометрические связи, cross-attention между парой форм явно выучивает плотные соответствия. Извлечение попарных точечных соответствий из выученных признаков идёт через функциональные карты, компактное спектральное представление соответствия в базисе Лапласа-Бельтрами.
Обучение двухстадийное. Сначала masked auto-encoding в духе MeshMAE: случайные 50% токенов маскируются, энкодер видит только оставшиеся, а лёгкий декодер восстанавливает замаскированную геометрию по лоссу из расстояния Чамфера и реконструкции признаков. Именно 50% оказалось оптимумом: меньше, и задача слишком лёгкая, больше, и контекста не хватает. Эта предтренировка по сути симулирует частичность: модель учится дорисовывать недостающую геометрию из неполных наблюдений, что напрямую переносится на сценарии частичного сопоставления. После предтренировки декодер выбрасывается. Затем идёт основное обучение с лоссами на функциональную карту, предсказание перекрытия и контрастивное сближение соответствующих признаков.
Важная деталь: модель не использует шаблоны и не делает предположений о категории объектов. Можно подать гуманоида, животное, что угодно, архитектура одинаковая.
Обучение на частичных, работа на полных
Самый сильный результат статьи связан с данными. TokenMatch обучали исключительно на датасете BeCoS: более 2500 уникальных форм, гуманоиды и животные, сильные неизометричные деформации, и только в режиме «частичная форма к частичной». Полных форм в обучении не было вообще.
На тесте модель сначала подтвердила превосходство в родном режиме. По среднему IoU на трёх частичных бенчмарках TokenMatch вышел вперёд везде: CP2P 85,56 против 84,72 у лучшего конкурента EchoMatch с признаками DINOv2, PSMAL 85,21 против 84,75, BeCoS 65,25 против 64,68. Отрыв скромный, но стабильный по всем трём наборам, а методы на сырых XYZ-координатах отстают на десятки пунктов: у GC-PPSM на BeCoS всего 49,34.
Потом авторы взяли ту же обученную модель и без всякого дообучения прогнали её на полных формах. Результат неожиданный: IoU на полном сопоставлении BeCoS составил 71,85, то есть выше, чем 65,25 в родном частичном режиме. На классических бенчмарках полных форм по средней геодезической ошибке (чем ниже, тем лучше): FAUST 1,72, SCAPE 2,09, SHREC'19 3,45. На SHREC'19, сложном из-за разной связности мешей, это лучший результат среди всех методов, включая специализированные. На FAUST и SCAPE TokenMatch входит в пятёрку точнейших, уступая методам вроде SmS и ULRSSM, которые, в отличие от него, обучались именно на полных формах.
По сути, работа показывает, что сопоставление полных форм это просто частный случай частичного. Обучение на частичных данных покрывает обе задачи и снимает зависимость от датасетов полных сканов, которые собирать значительно труднее.
Устойчивость там, где конкуренты ломаются
Два стресс-теста хорошо показывают характер модели. Первый: анизотропный ремешинг, когда треугольники в меше имеют сильно неравномерный размер, с плотными и редкими регионами. Методы, завязанные на консистентную связность или локальные окрестности, здесь заметно деградируют: DenoisFM получает 2,20 на анизотропном SCAPE, AttentiveFMaps 2,80, а 3D-CODED разваливается до 16,90. TokenMatch держит 1,70, лучший результат, и 1,90 на анизотропном FAUST. Причина прямо следует из архитектуры: токенизация по кривизне ловит геометрию независимо от разрешения сетки, а трансформер работает с токенами, а не с явной связностью.
Второй тест: кросс-категориальное обобщение на животных из SMAL. Средняя геодезическая ошибка 4,10 против 4,30 у ближайших конкурентов и 4,80 у ConsistentFMaps, причём TokenMatch выигрывает на всех категориях, кроме лисы, включая сложных бегемота и льва, где вариативность формы высокая, а локальная геометрия нестабильная. Добавьте к этому устойчивость к умеренному гауссову шуму вершин без какой-либо шумовой аугментации при обучении, заслуга маскированной предтренировки.
По скорости картина такая: обучение занимает около 6 часов на узле с двумя NVIDIA A100 80GB, что сопоставимо с DPFM (4,8 ч) и EchoMatch (5,4 ч), а инференс 0,16 секунды на пару форм против 0,18 и 0,20 секунды у тех же конкурентов. Оптимизационные методы, как упоминалось, живут в другой весовой категории: минуты и часы.
Почему это важно
TokenMatch читается как шаг к фундаментальной модели для геометрического сопоставления. Одна архитектура, один набор весов, и покрыты частичное, полное и смешанное сопоставление, разные категории объектов, разный ремешинг, шум. Для практики это означает, что конвейеры переноса текстур и деформаций, статистического моделирования форм, работы с медицинскими сканами или аватарами для AR могут опираться на единый feed-forward компонент вместо зоопарка методов под каждый режим. Субсекундный инференс делает задачу пригодной для интерактивных приложений, а не только офлайн-обработки.
Есть и честно названное ограничение: вычисление геодезических расстояний дорожает на мешах высокого разрешения, выше типичного для текущих датасетов. Это следующий фронт работы.
Часто задаваемые вопросы
Чем TokenMatch отличается от EchoMatch и DPFM?
Главное отличие в признаках. DPFM и EchoMatch используют готовые дескрипторы: сырые XYZ-координаты или признаки DINOv2, обученной на фотографиях. TokenMatch выучивает геометрически нативные признаки прямо из структуры меша через токенизацию по кривизне и маскированную предтренировку, что даёт более стабильное качество при смене домена.
Нужно ли переобучать модель для полных форм?
Нет. TokenMatch обучался только на парах частичных форм, но на полном сопоставлении BeCoS показал IoU 71,85, выше, чем в частичном режиме, и лучший результат на SHREC'19 без единой эпохи дообучения. Полное сопоставление оказалось частным случаем частичного.
Насколько это быстро на практике?
0,16 секунды на пару форм на двух A100, включая все стадии: предсказание перекрытия, функциональную карту и точечные соответствия. Это быстрее нейросетевых конкурентов и на порядки быстрее оптимизационных методов, которым нужны часы на одну пару.
Итог
TokenMatch показывает, что для сопоставления 3D-форм не нужны ни рукодельные дескрипторы, ни генеративные модели, ни отдельные методы под частичные и полные данные. Достаточно правильно нарезать поверхность на токены, а правильно значит по кривизне и спектральной энергии, и дальше трансформер с маскированной предтренировкой делает остальное: точнее конкурентов, за долю секунды, с обобщением на невиданные режимы. Если вы работаете с 3D-сканами и сопоставление форм до сих пор было узким местом конвейера, за этой работой стоит следить: код и веса подобных проектов обычно появляются следом за статьёй, и попробовать модель на своих данных будет разумным следующим шагом.