Subhomogeneous Deep Equilibrium Models

Official PyTorch implementation of Subhomogeneous Deep Equilibrium Models (ICML 2024).

Pietro Sittoni and Francesco Tudisco

Deep equilibrium (DEQ) models define their hidden state as the solution of a fixed-point equation. This repository implements SubDEQ: implicit layers whose activations are subhomogeneous, so existence, uniqueness, and linear convergence of the equilibrium follow from nonlinear Perron–Frobenius theory — without constraining the learnable weights.

The implicit layer used in the main experiments is

z* = tanh(W z*) + f_θ(x) + shift

optionally followed by an ℓ_∞ normalization. The same idea yields a nonlinear Personalized PageRank layer for graphs.

Installation

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
pip install -e .

A CUDA build of PyTorch is recommended. Graph experiments additionally need torch-sparse; install the wheel that matches your PyTorch / CUDA versions.

Quick start

Train a convolutional SubDEQ on CIFAR-10:

python train_vision.py --dataset cifar10 --architecture conv --variant tanh

Train the nonlinear APPNP / SubDEQ model on Cora citation:

python train_gnn.py --dataset cora-cit --variant tanh

Reproducing the paper

Image classification (Section 6.2)

# Main tables: MNIST (dense + conv), CIFAR-10, SVHN, Tiny ImageNet
python train_vision.py --reproduce

# Appendix activations (power-scaled tanh, absolute shifted tanh)
python train_vision.py --reproduce --appendix

Tiny ImageNet is not downloaded automatically:

python download_tiny_imagenet.py

MNIST, CIFAR-10, and SVHN are fetched by torchvision on the first run.

Dataset	Architecture	Hidden size	Epochs	Runs
MNIST	dense	87	30	5
MNIST	conv	32 channels	40	5
CIFAR-10 / SVHN	conv	48 channels	40	5
Tiny ImageNet	conv	64 channels	50	3

Variants:

• tanh — tanh(Wz) + f_θ(x) + 1.2
• normalized_tanh — ℓ_∞-normalize the same map with shift 1.603
• power_scale_tanh, abs_shift_tanh — appendix models

Graph node classification (Section 6.3)

python train_gnn.py --reproduce
python train_gnn.py --reproduce --appendix

Place the feature matrices and labels next to the hyperedges already in data_graph/:

data_graph/cora-cit/features.pickle
data_graph/cora-cit/node-labels-cora-cit.txt
data_graph/cora-author/features.pickle
data_graph/citeseer/features.pickle
data_graph/citeseer/labels.pickle
data_graph/dblp/features.pickle
data_graph/pubmed/features.pickle

Dataset	Labeled nodes	Variants
Cora citation	5.2%	tanh, normalized_tanh
Cora author	5.2%	tanh, normalized_tanh
CiteSeer	4.2%	tanh, normalized_tanh
DBLP	4.0%	tanh, normalized_tanh
PubMed	0.8%	tanh, normalized_tanh

--variant linear recovers standard APPNP. --appendix uses the held-out test split and α sweep from the appendix.

Repository layout

subdeq/
  solvers.py      Anderson acceleration and implicit-function backward
  models.py       Subhomogeneous FC / conv layers and classifiers
  data.py         MNIST, CIFAR-10, SVHN, Tiny ImageNet loaders
  train.py        Vision training loop
  gnn/            Nonlinear APPNP models and trainers
train_vision.py   CLI for image experiments
train_gnn.py      CLI for graph experiments
data_graph/       Hyperedges for the citation datasets

Citation

If you use this code, please cite:

@InProceedings{pmlr-v235-sittoni24a,
  title = 	 {Subhomogeneous Deep Equilibrium Models},
  author =       {Sittoni, Pietro and Tudisco, Francesco},
  booktitle = 	 {Proceedings of the 41st International Conference on Machine Learning},
  pages = 	 {45794--45812},
  year = 	 {2024},
  editor = 	 {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},
  volume = 	 {235},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {21--27 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://raw.githubusercontent.com/mlresearch/v235/main/assets/sittoni24a/sittoni24a.pdf},
  url = 	 {https://proceedings.mlr.press/v235/sittoni24a.html}
}

Acknowledgements

The DEQ solver follows the Deep Implicit Layers tutorial. The graph trainer is adapted from the APPNP implementation of Gasteiger, Bojchevski, and Günnemann.