Papers accepted to ICML 2026
Eight papers by the Yandex Research team (shown in bold below) and our collaborators have been accepted for publication at the International Conference on Machine Learning (ICML 2026).
On Efficient Scaling of GNNs via IO-Aware Layers Implementations (Spotlight) by Daria Fomina, Daniil Krasylnikov, Alexey Boykov, Andrey Dolgovyazov, Vyacheslav Zhdanovskiy, Fedor Velikonivtsev
In this work, we study how to efficiently scale Graph Neural Networks (GNNs), whose performance is often limited by irregular memory access patterns and excessive data movement on GPUs. We develop a set of IO-aware GPU kernels for major GNN layer families, including graph convolutions, aggregation operators, and attention-based architectures such as Graph Transformers and GATv2.
Our experiments on large-scale graph workloads demonstrate substantial speedups and memory savings compared to existing frameworks, highlighting the importance of hardware-aware optimization for scalable GNN training and inference. We also investigate the impact of node reordering and show that the benefits of placing neighboring nodes close to each other in GPU memory depend on both the graph structure and the underlying memory access pattern.
To support further research on efficient graph learning systems, we open-source all the proposed implementations.
GraphPFN: A Prior-Data Fitted Graph Foundation Model by Dmitry Eremeev, Oleg Platonov, Gleb Bazhenov, Artem Babenko, Liudmila Prokhorenkova
We introduce GraphPFN, a graph foundation model that extends the prior-data fitted network (PFN) framework to graphs. GraphPFN is pretrained on millions of specifically designed synthetic graph datasets and achieves state-of-the-art results on real-world graphs in both in-context and finetuned settings. See our blog post for more details.
TabPack: Efficient Hyperparameter Ensembles for Tabular Deep Learning by Yury Gorishniy, Akim Kotelnikov, Ivan Rubachev, Artem Babenko
In this work, we introduce TabPack, a new deep learning method for supervised learning on tabular data. TabPack is an efficiently implemented ensemble of MLPs with different model and optimizer hyperparameters, sampled at random. All MLPs are trained in parallel, and the final ensemble is built on the fly during training. This enables obtaining strong ensembles in a single run with little to no tuning, making modern tabular deep learning easier to use and more accessible on consumer-grade hardware. For example, our main TabPack experiments took less time on a modern MacBook than tuning some baselines methods on an industry-grade GPU.
Unveiling the Role of Data Uncertainty in Tabular Deep Learning by Nikolay Kartashev, Ivan Rubachev, Artem Babenko
In this work, we investigate why recent deep learning approaches for tabular data have achieved strong empirical performance across a wide range of tasks. We show that many successful design choices in modern tabular deep learning, including numerical feature embeddings, retrieval-augmented architectures, and advanced ensembling techniques, implicitly improve a model’s ability to handle high data uncertainty.
Building on this observation, we develop improved numerical feature embeddings and propose a unifying perspective on recent advances in tabular deep learning. Our results provide new insights into the mechanisms behind the success of modern tabular DL methods and suggest practical directions for designing more effective models.
Relevance-Based Embeddings: Lightweight Candidate Retrieval via Heavy-Ranker Calls by Kirill Shevkunov, Andrey Ploskonosov, Liudmila Prokhorenkova
We consider the problem of relevance search: finding the most relevant items for a given query. Relevance is usually defined by a complex and expensive similarity model, making exhaustive search infeasible.
A typical solution is to train another model that separately embeds queries and items into a vector space, where similarity is defined via the dot product or cosine similarity. This allows one to efficiently retrieve approximate nearest neighbors in this space, and then re-rank these candidates using the expensive relevance model. This approach has certain limitations, especially when individual query and item features are not expressive enough and the relevance model relies significantly on pairwise features (e.g., statistics of previous user-item interactions).
We propose an alternative approach that utilizes the scores of the expensive relevance model to improve the representations of queries and items. The idea is to describe each query (item) by its relevance to a set of support items (queries) and use these new representations to obtain query (item) embeddings. We prove that such embeddings are powerful enough to approximate any complex similarity model (under mild conditions). We also investigate the choice of support items, which is a crucial ingredient of the proposed approach. Our experiments support the claim that relevance vectors are powerful query and item representations which can be used in a variety of tasks where an arbitrary relevance function for query-item pairs is available.
Inverse Entropic Optimal Transport Solves Semi-supervised Learning via Data Likelihood Maximization by Mikhail Persiianov, Arip Asadulaev, Nikita Andreev, Nikita Starodubcev, Dmitry Baranchuk, Anastasis Kratsios, Evgeny Burnaev, Alexander Korotin
Learning conditional distributions often requires large amounts of paired data, which can be expensive or difficult to obtain in practice. In this work, we introduce a framework that combines a limited set of paired examples with additional unpaired samples in a unified semi-supervised learning setting. The approach is closely related to inverse entropic optimal transport, enabling us to leverage recent advances in computational OT to design an efficient end-to-end training procedure. We show that the method can approximate arbitrary conditional distributions and demonstrate its effectiveness on a variety of semi-supervised learning tasks.
One-Step Gradient Delay Is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining by Philip Zmushko, Egor Petrov, Nursultan Abdullaev, Mikhail Khrushchev, Samuel Horvath
Although asynchronous Pipeline Parallelism can remove pipeline bubbles and therefore substantially improve GPU utilization, it has largely been dismissed due to concerns about stale gradients and degraded convergence. In this work, we revisit asynchronous PP as a practical alternative to widely used synchronous pipeline training for large-scale LLM pretraining. We show that degradation from one-step gradient delay is not an inherent limitation of asynchronous training, but depends strongly on the optimizer. In a broad empirical study, AdamW is highly sensitive to delay, while several modern optimizers remain robust; among them, Muon offers the best practical trade-off between synchronous performance and asynchronous stability. We further propose an optimizer-agnostic Error-Feedback correction that consistently reduces the sync-async gap, and provide supporting convergence analysis for delayed Muon with and without Error Feedback. Finally, we validate the approach at scale on MoE models up to 10B parameters trained on 200B tokens. With Muon and Error Feedback, asynchronous pipeline parallelism matches the synchronous baseline loss.
Softsign: Smooth Sign in Your Optimizer for Better Parameter Heterogeneity Handling by Dmitrii Feoktistov, Timofey Belinsky, Andrey Veprikov, Amir Zainullin, Aleksandr Beznosikov
Sign-based optimization methods such as Signum and Muon have recently demonstrated strong performance across a variety of deep learning tasks. However, their updates use a fixed magnitude, making it difficult to adapt to the heterogeneous behavior of different model parameters.
To address this limitation, we introduce a smooth relaxation of sign-based optimization based on a temperature-controlled soft-sign transformation. The proposed approach enables a continuous transition between sign-like and SGD-like updates, combining the advantages of both optimization regimes.
Building on this idea, we develop SoftSignum and SoftMuon, a matrix-valued extension of the method. Across a range of deep learning workloads, including large language model pretraining, both optimizers consistently outperform standard sign-based methods and AdamW, demonstrating the benefits of adaptive update magnitudes.