Source-linked AI summary

Riemannian Structure and Optimization for a Class of Low-Parametric Orthogonal Matrices

Ali Aliev, Maxim Rakhuba

arXiv:2609.27982v1cs.LGcs.AIcs.NAmath.DGmath.NA

TL;DR

Efficient computation for structured matrices is needed, while their geometry remains under-studied despite empirical success. The paper analyzes manifold properties and develops Riemannian tools for orthogonal GS matrices; it establishes limitations for unconstrained block factorizations and results for orthogonal cases.

  • Problem

    GS matrices have empirical success, but their geometry remains mainly unstudied and efficient computational strategies are needed.

    Source: §1 · p. 1
  • Method

    The paper studies GS manifold properties and derives Riemannian tools for orthogonal matrices, including tangent projection, retraction, and shared-parameter projection.

    Source: §1.1 · p. 2, §3.1 · p. 6, §3.3.1 · p. 11
  • Results

    Unconstrained block factorizations with rectangular or square blocks need not form smooth manifolds, while orthogonal-block constructions yield embedded submanifold results and the m = 3 case connects to the CS decomposition.

    Source: §2.1 · p. 4, §2.2 · p. 4
  • Takeaways & Limitations

    For orthogonal GS matrices, the framework provides retractions when k1 ≥ b2 and computes shared parameters by averaging the corresponding gradient blocks.

    Source: §3.3.1 · p. 11, §3.4 · p. 13

    For arbitrary sharing patterns, an element-wise 1/2 mask is not a minimal-norm solution; exact overlap components require a small precomputable least-squares problem.

    Source: §3.3.2 · p. 12

Abstract

from arXiv · show

In this paper, we are concerned with matrices formed by block-diagonal factors interleaved with fixed permutations -- a flexible family of structured matrices. This class has recently drawn interest in deep learning architectures for its balanced expressivity-efficiency trade-off, yet efficient computational strategies for working with it remain to be found. We approach this problem through Riemannian geometry and examine under what conditions this class admits a smooth manifold structure. For the practically important case of orthogonal two-factor matrices, we derive the essential Riemannian tools and propose efficient algorithms for their implementation. The algorithms leverage automatic differentiation, support parameter sharing within each factor, and avoid explicit dense matrix construction. We test them within the Riemannian optimization framework on the best matrix approximation problem and for parameter-efficient fine-tuning of large language models. Beyond the two-factor setting, we study the geometric and matrix-theoretic properties of factorizations with a larger number of block-diagonal factors.

1. Introduction.

The Group-and-Shuffle family generalizes block-diagonal and butterfly constructions, while its geometry remains understudied despite empirical success. The paper analyzes its manifold properties and develops Riemannian optimization methods, emphasizing orthogonal two-factor matrices and efficient computation.

1. Introduction.

  • GS matrices generalize butterfly and block-diagonal constructions by alternating block-diagonal factors with fixed permutations.The family unifies approaches including Monarch matrices and block-butterfly factorizations.Source: §1 · p. 1
  • Euclidean optimization ignores GS-matrix geometry, which remains mainly unstudied despite empirical success.The paper focuses on orthogonal two-factor matrices, a setting described as empirically successful and computationally friendly for deep learning.Source: §1 · pp. 1–2

1.1. Contributions.

  • The paper derives Riemannian optimization algorithms for orthogonal GS matrices, using automatic differentiation for tangent projection without dense matrix construction and supporting parameter sharing.It also examines manifold conditions and higher-factor properties, including geodesics for m = 2 under additional assumptions and a CS-decomposition connection for m = 3.Source: §1.1 · p. 2

1.2. Structure of the paper.

  • The paper is organized around manifold properties, orthogonal-case geometry and algorithms, numerical experiments, and higher-order factorization results.Section 3 includes retraction, tangent-space projection, automatic differentiation, and shared-parameter formulas.Source: §1.2 · p. 2

Related Work.

The GS class extends structured matrix approaches, while manifold optimization and orthogonal parameterizations provide relevant context for constrained optimization and parameter-efficient fine-tuning.

