WaveletViT: Discrete Wavelet Bank D-dimensional I/O Vision Transformer

waveletvit provides TensorFlow/Keras WaveletViT layers for 1D, 2D, and 3D multilevel wavelet-band tokenization, attention, and reconstruction.

Capabilities

• WaveletTokenizer1D, WaveletTokenizer2D, and WaveletTokenizer3D.
• WaveletViTBlock1D, WaveletViTBlock2D, and WaveletViTBlock3D.
• WaveletAssembler1D, WaveletAssembler2D, and WaveletAssembler3D.
• UWaveletViTNet3D for a 3D U-Net with a WaveletViT bottleneck.

Limitations

• Inputs are TensorFlow tensors with channel-last layout.
• Spatial sizes should be compatible with the configured DWT levels.

Dependencies

• tensorflow
• TFDWT

Installation

pip install waveletvit

Minimal Example

import tensorflow as tf
from waveletvit import UWaveletViTNet3D

model = UWaveletViTNet3D(
    input_shape=(32, 32, 32, 1),
    config=(16, 32, 64, 128),
    n_classes=4,
    levels=1,
    wave="haar",
)

x = tf.random.normal([1, 32, 32, 32, 1])
y = model(x)

print(y.shape)

Citation

This software is released for broad research, educational, and engineering use. If this package proves useful in related work, please cite the following thesis, whose Chapter 4 presents the underlying theory and computational details:

@misc{tarafdar2026interpretablefrugallearningsystems,
      title={Interpretable and Frugal Learning Systems Employing Multiresolution Pyramids and Volterra Kernels},
      author={Kishore Kumar Tarafdar},
      year={2026},
      eprint={2606.15011},
      archivePrefix={arXiv},
      primaryClass={eess.SP},
      url={https://arxiv.org/abs/2606.15011},
}

License

Apache License 2.0. See LICENSE.

Copyright 2026 Kishore Kumar Tarafdar.