Fusedmax: Continuous Sparsity and Transformer Acceleration

• Fusedmax is a framework that applies fused-lasso regularization to generate sparse continuous probability densities and optimize transformer attention hardware.
• It uses total variation and Sobolev regularizers to produce interpretable, contiguous support in attention densities, improving model performance.
• Fusedmax accelerates transformer inference by mapping extended Einsum computations onto spatial arrays, reducing memory buffers and energy consumption.

Fusedmax refers to two technically distinct but thematically related concepts in the machine learning literature: (1) a sparse mapping from scores to continuous probability densities that extends fused-lasso-like regularization to infinite-dimensional domains, and (2) a mapping of fused attention computation onto spatial hardware arrays for transformer acceleration. These advances leverage structured sparsity and efficient composition of linear and nonlinear operations, resulting in interpretable attention, improved computational efficiency, and buffer/memory reductions in both continuous optimization and hardware realization contexts (Martins et al., 2021, Nayak et al., 2024).

1. Continuous Fusedmax: Problem Definition and Regularization

Continuous fusedmax, as introduced in "Sparse Continuous Distributions and Fenchel-Young Losses," generalizes the fused-lasso principle to infinite-dimensional probability densities. For a domain $S \subset \mathbb{R}$, let $$\mathcal{Y} = \{p: S \to \mathbb{R}_+ \mid \int_S p(t)dt = 1, p \in H^1(S)\}$$ denote the set of densities with integrable weak derivative $p'$. Given a real-valued score or energy function $\eta : S \to \mathbb{R}$, fusedmax is formulated as the $\Omega$-regularized prediction map:

$$\operatorname{fusedmax}(\eta) = \underset{p \in \mathcal{Y}}{\arg\min} \left[ -\int_S \eta(t) p(t) dt + \Omega(p) \right]$$

Two major regularizers $\Omega$ are considered:

• Total Variation (TV) / Rudin-Osher-Fatemi (ROF)-style: $\Omega_{TV}(p) = \gamma \int_S |p'(t)| dt$
• Sobolev (quadratic gradient): $\Omega_{S}(p) = \gamma \int_S |p'(t)|^2 dt$ where $\gamma > 0$ is the smoothing parameter. The TV regularizer encourages piecewise constant solutions, while the Sobolev regularizer induces smooth densities (Martins et al., 2021).

2. Analytical Structure and Solution Properties

The continuous fusedmax solution, when equipped with the TV regularizer, is characterized by the Euler–Lagrange or subgradient conditions, incorporating a Lagrange multiplier to enforce $\int p = 1$. The unconstrained ROF-type problem

$$\min_p \tfrac{1}{2}\|p - \eta\|^2_{L_2} + \gamma\, TV(p)$$

possesses a "taut-string" solution $u^*(t)$; imposing normalization shifts $\eta$ by $\tau$ and thresholding for nonnegativity:

$p^*(t) = [u^*(t) - \tau]_+$

For even, unimodal $\eta$, $u^*$ is constant ("clipped") on an interval $|t| \leq a$, equaling $\eta$ elsewhere. The parameters $a, \tau$ solve scalar equations that relate the "fused" (flat) interval to the TV budget $\gamma$. For Sobolev regularization, the solution reduces to solving the linear ODE $p - \gamma p'' = \eta - \tau$ under nonnegativity and normalization constraints, which yields closed-form solutions in terms of hyperbolic functions for symmetric $\eta$ (Martins et al., 2021).

Typical closed-form instances include:

• For $\eta(t) = -|t|/\sigma$: $a = \sqrt{2 \sigma \gamma}$, $\tau = -a/\sigma$, support $$=[-\sqrt{\sigma(1+2\gamma)}, \sqrt{\sigma(1+2\gamma)}]$$
• For $\eta(t) = -t^2/(2\sigma^2)$: $a = (3\sigma^2 \gamma)^{1/3}$, $$\tau = -\frac{1}{2} \big[ (3(1+2\gamma)/2\sigma)^{2/3} \big]$$

This construction yields sparse, contiguous support for the resulting density, in contrast to the diffuse support of softmax.

3. Fenchel–Young Loss for Fusedmax

For any convex regularizer $\Omega$, the Fenchel–Young loss is

$$L_\Omega(\eta; q) = \Omega^*(\eta) - \int \eta\,q + \Omega(q)$$

with $L_\Omega(\eta; p^*) = 0$ if and only if $p^* = \operatorname{fusedmax}(\eta)$, and $\Omega^*$ denotes the convex conjugate. For the TV case,

$$\Omega^*_{TV}(\eta) = \tfrac{1}{2} \|u^*(\eta) - \eta\|_2^2 + \gamma\, TV(u^*(\eta))$$

where $u^*(\eta)$ is the pre-rectification ROF solution. For the Sobolev regularizer,

$$\Omega_S^*(\eta) = \int_S \eta\,u^*(\eta) - \tfrac{1}{2} \|u^*(\eta) - \eta\|_2^2 - \tfrac{\gamma}{2} \|u^{*\prime}(\eta)\|_2^2$$

Thus, the Fenchel–Young loss embodies the regression residual plus a TV or Sobolev penalty (Martins et al., 2021).

4. Efficient Computation and Differentiation

Numerical implementation of fusedmax proceeds by discretizing $S$ on a regular grid. The discrete TV-regularized problem becomes

$$\min_{p \geq 0, \sum p = 1/h} \tfrac{1}{2} \|p - f\|^2 + h\gamma \sum_i |p_i - p_{i-1}|$$

matching Euler's finite-difference discretization of ROF denoising (see Prop. C.1). O($n$) complexity is achieved via fused-lasso solvers, notably the taut-string algorithm. The Lagrange multiplier $\tau$ is recovered by a one-dimensional root-finding on the normalization constraint.

Gradient propagation is straightforward: if $y = \operatorname{fusedmax}(\eta)$, then $\frac{\partial L_\Omega}{\partial \eta} = y - q$; one simply differentiates through the clip operation. Differentiation through the ROF-solver can leverage implicit differentiation of the KKT system (as in [28] of (Martins et al., 2021)) or a truncated unrolled primal-dual algorithm. In the Sobolev case, an additional linear ODE is solved during gradient computation (Martins et al., 2021).

5. Applications and Empirical Results

Continuous fusedmax yields interpretable and parsimonious attention densities:

• Audio classification (UrbanSound8K): Substituting standard softmax attention with continuous fusedmax attention (with $\eta$ as a learnable Gaussian score) increases accuracy by approximately 3 percentage points. The resulting attention identifies contiguous bursts of audio, suppressing isolated frames/noise.
• Visual question answering (VQA-v2): Replacing the discrete $14 \times 14$ attention grid with a single continuous fusedmax density over $[0,1]^2$—using either TV or Sobolev regularization—achieves comparable or superior accuracy with significantly fewer parameters. Fusedmax regions are observed to be compact ellipses matching queried objects.

An important advantage is enhanced interpretability: the attention density is forced to concentrate on a small number of contiguous support blocks, making the locus of "attention" in input space explicit (Martins et al., 2021).

6. FuseMax for Transformer Attention Acceleration

A separate but nomenclaturally related line of work, "FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design" (Nayak et al., 2024), proposes a hardware accelerator for transformer attention using the "cascade of Einsums" abstraction.

Key aspects include:

• Modeling attention as a cascade of extended Einsums, allowing fine-grained analysis of data access patterns and computational dependencies.
• Taxonomy via passes over input fibers: Traditional stable softmax implementations require three passes; optimizations reduce this to two or (in the case of FlashAttention-2 and FuseMax) one pass.
• Spatial array design: FuseMax maps the 1-pass softmax cascade onto a spatial array with partitioned sequence and projection dimensions, achieving buffer requirements independent of full sequence length. Only blockwise tiles need to reside on-chip.
• High utilization: Both 2D and 1D processing elements perform fused mixed operations (MACC, EXP); pipelining and tiling deliver sustained >95% utilization on both arrays for all sequence lengths.

In cycle-accurate simulation on workloads including BERT-Base and T5-small, FuseMax achieves an average attention-only speedup of 6.7$\times$ over FLAT (the prior state-of-the-art), using only 79% of the energy. End-to-end, it yields average 5.3$\times$ transformer inference speedup with 83% of the energy (Nayak et al., 2024).

7. Significance and Connections

Fusedmax provides a principled framework for structured sparsity and spatial localization within both continuous probability mapping and hardware acceleration. In the continuous case, it generalizes fused-lasso denoising to infinite-dimensional domains, retaining a well-behaved convex loss and efficient solvers, and admits interpretable sparsity for tasks such as attention-based sequence modeling. For hardware, FuseMax formalizes and fuses the core computation stages, enabling near-optimal utilization and bounded on-chip memory that breaks prior scaling laws.

These two threads illustrate the power of leveraging fusion—either in the functional or hardware domain—for both statistical expressivity and computational efficiency (Martins et al., 2021, Nayak et al., 2024).