Related Work.

  • GS generalizes Monarch matrices by allowing flexible block sizes and varying numbers of factors.Monarch matrices combine block-diagonal factors with permutations.Source: Related Work · p. 3
  • Manifold optimization is used for constrained spaces, and multiplicative orthogonal updates such as OFT and BOFT have shown better preservation of pretrained weight structure than additive low-rank updates like LoRA.Matrix manifolds are described as helping preserve gradient norms and control Lipschitz constants.Source: Related Work · p. 3

2. Conditions for the Smooth Manifold Structure of GS-matrices.

For two-factor GS matrices, arbitrary rectangular or square blocks need not produce a smooth manifold. Orthogonal or unitary blocks give a smooth embedded manifold for every fixed permutation.

2.1. Arbitrary Rectangular and Square Blocks.

  • Arbitrary rectangular and even arbitrary square blocks can fail to yield a smooth manifold, as explicit counterexamples show.In the examples, output blocks reduce to bounded-rank sets whose Cartesian product is not smooth.Source: §2 · p. 3, §2.1 · p. 4
  • The rectangular example produces four 3×3 blocks, each a single outer product and therefore of rank at most one.Because the vectors are unconstrained, both rank zero and rank one are attainable.Source: §2.1 · p. 4
  • The paper consequently restricts the remainder of its analysis to square blocks.Source: §2.1 · p. 4

2.2. Unitary or Orthogonal blocks.

  • For two factors, orthogonal or unitary blocks yield a smooth embedded manifold for every fixed permutation, via a compact Lie group action.For three factors, block sizes and permutations allow considerable freedom; Appendix A relates perfect-shuffle permutations to the classical Cosine-Sine decomposition.Source: §2.2 · p. 4

3. Riemannian Optimization Framework for Orthogonal Blocks.

The framework optimizes GS-orthogonal matrices on their manifold by projecting gradients onto tangent spaces and retracting updates. Its projection handles disjoint or overlapping factor supports, while automatic differentiation avoids dense matrix operations.

3. Riemannian Optimization Framework for Orthogonal Blocks.

  • Retraction keeps each gradient-descent update on the manifold, while an auxiliary tangent-space function yields the Riemannian gradient without forming dense matrices.The update has the form X_{k+1}=R_{X_k}(-τ_k grad f(X_k)); differentiating the auxiliary function at zero supplies the gradient.Source: §3 · p. 5
  • The framework uses perfect-shuffle permutations, which maximize manifold dimension for fixed block sizes when k2 ≥ b1.The paper adopts this permutation in the remainder because of properties considered beneficial in neural networks.Source: §3 · p. 5

3.1. Projection onto the Tangent Space.

  • Tangent projection uses the block-diagonal skew-symmetric structure of each factor to represent admissible tangent vectors and project a Euclidean gradient.The tangent vectors to B_i are B_i C_i, where C_i is block-diagonal and skew-symmetric.Source: §3.1 · p. 6
  • When k1 ≥ b2, the two parameter supports are disjoint and the projection separates into independent subproblems.This condition is equivalent to k2 ≥ b1.Source: §3.1 · p. 7
  • When k1 < b2, the supports overlap, so projection is solved index by index, with shared entries satisfying (C2)ij + (bΩ1)ij = bZij.On overlapping indices, the method chooses the parameter pair with minimum product Frobenius norm.Source: §3.1 · p. 8

3.2. Riemannian Gradient via Automatic Differentiation.

The section derives Riemannian gradients from an auxiliary function and extends projection and automatic differentiation (AD) to tied factor blocks. It also specifies corrections for overlapping supports and arbitrary sharing patterns.

3.2. Riemannian Gradient via Automatic Differentiation.

  • Differentiating an auxiliary function on tangent parameters and applying blockwise skew-symmetric projection recovers the exact Riemannian gradient.The derivation also treats the case k1 < b2, where overlapping elements require a factor-of-two correction.Source: §3.2 · pp. 9–10

3.3. Sharing.

  • With parameter sharing, tied blocks constrain the factors and their tangent vectors; the resulting manifold remains smooth and retraction remains blockwise.Blocks may be shared across the left and right factors, but one block cannot be shared between both factors.Source: §3.3 · p. 10

