Details

Настоящая работа посвящена разработке и исследованию модели трансформера с механизмом группового внимания для задач регрессии на табличных данных с групповой структурой. Предметом исследования являются методы и архитектуры моделей машинного обучения, использующие групповую структуру данных для повышения качества регрессионных предсказаний и оценки эффекта воздействия. Предложен механизм группового внимания, расширяющий стандартный multi-head attention за счёт аддитивной и мультипликативной модуляции весов внимания на основе матрицы групповых отношений. Матрица отношений строится локально для каждого батча на основе категориальных признаков, задающих групповую принадлежность объектов. Поддерживается несколько независимых каналов групповых отношений. Проведены эксперименты на синтетических и реальных наборах данных. Сравнение с трансформером без группового механизма и методами градиентного бустинга (XGBoost, CatBoost) показало, что предложенная модель превосходит конкурентов по метрике R^2 в большинстве рассмотренных конфигураций, демонстрируя особенно высокую устойчивость к шуму в признаках. На реальных данных с таксономической групповой структурой (Bird Genetic Diversity) трансформер с групповым вниманием достиг R^2 = 0.61 против 0.44 у градиентного бустинга. Дополнительно предложены два подхода к адаптации модели для задачи оценки условного среднего эффекта воздействия: двухэтапный, в котором T-learner формирует псевдо-аутком как начальное приближение CATE, а трансформер уточняет оценку используя индикатор воздействия в качестве группового канала механизма внимания, и двухголовая архитектура с двумя независимыми выходными слоями, обучаемая непосредственно на наблюдаемых исходах без вспомогательных моделей.

This work is devoted to the development and research of a transformer model with a group attention mechanism for regression tasks on tabular data with group structure. The subject of the research is methods and model architectures that exploit the group structure of data to improve regression predictions and treatment effect estimation. A group attention mechanism is proposed that extends the standard multi-head attention through additive and multiplicative modulation of attention scores based on a matrix of group relations. The relation matrix is constructed locally for each batch using categorical features that define group membership. Multiple independent channels of group relations are supported. Experiments were conducted on synthetic and real-world datasets. Comparison with a transformer without group mechanism and gradient boosting methods (XGBoost, CatBoost) showed that the proposed model outperforms competitors in terms of R^2 in most configurations, demonstrating particularly high robustness to feature noise. On real-world data with taxonomic group structure (Bird Genetic Diversity), the transformer with group attention achieved R^2 = 0.61 compared to 0.44 for gradient boosting. Additionally, two approaches to adapting the model for conditional average treatment effect estimation are proposed: a two-stage approach in which a T-learner generates a pseudo-outcome as an initial CATE approximation and the transformer refines the estimate using the treatment indicator as a group channel, and a two-head architecture with two independent output layers trained directly on observed outcomes without auxiliary models. The results can be applied in medicine, economics and biology where data exhibit natural group structure.