Interpretable Lightweight Transformer via Unrolling of Learned Graph Smoothness Priors

We build interpretable and lightweight transformer-like neural networks by unrolling iterative optimization algorithms that minimize graph smoothness priors -- the quadratic graph Laplacian regularizer (GLR) and the $\ell_1$-norm graph total variation (GTV) -- subject to an interpolation constraint. The crucial insight is that a normalized signal-dependent graph learning module amounts to a variant of the basic self-attention mechanism in conventional transformers. Unlike"black-box"transformers that require learning of large key, query and value matrices to compute scaled dot products as affinities and subsequent output embeddings, resulting in huge parameter sets, our unrolled networks employ shallow CNNs to learn low-dimensional features per node to establish pairwise Mahalanobis distances and construct sparse similarity graphs. At each layer, given a learned graph, the target interpolated signal is simply a low-pass filtered output derived from the minimization of an assumed graph smoothness prior, leading to a dramatic reduction in parameter count. Experiments for two image interpolation applications verify the restoration performance, parameter efficiency and robustness to covariate shift of our graph-based unrolled networks compared to conventional transformers.

Paper

Similar papers

Peer review

Reviewer rsiC7/10 · confidence 5/52024-06-27

Summary

The authors propose a transformer-like NN by unrolling iterative optimization algorithms that minimize graph smoothness, which is used for imaging tasks.

Strengths

1. The method is well-illustrated, and the theoretical details are convinced. 2. Experiments in imaging tasks show superior experimental results compared with SOTA methods.

Weaknesses

Some more discussions on the proposed method are needed. See my questions for details.

Questions

1. The ideas of unfolding traditional graph optimization methods are popular in recent years. The authors may consider discuss more related works in GNNs, e.g., [1] Graph neural networks inspired by classical iterative algorithms Y Yang, T Liu, Y Wang, J Zhou, Q Gan, Z Wei, Z Zhang, Z Huang, D Wipf, International Conference on Machine Learning, 11773-11783. [2] Simple and deep graph convolutional networks, M Chen, Z Wei, Z Huang, B Ding, Y Li International conference on machine learning, 1725-1735 2. Could you provide a convergence guarantee for the proposed model. 3. It seems that all parameters are trained and updated following ADMM. However, in LISTA and other LISTA-like methods, most parameters are trained with back propagation. ADMM is just a way to build optimization-inspired networks, where each iteration corresponds to a single block in the network. What’s the difference between your model and these LISTA-like methods? And what’s the advantage of the proposed model.

Rating

7

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

No limitations are discussed in the paper.

Authorsrebuttal2024-08-09

Response to Reviewer rsiC on convergence of unrolled ADMM

We thank the reviewer for the 2nd round comment, which clarifies an earlier question. In our implementation, we first select the number of iterations, say K, until the model-based iterative ADMM algorithm typically converges, and then we correspondingly “unroll” the K iterations into a fixed number of K neural layers. This means that before back-prop and SGD, the unrolled neural net already achieves convergence to ADMM solution x* typically. Then, during SGD in our supervised learning setting, we empirically observe that the updated network parameters (including the graph learning parameters, conjugate gradient parameters, Lagrange multipliers, etc) only guide the converged x* towards an improved output z* that further reduces the defined end-to-end loss function, resulting in better signal interpolation performance. Thus, SGD does not affect the convergence of ADMM to a fixed point solution. We are aware of recent works that theoretically study the convergence of specific unrolled ADMM algorithms, for example: W. Pu, Y. C. Eldar and M. R. D. Rodrigues, "Optimization Guarantees for ISTA and ADMM Based Unfolded Networks," ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Singapore, Singapore, 2022, At a future work, we can similarly study theoretical convergence of our graph-based unrolled ADMM network, in the particular case when a graph learning module is periodically inserted.

Reviewer rsiC2024-08-12

Thanks for the reply. This solves my problems and I have no more questions. I will increase my rating.

Reviewer 1kQ96/10 · confidence 3/52024-07-09

Summary

This paper proposes a novel approach to build interpretable and lightweight transformer-like networks by unrolling graph-based algorithms. The authors propose unrolling iterative graph-based algorithms for signal restoration with graph smoothness priors (minimizing roughness) into interpretable neural network layers. Two graph smoothness priors are explored: graph Laplacian regularizer (GLR) and graph total variation (GTV). Experimental results demonstrate the advance of proposed methods.

Strengths

S1. The paper details the mathematical formulation for both GLR and GTV based optimization problems. S2. The proposed methods achieve good restoration performance with significantly fewer parameters compared to conventional transformer-based methods. S3. The authors show robustness of the proposed approach to covariate shifts.