3.3.1. Projection with Sharing Constraints.

  • When k1 ≥ b2, projecting onto a shared parameter yields the arithmetic mean of the gradient blocks assigned to that parameter.The block indices assigned to a unique parameter q form Jq.Source: §3.3.1 · p. 11

3.3.2. AD Implementation for Shared Parameters.

  • For shared parameters, AD sums gradients over tied instances, so each shared block gradient must be divided by its multiplicity.The shared-parameter computation treats the objective as a composition of a sharing map and the objective on full factors.Source: §3.3.1 · p. 11, §3.3.2 · p. 12
  • With overlapping supports and sharing, the projection halves gradients on the overlap set; arbitrary sharing patterns require a small, precomputable least-squares problem for exact overlap components.For arbitrary patterns, the element-wise one-half mask does not provide a minimal-norm solution.Source: §3.3.2 · p. 12

3.4. Retraction.

The section constructs a manifold retraction by applying blockwise retractions to factor tangent vectors. A dimension condition ensures the factor vectors are unique; below that condition, a minimum-norm choice resolves ambiguity.

3.4. Retraction.

  • The retraction applies standard blockwise retractions to factor tangent vectors and maps their composition back onto the orthogonal matrix manifold.The factors’ tangent vectors combine as ξ = η2P1B1 + B2P1η1.Source: §3.4 · pp. 12–13
  • When k1 ≥ b2, each tangent vector has a unique, smooth pair of factor tangent vectors, making the composed blockwise retraction valid.The composition also approximates the manifold exponential map because the factor retractions approximate their blockwise exponential maps and the factor map is a local isometry.Source: §3.4 · p. 13
  • For k1 < b2, multiple factor tangent pairs can represent the same tangent vector, so the construction needs a minimum-norm choice to define a unique map.The resulting retraction is independent of the particular factorization of the point.Source: §3.4 · pp. 13–14

3.5. Complexity Analysis.

The section gives arithmetic and storage bounds for the Riemannian optimizer, explains how fixed boundary permutations fit the framework, and develops efficient vector transport. Experiments report accuracy comparable to or better than the Cayley baseline.

3.5. Complexity Analysis.

  • A Riemannian optimization step costs O(Cf + N(b1 + b2) + u1b1^3 + u2b2^3), covering objective-gradient work, tangent projection, and retraction.Retraction acts on the ui unique blocks, while tangent projection costs O(N(b1 + b2)).Source: §3.5 · p. 14
  • Optimizer memory is O(Sf + u1b1^2 + u2b2^2), with unique-block storage and a compact overlap mask whose persistent overhead is almost zero.The mask is generated once before training.Source: §3.5 · p. 14

3.6. Fixed Boundary Permutations.

  • Fixed boundary permutations preserve the framework’s geometry as an isometry; applying them costs linear-time reindexing and permits identity initialization.The gradient is adjusted with inverse permutations before the core Riemannian optimization logic.Source: §3.6 · p. 15

3.7. Vector transport.

  • Vector transport reduces to block-diagonal operations, costing O(u1b1^3 + u2b2^3) with O(u1b1^2 + u2b2^2) temporary memory; the orthogonal framework extends to other compact Lie groups with minor adjustments.The transported parameters are recovered by enforcing skew-symmetry, with an extra mask required in the overlap case.Source: §3.7 · p. 16
  • For non-overlapping blocks with bi = ki = 32, the figures plot loss versus iterations and time, and compare Cayley Euclidean-gradient norms with projected Riemannian-gradient norms for GS-target and arbitrary-target Procrustes.Source: Figure 1 · p. 17, Figure 2 · p. 17

4. Experiments.

The experiments evaluate Riemannian optimization on synthetic GS approximation tasks and RoBERTa fine-tuning. Riemannian SGD outperforms Euclidean SGD in performance and stability, though it may not reach Adam-level results.

4.1. Synthetic tasks.

  • The Riemannian method achieves comparable or better accuracy than the Cayley-transform baseline on both GS and arbitrary target approximation tasks.The baseline uses Euclidean optimization with Cayley parameterization; the proposed method uses Riemannian gradient descent.Source: §4.1 · p. 17
  • Additional experiments cover parameter sharing, overlap, and complex cases, while smaller RoBERTa tasks explore broader configurations and block sizes.The first set is omitted from the paper with qualitatively similar results; the smaller-task settings and results are reported in Table 2.Source: §4.1 · p. 17, §4.2 · p. 18

