[Paper] Approximating Matrix Functions with Deep Neural Networks and Transformers

Source: arXiv - 2602.07800v1

Overview

This paper investigates how modern deep learning architectures—especially ReLU‑based feed‑forward networks and transformer encoder‑decoders—can be trained to approximate matrix‑valued functions such as the matrix exponential and matrix sign. By providing both theoretical guarantees (network size vs. accuracy) and empirical evidence (transformers achieving ≈5 % relative error), the authors open a new avenue for using neural nets as fast, differentiable surrogates for classic numerical linear‑algebra routines.

Key Contributions

• Theoretical depth/width bounds for ReLU networks that approximate the matrix exponential to any prescribed precision.
• Empirical demonstration that a transformer encoder‑decoder, equipped with carefully designed numeric encodings, can learn a variety of matrix functions with modest relative error.
• Systematic study of encoding schemes, showing that the choice of representation (e.g., flattened entries, spectral features, positional encodings) dramatically influences learning success for different functions.
• Open‑source code and benchmark datasets for reproducibility and for the community to build upon.

Methodology

1. Problem formulation – The authors treat a matrix function (F:\mathbb{R}^{n\times n}\to\mathbb{R}^{n\times n}) as a mapping from a vector of size (n^2) (the flattened input matrix) to another vector of size (n^2).
2. ReLU network analysis – Using classical approximation theory for Lipschitz functions, they derive explicit formulas for the required depth and width of a fully‑connected ReLU network that guarantees a uniform error ≤ ε on the matrix exponential over a bounded set of inputs.
3. Transformer design – A standard encoder‑decoder architecture is adapted for numeric data:
   • Input encoding: each matrix entry is embedded with a learnable linear projection plus a positional encoding that respects the 2‑D grid structure.
   • Self‑attention: the model can capture global interactions among all entries, which is crucial for functions like the exponential that involve infinite series.
   • Output decoding: the decoder produces a sequence that is reshaped back into a matrix.
4. Training regime – Synthetic datasets are generated by sampling random matrices (e.g., from a Gaussian ensemble) and computing the exact target function via high‑precision libraries. Models are trained with mean‑relative‑error loss and evaluated on held‑out test sets.
5. Encoding comparison – The authors experiment with several schemes (raw flattening, eigenvalue‑based features, block‑wise encodings) and report which works best for each target function.

Results & Findings

• The theoretical bounds predict that a ReLU network with depth (O(\log(1/ε))) and width polynomial in (n) suffices for the exponential; experiments confirm that modest networks already hit the predicted error regime.
• Transformers outperform plain MLPs on larger matrices because self‑attention efficiently aggregates information across all entries.
• Encoding matters: a naïve flatten‑only representation leads to poor convergence for the sign function, while adding spectral cues restores performance.

Practical Implications

• Fast surrogate models: In simulation pipelines (e.g., stochastic chemical kinetics, control‑system design) where the matrix exponential must be evaluated thousands of times, a pre‑trained transformer can replace a costly eigendecomposition or scaling‑and‑squaring routine, delivering speed‑ups of 10×–100× on GPU hardware.
• Differentiable pipelines: Because the neural approximator is fully differentiable, it can be embedded inside end‑to‑end learning loops (e.g., reinforcement learning for dynamical systems) where gradients of the matrix function are required.
• Hardware‑aware computing: On edge devices lacking high‑precision BLAS libraries, a compact ReLU network can provide acceptable accuracy with a tiny memory footprint.
• Algorithmic research: The encoding insights suggest new ways to feed structured numerical data to transformers, potentially benefiting other tasks such as solving PDEs, graph‑based spectral analysis, or quantum‑state evolution.

Limitations & Future Work

• Scalability: Experiments are limited to matrices up to size (n=8); extending to larger dimensions will require architectural tweaks (e.g., hierarchical attention or low‑rank factorization).
• Error guarantees: While the ReLU analysis gives provable bounds, the transformer results are empirical; formal approximation guarantees for attention‑based models remain an open question.
• Encoding generality: The best encoding varies per function, indicating a need for automated or adaptive encoding strategies.
• Numerical stability: Neural approximators can produce outputs that violate known matrix properties (e.g., non‑orthogonality for matrix exponentials of skew‑symmetric inputs); incorporating physics‑informed constraints is a promising direction.

Bottom line: By marrying rigorous approximation theory with modern transformer architectures, this work demonstrates that deep nets can become practical, differentiable stand‑ins for classic matrix operations—potentially reshaping how developers integrate heavy linear‑algebra kernels into AI‑driven systems.*

Authors

• Rahul Padmanabhan
• Simone Brugiapaglia

Paper Information

• arXiv ID: 2602.07800v1
• Categories: cs.LG, cs.NE, math.NA
• Published: February 8, 2026
• PDF: Download PDF