Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Multiscale Training of Convolutional Neural Networks

Shadab Ahamed, Niloufar Zakariaei, Eldad Haber, Moshe Eliasof

TL;DR

Multiscale Gradient Estimation (MGE), a Multilevel Monte Carlo-inspired estimator that expresses the expected gradient on the finest mesh as a telescopic sum of gradients computed on progressively coarser meshes, and Full-Multiscale, a principled, architecture-agnostic route to accelerate CNN training on high-resolution data without sacrificing accuracy.

Abstract

Training convolutional neural networks (CNNs) on high-resolution images is often bottlenecked by the cost of evaluating gradients of the loss on the finest spatial mesh. To address this, we propose Multiscale Gradient Estimation (MGE), a Multilevel Monte Carlo-inspired estimator that expresses the expected gradient on the finest mesh as a telescopic sum of gradients computed on progressively coarser meshes. By assigning larger batches to the cheaper coarse levels, MGE achieves the same variance as single-scale stochastic gradient estimation while reducing the number of fine mesh convolutions by a factor of 4 with each downsampling. We further embed MGE within a Full-Multiscale training algorithm that solves the learning problem on coarse meshes first and "hot-starts" the next finer level, cutting the required fine mesh iterations by an additional order of magnitude. Extensive experiments on image denoising, deblurring, inpainting and super-resolution tasks using UNet, ResNet and ESPCN backbones confirm the practical benefits: Full-Multiscale reduces the computation costs by 4-16x with no significant loss in performance. Together, MGE and Full-Multiscale offer a principled, architecture-agnostic route to accelerate CNN training on high-resolution data without sacrificing accuracy, and they can be combined with other variance-reduction or learning-rate schedules to further enhance scalability.

Multiscale Training of Convolutional Neural Networks

TL;DR

Multiscale Gradient Estimation (MGE), a Multilevel Monte Carlo-inspired estimator that expresses the expected gradient on the finest mesh as a telescopic sum of gradients computed on progressively coarser meshes, and Full-Multiscale, a principled, architecture-agnostic route to accelerate CNN training on high-resolution data without sacrificing accuracy.

Abstract

Training convolutional neural networks (CNNs) on high-resolution images is often bottlenecked by the cost of evaluating gradients of the loss on the finest spatial mesh. To address this, we propose Multiscale Gradient Estimation (MGE), a Multilevel Monte Carlo-inspired estimator that expresses the expected gradient on the finest mesh as a telescopic sum of gradients computed on progressively coarser meshes. By assigning larger batches to the cheaper coarse levels, MGE achieves the same variance as single-scale stochastic gradient estimation while reducing the number of fine mesh convolutions by a factor of 4 with each downsampling. We further embed MGE within a Full-Multiscale training algorithm that solves the learning problem on coarse meshes first and "hot-starts" the next finer level, cutting the required fine mesh iterations by an additional order of magnitude. Extensive experiments on image denoising, deblurring, inpainting and super-resolution tasks using UNet, ResNet and ESPCN backbones confirm the practical benefits: Full-Multiscale reduces the computation costs by 4-16x with no significant loss in performance. Together, MGE and Full-Multiscale offer a principled, architecture-agnostic route to accelerate CNN training on high-resolution data without sacrificing accuracy, and they can be combined with other variance-reduction or learning-rate schedules to further enhance scalability.
Paper Structure (23 sections, 2 theorems, 39 equations, 10 figures, 8 tables, 2 algorithms)

This paper contains 23 sections, 2 theorems, 39 equations, 10 figures, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Multiscale Gradient Estimation
  3. Efficient Training with Multiscale Estimation of the Gradient
  4. Computational Complexity of using Multiscale Gradient Estimation.
  5. Multiscale Analysis of Convolutional Neural Networks
  6. The Full-Multiscale Training Algorithm
  7. Convergence rate for Full-Multiscale.
  8. Experimental Results and Discussion
  9. Practical guidelines for hyperparameters ($L$ and batch-size) selection.
  10. Potential challenges with extension to attention-based networks.
  11. Conclusions
  12. Proofs
  13. Proof to \ref{['lemma:standardConv_error']}
  14. Convergence of standard convolution kernels.
  15. Why coarsening is better than cropping under a multiscale training framework
  16. ...and 8 more sections

Key Result

Lemma 1

Let ${\bf u}^h, {\bf y}^h$ be continuously differentiable grid functions, and let ${\bf u}^{2h} = {\bf R}_h^{2h} {\bf u}^h$, and ${\bf y}^{2h} = {\bf R}_h^{2h} {\bf y}^h$ be their interpolation on a mesh with resolution $2h$. Let ${\bf g}^h$ and ${\bf g}^{2h}$ be the gradients of the function in eq:

Figures (10)

  • Figure 1: Illustration of our Multiscale Gradient Estimation (MGE) algorithm introduced in \ref{['sec:multilevel_SGD']}. This figure shows a schematic of a 3-level MGE algorithm with resolutions $h$ (finest), $2h$, and $4h$ (coarsest) with batch sizes $N_3 > N_2 > N_1$.
  • Figure 2: The difference between the gradients on different mesh sizes and for different noise levels $\sigma$. The difference remains small even for large noise levels.
  • Figure 3: Examples of image recovery over different tasks such as image denoising (top-left), deblurring (top-right), inpainting (bottom-left), and super-resolution (bottom-right) under different training strategies (Single-scale, Multiscale, and Full-Multiscale) using various base networks such as UNet, ResNet, and ESPCN.
  • Figure 4: Multiscale Gradient Estimation using coarsening-based subsampling of images. Here, the left figure shows the discretization of an image on a mesh of resolution $h$ (finest mesh). With subsequent coarsening, the image can be downsampled to a mesh of resolution $2h$ (and $4h, 8h$, and so on). For the sake of proving \ref{['lemma:standardConv_error']}, we define non-overlapping patches $B^{2h}_k$ on the image with resolution $2h$ (with centers $c^{2h}_k$) each of which contains four pixels from the image of resolution $h$. After coarsening, the image looks like the figure on the right with each pixel of resolution $2h$.
  • Figure 5: A comparison of different network predictions for Single-scale, Multiscale, and Full-Multiscale training for an image from the STL10 dataset for the denoising task. The first two columns display the original image and data (same for all rows), followed by results from UNet and ResNet. Here, the Multiscale and Full-Multiscale training utilize only the coarsening strategy for image subsampling.
  • ...and 5 more figures

Theorems & Definitions (6)

  • Lemma 1: Convergence of standard convolution kernels
  • Example 1
  • Theorem 1: Hotstarting SGD
  • proof
  • Definition 1: Working Unit (WU)
  • Remark 1