4.2. Performance in Large-Scale Fine-Tuning.

  • The GS adapter replaces a frozen linear layer W with AW and applies a learnable diagonal scaling matrix S to the injected output.The experiments use the same optimizer, learning rate, schedule, and weight decay for S in Euclidean and Riemannian setups; only the Euclidean Cayley transform receives added regularization.Source: §4.1 · p. 17, §4.2 · pp. 17–18
  • Riemannian SGD outperforms Euclidean SGD in performance and stability on GLUE, though it may not quite reach Adam-level results.Experiments fine-tune RoBERTa with block sizes b_i = 8 and k_i = 96.Source: Table 1 · p. 18, §4.2 · p. 18
  • Riemannian methods add an average runtime overhead of 5–7% over Euclidean baselines across tasks.Vector-Jacobian products through frozen parameters dominate computation, reducing the runtime gap.Source: §4.2 · p. 18
  • The paper presents a Riemannian optimization framework for GS-orthogonal matrices, with efficient automatic-differentiation-based gradients and limited overhead.Future work includes higher-order optimizers, different block constraints, and further study of factorizations with more than two block-diagonal factors.Source: §5 · p. 18

Appendix A. GS decomposition with orthogonal blocks for m = 3.

For three-factor orthogonal matrices, the appendix studies the structured factorization class and the Perfect Shuffle case, including its expressivity and dimensional properties. It also examines the local geometry at M0, modeling the space by a determinantal cone and showing it is not a topological manifold when d is odd.

Appendix A. GS decomposition with orthogonal blocks for m = 3.

  • For arbitrary fixed permutations and potentially different block sizes, A3 ∩ O(N) = Aorth_3.The reverse inclusion is obtained by modified QR and RQ factorizations that yield orthogonal block-diagonal factors.Source: Appendix A · p. 20
  • Conjugating the middle factor by the Perfect Shuffle arranges it as a q × q grid of p × p diagonal blocks, yielding blocks M_ij = W_iD_ijV_j in the product.This explicit structure underlies the appendix’s expressivity analysis.Source: Appendix A · p. 21
  • For q = 2 and any p ≥ 2, the three-factor Perfect Shuffle decomposition gives a smooth manifold and represents every matrix in O(2p), by the classical CS decomposition.The appendix identifies this case with the classical Cosine-Sine decomposition.Source: Appendix A · p. 21
  • For p > 1 and q > 2, the three-factor Perfect Shuffle family is a proper subset of O(N), because its parameter space lacks enough degrees of freedom.The appendix compares the parameter-space dimension with dim O(N) and concludes that the deficiency is positive for q > 2.Source: Appendix A · p. 22
  • Near M0, Aorth_3 is locally homeomorphic to Re × C2,d, where e = 2q − 1 + d and C2,d consists of 2 × d matrices of rank at most one.The proof establishes that M0 and 0 are interior points in the corresponding local models.Source: Appendix A · pp. 27–28
  • For odd d, the link Ld of the rank-one cone is nonorientable, contradicting the orientability of an open Euclidean subset and showing Re × C2,d is not a manifold at (0, 0).The argument uses the cone structure to relate the link to neighborhoods of the origin.Source: Appendix A · pp. 28–29
  • When d is odd, Aorth_3 is not a topological manifold at M0.The local model and the non-manifold result for Re × C2,d together yield this conclusion.Source: Appendix A · pp. 28–29

Appendix B. Geodesics in GS-orthogonal matrices for k1 ≥b2 and m = 2.

For m = 2 and k1 ≥ b2, the standard product Frobenius metric makes the group-action map a local isometry, and in this setting geodesics of GS-orthogonal matrices are obtained from blockwise geodesics in the product group.