Weaknesses

W1. My major concern is about the claim related to interpretability. The paper claims the proposed network is interpretable. However, the specific mechanisms for interpretability are unclear to me. Could the authors provide additional details or visualizations (e.g., activation maps) to support this claim? W2. While the paper explains the connection between the proposed graph learning module and the self-attention mechanism in transformers, there is still a large gap between these two things as the authors acked on page 7. Furthermore, It's unclear why the authors highlight the similarity. Is it because the performance advantage directly stems from this connection? Can the authors clarify the purpose of this comparison? W3. While the paper shows a reduce of proposed methods on parameter size, a runtime analysis of the proposed method compared to existing methods would be beneficial. W4. The inclusion of commonly used metrics like cPSNR and SSIM would strengthen the evaluation. Besides, adding more baselines can be more convincing. W5. The paper explores two priors, GLR and GTV. Can the authors discuss the advantages and disadvantages of each, and provide guidance on when to choose one over the other? W6. Why is there only "iGTV" in the experimental results, how about "iGLR"?

Questions

See weaknesses section

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors don't discuss the limitations of the methods

Reviewer oH216/10 · confidence 2/52024-07-12

Summary

The authors build a "white-box" Transformer-like neural network through algorithm unrolling and graph signal processing. This network utilizes convolutions to learn low-dimensional features per node, constructs sparse similarity graphs, and employs low-pass filtering at each layer, significantly reducing the parameter count compared to conventional transformers. The experimental results demonstrate its parameter efficiency and robustness to covariate shift in image interpolation tasks.

Strengths

1. This paper utilizes algorithm unrolling and graph signal processing to construct a lightweight white-box Transformer-like model.

Weaknesses

1. Considering directed graph signal processing (DGSP) is much better than graph signal processing (GSP) in reality. The later is limited to undirected graphs. Therefore, the method proposed in this paper merely approximates symmetric self-attention, where the query equals the key, and cannot mimic or surpass general self-attention. 2. The experiments are not enough to fully validate the theory.

Questions

1. Since this is a white-box Transformer-like neural network, why not use some Transformer benchmarks such as long-range arena for the experiments?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Reviewer oH212024-08-08

Thank you for the authors response. Unfortunately, the authors did not base the theoretical explanation on directed graph signal processing, which makes the theoretical explanation in this paper insufficient.

Authorsrebuttal2024-08-08

Response to Reviewer oH21 on sufficiency of undirected graph signal processing

We thank the reviewer for the 2nd round comment. With due respect, we fail to understand why theoretical analysis on undirected graphs is insufficient. We reiterate that our goal (as stated in our paper title) is to unroll an iterative optimization algorithm based on an undirected graph smoothness prior (GLR or GTV) into an interpretable “white-box” neural net, and show it is a form of transformer achieving SOTA results in signal interpolation. In so doing, for the first time in the literature, we provide a clear and intuitive connection between the plethora of undirected graph-based algorithms in the vibrant GSP community that study them and the latest deep learning architectures like transformers. Further, our low-pass filters from theoretical derivation to interpolate signals enables us to eliminate the value matrix V in conventional self-attention, resulting in drastic reduction in parameter size. OUR GOAL IS NOT TO ANALYZE THE CONVENTIONALLY DEFINED SELF-ATTENTION MECHANISM using query and key matrices in standard “black-box” transformers, which is a form of directed graph. In fact, we demonstrate that (at least for image demosaicking / interpolation) the symmetric definition of self-attention is already SUFFICIENT to achieve SOTA performance—further generalization to a directed graph is not necessary here. Thus, we are confused why the reviewer presumes that “directed graph signal processing (DGSP) is much better … in reality” (1st round comment) and our “theoretical explanation is insufficient” (2nd round comment). As already stated in our rebuttal, defining frequencies on directed graphs is still an open problem in the GSP community, and thus expecting spectral analysis of signals on directed graphs is not reasonable, in our humble opinion. For more background, see: A. G. Marques, S. Segarra and G. Mateos, "Signal Processing on Directed Graphs: The Role of Edge Directionality When Processing and Learning From Network Data," in IEEE Signal Processing Magazine, vol. 37, no. 6, pp. 99-116, Nov. 2020.

Reviewer rsiC2024-08-09

Thanks very much for the response. I know that the convergence of the ADMM algorithm has been well established by prior work. However, I wonder if the convergence of the objective still holds when all network parameters are updated by the back propagation and SGD as you mentioned.

Reviewer 1kQ92024-08-13

Thanks for the response. The authors' rebuttal has addressed my concerns, and I have no further questions.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC