FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness

Optimizing deep learning algorithms currently requires slow, manual derivation, potentially leaving much performance untapped. Methods like FlashAttention have achieved a x6 performance improvement over native PyTorch by avoiding unnecessary data transfers, but required three iterations over three years. Automated compiled methods have consistently lagged behind. GPUs are limited by both transfers to processors and available compute, with transfer bandwidth having improved at a far slower pace. Already, transfer bandwidth accounts for 46% of GPU energy costs. This indicates the future of energy and capital-efficient algorithms relies on improved consideration of transfer costs (IO-awareness) and a systematic method for deriving optimized algorithms. In this paper, we present a diagrammatic approach to deep learning models which, with simple relabelings, derive optimal implementations and performance models that consider low-level memory. Diagrams generalize down the GPU hierarchy, providing a universal performance model for comparing hardware and quantization choices. Diagrams generate pseudocode, which reveals the application of hardware-specific features such as coalesced memory access, tensor core operations, and overlapped computation. We present attention algorithms for Ampere, which fits 13 warps per SM (FlashAttention fits 8), and for Hopper, which has improved overlapping and may achieve 1.32 PFLOPs.

Generalizable Spacecraft Trajectory Generation via Multimodal Learning with Transformers

Effective trajectory generation is essential for reliable on-board spacecraft autonomy. Among other approaches, learning-based warm-starting represents an appealing paradigm for solving the trajectory generation problem, effectively combining the benefits of optimization- and data-driven methods. Current approaches for learning-based trajectory generation often focus on fixed, single-scenario environments, where key scene characteristics, such as obstacle positions or final-time requirements, remain constant across problem instances. However, practical trajectory generation requires the scenario to be frequently reconfigured, making the single-scenario approach a potentially impractical solution. To address this challenge, we present a novel trajectory generation framework that generalizes across diverse problem configurations, by leveraging high-capacity transformer neural networks capable of learning from multimodal data sources. Specifically, our approach integrates transformer-based neural network models into the trajectory optimization process, encoding both scene-level information (e.g., obstacle locations, initial and goal states) and trajectory-level constraints (e.g., time bounds, fuel consumption targets) via multimodal representations. The transformer network then generates near-optimal initial guesses for non-convex optimization problems, significantly enhancing convergence speed and performance. The framework is validated through extensive simulations and real-world experiments on a free-flyer platform, achieving up to 30% cost improvement and 80% reduction in infeasible cases with respect to traditional approaches, and demonstrating robust generalization across diverse scenario variations.

Categorification of Negative Information using Enrichment

In many applications of category theory it is useful to reason about "negative information". For example, in planning problems, providing an optimal solution is the same as giving a feasible solution (the "positive" information) together with a proof of the fact that there cannot be feasible solutions better than the one given (the "negative" information). We model negative information by introducing the concept of "norphisms", as opposed to the positive information of morphisms. A "nategory" is a category that has "Nom"-sets in addition to hom-sets, and specifies the compatibility rules between norphisms and morphisms. With this setup we can choose to work in "coherent" "subnategories": subcategories that describe a potential instantiation of the world in which all morphisms and norphisms are compatible. We derive the composition rules for norphisms in a coherent subnategory; we show that norphisms do not compose by themselves, but rather they need to use morphisms as catalysts. We have two distinct rules of the type $\text{morphism} + \text{norphism} \rightarrow \text{norphism}$. We then show that those complex rules for norphism inference are actually as natural as the ones for morphisms, from the perspective of enriched category theory. Every small category is enriched over $\text{P}= \langle \text{Set}, \times, 1\rangle$. We show that we can derive the machinery of norphisms by considering an enrichment over a certain monoidal category called PN(for "positive"/"negative"). In summary, we show that an alternative to considering negative information using logic on top of the categorical formalization is to "categorify" the negative information, obtaining negative arrows that live at the same level as the positive arrows, and suggest that the new inference rules are born of the same substance from the perspective of enriched category theory.

Task-driven Modular Co-design of Vehicle Control Systems

When designing autonomous systems, we need to consider multiple trade-offs at various abstraction levels, and the choices of single (hardware and software) components need to be studied jointly. In this work we consider the problem of designing the control algorithm as well as the platform on which it is executed. In particular, we focus on vehicle control systems, and formalize state-of-the-art control schemes as monotone feasibility relations. We then show how, leveraging a monotone theory of co-design, we can study the embedding of control synthesis problems into the task-driven co-design problem of a robotic platform. The properties of the proposed approach are illustrated by considering urban driving scenarios. We show how, given a particular task, we can efficiently compute Pareto optimal design solutions.

Posetal Games: Efficiency, Existence, and Refinement of Equilibria in Games with Prioritized Metrics

Modern applications require robots to comply with multiple, often conflicting rules and to interact with the other agents. We present Posetal Games as a class of games in which each player expresses a preference over the outcomes via a partially ordered set of metrics. This allows one to combine hierarchical priorities of each player with the interactive nature of the environment. By contextualizing standard game theoretical notions, we provide two sufficient conditions on the preference of the players to prove existence of pure Nash Equilibria in finite action sets. Moreover, we define formal operations on the preference structures and link them to a refinement of the game solutions, showing how the set of equilibria can be systematically shrunk. The presented results are showcased in a driving game where autonomous vehicles select from a finite set of trajectories. The results demonstrate the interpretability of results in terms of minimum-rank-violation for each player.

A Compositional Sheaf-Theoretic Framework for Event-Based Systems

A compositional sheaf-theoretic framework for the modeling of complex event-based systems is presented. We show that event-based systems are machines, with inputs and outputs, and that they can be composed with machines of different types, all within a unified, sheaf-theoretic formalism. We take robotic systems as an exemplar of complex systems and rigorously describe actuators, sensors, and algorithms using this framework.

Co-Design of Autonomous Systems: From Hardware Selection to Control Synthesis

Designing cyber-physical systems is a complex task which requires insights at multiple abstraction levels. The choices of single components are deeply interconnected and need to be jointly studied. In this work, we consider the problem of co-designing the control algorithm as well as the platform around it. In particular, we leverage a monotone theory of co-design to formalize variations of the LQG control problem as monotone feasibility relations. We then show how this enables the embedding of control co-design problems in the higher level co-design problem of a robotic platform. We illustrate the properties of our formalization by analyzing the co-design of an autonomous drone performing search-and-rescue tasks and show how, given a set of desired robot behaviors, we can compute Pareto efficient design solutions.

A Formal Approach to the Co-Design of Embodied Intelligence

We consider the problem of formally co-designing embodied intelligence as a whole, from hardware components such as chassis and sensors to software modules such as control and perception pipelines. We propose a principled approach to formulate and solve complex embodied intelligence co-design problems, leveraging a monotone co-design theory. The methods we propose are intuitive and integrate heterogeneous engineering disciplines, allowing analytical and simulation-based modeling techniques and enabling interdisciplinarity. We illustrate through a case study how, given a set of desired behaviors, our framework is able to compute Pareto efficient solutions for the entire hardware and software stack of a self-driving vehicle.