Appendix B. Geodesics in GS-orthogonal matrices for k1 ≥b2 and m = 2.

  • For k1 ≥ b2, the product of standard Frobenius metrics makes Φ a local isometry and a Riemannian covering map.The target inherits the ambient Frobenius metric.Source: Appendix B · p. 30
  • The appendix also gives a simpler proof that Φ is a Riemannian covering map, avoiding dependence on.The covering-map result supports the geodesic projection argument.Source: Appendix B · p. 31
  • Because G is connected and complete, GS-orthogonal geodesics are projections of its explicitly known, blockwise geodesics.The proof uses the local-isometry and completeness results stated in the appendix.Source: Appendix B · pp. 29–31
  • Fixed orthogonal boundary matrices preserve the construction: the resulting map is a global isometry, so its geodesics also come from blockwise geodesics.The section states this for fixed permutations Pin and Pout.Source: Appendix B · p. 31

Appendix C. Hierarchical and Higher-order GS-Decompositions.

The appendix analyzes hierarchical and higher-order GS decompositions, showing both flexibility in their manifold properties and practical limits of recursive constructions. It also relates nested decompositions to standard higher-order GS forms and identifies extensions to other settings.

Appendix C. Hierarchical and Higher-order GS-Decompositions.

  • Nested GS blocks can be flattened algebraically into a higher-order GS decomposition, though block-perfect-shuffle permutations need not be standard perfect shuffles.The inner block-diagonal factors concatenate into global block-diagonal factors.Source: Appendix C · p. 31
  • Recursive decomposition choices yield different decompositions without adding excess degrees of freedom, but the resulting 2 × 2-block constructions are memory-heavy and impractical for learning orthogonal matrices.For N = 2^k, recursion can continue to 2 × 2 base blocks.Source: Appendix C · pp. 31–32
  • For m = 4, the appendix constructs a matrix in A4 ∩ O(N) that is not in Aorth4, and gives an explicit counterexample for N = 6.The minimal example uses four factors, each with three 2 × 2 blocks and the standard Perfect Shuffle between layers.Source: Appendix C · pp. 32, 34
  • Adding a factor via a block-coordinate swap preserves the manifold property, and for every m ≥ 4 some decompositions have Aorth_m = A_m ∩ O(N) while others do not.The appendix establishes the equivalence of the manifold properties before stating this flexibility across decompositions.Source: Appendix C · p. 34
  • With m, permutations, and block sizes fixed, proving or disproving topological properties and sufficiency of orthogonal blocks becomes cumbersome, and simplifications are unclear.Source: Appendix C · p. 35
  • The block-structure arguments extend to some other permutation matrices and to unitary blocks, including the counterexamples and hierarchical GS decompositions.The cited property is that conjugating by the permutation yields diagonal blocks.Source: Appendix C · p. 35
  • For the broader group-action construction, the authors guarantee an immersed submanifold, not necessarily an embedded one, so its induced topology may differ from that of GL(N).The tools establish immersion but do not rule out that the set may still be embeddable.Source: Appendix D · p. 35

Appendix D. Invertible blocks.

For invertible blocks with k1 ≥ b2, the perfect shuffle maximizes the dimension of the decomposition set. The appendix notes that RGD can also be run for m = 2, but leaves that extension beyond its scope.

Appendix D. Invertible blocks.

  • For invertible blocks and k1 ≥ b2, the perfect shuffle maximizes the dimension of the decomposition set.The diagonal stabilizer has dimension at least N under this condition.Source: Appendix D · p. 35
  • For m = 2, RGD can also be run in the invertible-block case, but this extension is beyond the paper's scope.Source: Appendix D · p. 35

Appendix E. Generalization to Other Compact Matrix Groups.

The Riemannian framework extends to some other compact matrix groups with modifications to their tangent-space projections and metrics.

Appendix E. Generalization to Other Compact Matrix Groups.

  • For SO(b), the tangent space and gradient projection match the orthogonal case, and several retractions preserve determinant +1 when initialized in SO(b).Source: Appendix E · p. 35
  • For U(b), the metric uses the real Frobenius inner product and projection uses skew-Hermitian matrices; overlaps are always present.Source: Appendix E · p. 35
  • These compact-group modifications retain the algorithm's core efficiency and structural design.Source: Appendix E · p. 35
Loading 2609.27982v1…