arXiv ScienceSearch

arXiv subjects

Dmitry Rybin

Publications and source records attributed to Dmitry Rybin.

6 recordsLinked to original sources

Exact Causal Attention with 10% Fewer Operations

We present Exact Causal Attention (ECA), a Strassen-style algorithm that computes exact Causal Attention using 10\% fewer operations. ECA improves a special class of matrix multiplications where either one operand or the output matrix is upper- or lower-triangular. This includes all matrix multiplication operations in the forward and backward pass of Causal Attention, such as masked product $\mathrm{Mask}(QK^{T})$. ECA is built upon algebraic identities discovered via machine learning and combinatorial search. We note that ECA cannot accelerate fused kernels such as FlashAttention on GPU. This is because ECA requires materialization of large intermediate expressions in the memory, while FlashAttention does not. However, it provides an alternative approach for compute-bound applications and can potentially be useful in scenarios with FLOPs considerations.

cs.LG

$XX^{t}$ Can Be Faster

We present RXTX, a new algorithm for computing the product of matrix by its transpose $XX^{t}$ for $X\in \mathbb{R}^{n\times m}$. RXTX uses $5\%$ fewer multiplications and $5\%$ fewer operations (additions and multiplications) than State-of-the-Art algorithms. Note that the accelerations not only holds asymptotically for large matrices with $n \rightarrow \infty$, but also for small matrices including $n = 4$. The algorithm was discovered by combining Machine Learning-based search methods with Combinatorial Optimization.

cs.DS

Finite Horizon Optimization: Framework and Applications

In modern engineering scenarios, there is often a strict upper bound on the number of algorithm iterations that can be performed within a given time limit. This raises the question of optimal algorithmic configuration for a fixed and finite iteration budget. In this work, we introduce the framework of finite horizon optimization, which focuses on optimizing the algorithm performance under a strict iteration budget $T$. We apply this framework to linear programming (LP) and propose Finite Horizon stepsize rule for the primal-dual method. The main challenge in the stepsize design is controlling the singular values of $T$ cumulative product of non-symmetric matrices, which appears to be a highly nonconvex problem, and there are very few helpful tools. Fortunately, in the special case of the primal-dual method, we find that the optimal stepsize design problem admits hidden convexity, and we propose a convex semidefinite programming (SDP) reformulation. This SDP only involves matrix constraints of size $4 \times 4$ and can be solved efficiently in negligible time. Theoretical acceleration guarantee is also provided at the pre-fixed $T$-th iteration, but with no asymptotic guarantee. On more than 90 real-world LP instances, Finite Horizon stepsize rule reaches an average 3.9$\times$ speed-up over the optimal constant stepsize, saving 75\% wall-clock time. Our numerical results reveal substantial room for improvement when we abandon asymptotic guarantees, and instead focus on the performance under finite horizon. We highlight that the benefits are not merely theoretical - they translate directly into computational speed-up on real-world problems.

math.OC

Zeroth-Order Optimization Meets Human Feedback: Provable Learning via Ranking Oracles

In this study, we delve into an emerging optimization challenge involving a black-box objective function that can only be gauged via a ranking oracle-a situation frequently encountered in real-world scenarios, especially when the function is evaluated by human judges. Such challenge is inspired from Reinforcement Learning with Human Feedback (RLHF), an approach recently employed to enhance the performance of Large Language Models (LLMs) using human guidance. We introduce ZO-RankSGD, an innovative zeroth-order optimization algorithm designed to tackle this optimization problem, accompanied by theoretical assurances. Our algorithm utilizes a novel rank-based random estimator to determine the descent direction and guarantees convergence to a stationary point. Moreover, ZO-RankSGD is readily applicable to policy optimization problems in Reinforcement Learning (RL), particularly when only ranking oracles for the episode reward are available. Last but not least, we demonstrate the effectiveness of ZO-RankSGD in a novel application: improving the quality of images generated by a diffusion generative model with human ranking feedback. Throughout experiments, we found that ZO-RankSGD can significantly enhance the detail of generated images with only a few rounds of human feedback. Overall, our work advances the field of zeroth-order optimization by addressing the problem of optimizing functions with only ranking feedback, and offers a new and effective approach for aligning Artificial Intelligence (AI) with human intentions.

cs.LG

Invariant Layers for Graphs with Nodes of Different Types

Neural networks that satisfy invariance with respect to input permutations have been widely studied in machine learning literature. However, in many applications, only a subset of all input permutations is of interest. For heterogeneous graph data, one can focus on permutations that preserve node types. We fully characterize linear layers invariant to such permutations. We verify experimentally that implementing these layers in graph neural network architectures allows learning important node interactions more effectively than existing techniques. We show that the dimension of space of these layers is given by a generalization of Bell numbers, extending the work (Maron et al., 2019). We further narrow the invariant network design space by addressing a question about the sizes of tensor layers necessary for function approximation on graph data. Our findings suggest that function approximation on a graph with $n$ nodes can be done with tensors of sizes $\leq n$, which is tighter than the best-known bound $\leq n(n-1)/2$. For $d \times d$ image data with translation symmetry, our methods give a tight upper bound $2d - 1$ (instead of $d^{4}$) on sizes of invariant tensor generators via a surprising connection to Davenport constants.

cs.LG

FFLV polytopes for odd symplectic Lie algebras

We consider "odd symplectic Lie algebras" defined in terms of maximal rank skew-symmetric forms. We provide FFLV polytopes for these algebras and prove their standard properties. In particular, we obtain a new graded character formula and PBW bases for an analogue of irreducible representations in terms of lattice points of FFLV polytopes.

math.RT