A Muon-Accelerated Algorithm for Low Separation Rank Tensor Generalized Linear Models

TL;DR

This paper introduces LSRTR-M, a scalable algorithm for tensor generalized linear models that accelerates convergence using Muon updates, improving efficiency and accuracy in high-dimensional tensor data analysis.

Contribution

It proposes a novel Muon-accelerated method for low separation rank tensor GLMs, enhancing scalability and convergence speed over existing approaches.

Findings

LSRTR-M converges faster than traditional methods in synthetic experiments.

LSRTR-M achieves lower estimation and prediction errors.

On Vessel MNIST, it improves computational efficiency while maintaining accuracy.

Abstract

Tensor-valued data arise naturally in multidimensional signal and imaging problems, such as biomedical imaging. When incorporated into generalized linear models (GLMs), naive vectorization can destroy their multi-way structure and lead to high-dimensional, ill-posed estimation. To address this challenge, Low Separation Rank (LSR) decompositions reduce model complexity by imposing low-rank multilinear structure on the coefficient tensor. A representative approach for estimating LSR-based tensor GLMs (LSR-TGLMs) is the Low Separation Rank Tensor Regression (LSRTR) algorithm, which adopts block coordinate descent and enforces orthogonality of the factor matrices through repeated QR-based projections. However, the repeated projection steps can be computationally demanding and slow convergence. Motivated by the need for scalable estimation and classification from such data, we propose LSRTR-M, which incorporates Muon (MomentUm Orthogonalized by Newton-Schulz) updates into the LSRTR framework. Specifically, LSRTR-M preserves the original block coordinate scheme while replacing the projection-based factor updates with Muon steps. Across synthetic linear, logistic, and Poisson LSR-TGLMs, LSRTR-M converges faster in both iteration count and wall-clock time, while achieving lower normalized estimation and prediction errors. On the Vessel MNIST 3D task, it further improves computational efficiency while maintaining competitive classification performance.