Efficient Visual Representation Learning with Heat Conduction Equation

1 Introduction

The dominance of convolutional neural networks (CNNs) in computer vision has been long-standing. Transformers [Vaswani et al., 2017], originated from natural language processing (NLP), have recently demonstrated state-of-the-art performance in visual learning. Despite tremendous successes, the full self-attention has computational and memory requirement that is quadratic in the sequence length, thereby limiting its applicability to high-resolution vision tasks. Moreover, modeling 2-D images as sequences impairs interpretability. Recently, using classical physical models such as state space [Gu and Dao, 2024; Liu et al., 2024], gravitational field [Xu et al., 2022], and heat conduction equation [Chen et al., 2022] to inspire architecture design has become an effective way to explore new model architectures.

∗Corresponding Author

Our motivation is primarily driven by two key observations. First, the physical heatmap computed from heat equation and the classification heatmap learned from supervised learning are very similar, as shown in Figure 1(a). This may be due to the propagation of thermal energy through the collision of neighboring particles during heat conduction, resulting in similar temperatures in local space, which is similar to the property that neighboring regions tend to contain related information in images. Secondly, the vanilla feature extraction process maps the image to a low-dimensional feature space, which can be regarded as the compression and abstraction of the image. In this process, the model retains the most important and discriminative global information in the image, while ignoring redundant details, as shown in Figure 1(b). This mechanism is similar to heat conduction, where highfrequency information gradually decays, and low-frequency global information becomes dominant over time, as shown in Figure 1(c). Based on the above inspiration, connecting heat conduction and visual semantic propagation, we propose HcNet. This work shows that many modern model components, such as residual structures [He et al., 2016], SE block [Hu et al., 2018] and feed-forward network (FFN) [Vaswani et al., 2017], can be interpreted in terms of the heat conduction equation. From the perspective of the heat equation, the success of deep neural networks is mainly due to their ability to efficiently approximate heat conduction processes with complex initial conditions.

By establishing a connection between model architecture and the heat conduction equation, we can take advantage of the rich knowledge in the heat conduction equation to guide us in designing novel and more interpretable models. As an example, we propose the Heat Conduction Layer and the Refinement Approximation Layer inspired by solving the heat conduction equation using Finite Difference Method and Fourier series, respectively. Both layers can be explained mathematically using the concept of discrete heat conduction equation. Building upon the aforementioned components, we design a general vision backbone with a hierarchical architecture, named as Heat Conduction Network (HcNet). Compared to vision backbones with various architectures, HcNet consistently achieves competitive performance on image classification, object detection, and semantic segmentation across model scales. Specifically, HcNet-Base achieves 84.3% top-1 accuracy on ImageNet-1K, surpassing Swin [Liu et al., 2021] by 0.8%. In addition, HcNet has lower computational cost compared to ViT-based models.

Our contributions can be summarized as follows:

• We propose the Heat Conduction Layer, a network layer inspired by the finite difference solution of the heat conduction equation, which follows the heat conduction principle to transfer information with low computational complexity and high interpretability.
• We propose the Refinement Approximation Layer, a network layer inspired by the Fourier series solution of the heat conduction equation, which improves the fitting ability of the model by linearly combining the features of different channels. The Refinement Approximation Layer explains the role of the FFN layer in terms of Fourier series approximation.
• Unlike previous work that applied the heat conduction equation to a specific module of the model, HcNet attempts to guide the design of all model components through the heat conduction equation. This approach helps to introduce rich knowledge in partial differential equations into neural networks.

2 RelatedWork

Convolution Neural Networks (CNNs). In the early eras, ConvNet [LeCun et al., 1998] serves as the de-facto standard network design for computer vision. The distinctive characteristics of CNNs are encapsulated in the convolution kernels, which enjoy high computational efficiency given specifically designed GPUs. Therefore, many convolutional neural architectures [Radosavovic et al., 2020; Tan and Le, 2019; Radosavovic et al., 2020; Tan and Le, 2021] have been proposed as the vision backbone for various visual applications. To address the born limitations of local receptive fields, recent works [Chen et al., 2024; Ding et al., 2024] have been focused on developing large convolution kernels to obtain larger receptive fields.

Vision Transformers (ViTs). Transformers were proposed by Vaswani et al. [Vaswani et al., 2017] for machine translation, and have since become the state-of-the-art method in many NLP tasks. The pioneering work, Vision Transformer [Dosovitskiy et al., 2021] demonstrates that pure Transformer-based architectures can also achieve very competitive results in computer vision. One challenge for Vision Transformer-based models is data efficiency. Although ViT can perform better than convolutional networks with hundreds of millions of images for pre-training, such a data requirement is difficult to meet in many cases. To improve data efficiency, many recent works [Liu et al., 2021; Tian et al., 2023; Zhu et al., 2023; Wang et al., 2024; Zhang et al., 2024] have focused on introducing the locality and hierarchical structure into ViT. On the other hand, many works have blended attention model with convolution, and accordingly propose a series of hybrid Vision Transformers [Hatamizadeh et al., 2023; Yu et al., 2023].

Physics Inspired Models (PIMs). The principles of physics bring a novel perspective to foundation model research, taking advantage of the rich knowledge in physics to guide us in designing new and potentially more effective models. The diffusion model [Xu et al., 2022], motivated by the gravitational field, interprets the data points as electrical charges on a hyperplane and devises a backward ODE to generate the image. Mamba [Gu and Dao, 2024] models sequences based on state space, and Vision Mamba [Zhu et al., 2024; Liu et al., 2024] introduces state space into image modeling inspired by mamba. QB-Heat [Chen et al., 2022] extends the heat equation into high dimensional feature space to predict surrounding image patches for masked image modeling self-supervised task.

Unlike previous work that applied the heat conduction equation to a specific module of the model, HcNet attempts to guide the design of all model components through the heat conduction equation.

3 Methodology

3.1 Preliminaries: Heat Conduction Equation

Compared with the simple homogeneous diffusion process, the anisotropic heat equation can control the diffusion intensity in different directions, thereby enabling a more precise simulation of the complex heat conduction process. This capability establishes a foundation for the model to effectively capture complex semantic information within images. Considering that images are typically anisotropic, we employ the anisotropic heat equation [Widder, 1976] for modeling:

where u(x, y, t) is the temperature of point (x, y) at time t within a 2-D region, αx and αy denote the thermal diffusivity [Bird, 2002] of the material in the x- and y-directions, respectively.

Discretization. We use the Finite Difference Method (FDM) to solve the Heat Conduction Equation approximately. Divide the spatial domain into a grid of points. Let the grid points be (xi, yj), where i = 1, 2, · · · ,Nx and j = 1, 2, · · · ,Ny. The time domain is also discretized into time steps tn , where n = 0, 1, 2, · · · ,Nt. Substitute the finite difference approximations into the heat conduction equation:

Rearrange to solve for un+1 i,j :

3.2 Heat Conduction Layer

Suppose Z ∈ RC×H×W is the image feature map in networks, C is the number of channels, H is the height of the feature, and W is the width of the feature. We draw an analogy between feature information propagation in feature maps and temperature conduction, treating each channel as a separate 2-D heat conduction process, denoted as z ∈ RH×W .

The heat conduction layer embodies this concept of information conduction through temperature. Based on this idea, we first replace temperature u in the original heat equation with feature z, and further set the thermal diffusivity α as a learnable parameter, so that the model can adaptively adjust the diffusion behavior according to the local structure of the feature map. Since the model works with discrete image patch features, it naturally leads to a grid of points where ∆x = ∆y = ∆d. Substituting into Eq.3 yields:

where w1 = αx,1, w2 = αx,2, w3 = αy,1, w4 = αy,2, wi is a learnable parameter. Eq.4 is further simplified as follows:

where N (zn i,j ) denotes calculating the weighted sum of the neighboring points of zn i,j , W = w1+w2+w3+w4. Based on Eq.5, we propose the Heat Conduction Layer, which considers the layer dimension of the model as the time step, i.e., the feature passes through a Heat Conduction Layer and tn becomes tn+1 . Meanwhile, inspired by Mamba [Gu and Dao, 2024], we make k input-dependent to increase the contextaware ability of the model. Heat Conduction Layer:

where Z n ∈ RC×H×W , GAP:RC×H×W ↦→ RC denotes global average pooling, Linear denotes the mapping functions like fully connected layer, and k is a scalar. σ is a sigmoid function, it ensures k > 0 and limits the size of k, thereby improving the stability of the FDM solution. From the principles of heat conduction, we naturally derive the residual structure [He et al., 2016] and the SE block [Hu et al., 2018], used in modern model architectures. The overall structure of the Heat Conduction Layer is shown in Figure 2.

Eq.6 illustrates the calculation process for a single channel in the feature map Z n+1 , and the other channels are calculated in the same way.

3.3 Preliminaries: Fourier Series

Principle of Superposition. Linear partial differential equations follow the superposition principle. For the heat conduction equation, this principle states that if u1 (x, y, t) , u2 (x, y, t) , · · · , un (x, y, t) are both solutions to the heat equation, any linear combination of these solutions is also a solution:

where u1, u2, · · · , un satisfy Eq.1, u(x, y, t) is also a solution for any constants a1, a2, · · · , an. The principle of superposition allows us to construct complex solutions to the heat equation from simpler ones.

Fourier Series. The Fourier series is a powerful tool for solving the heat equation. The Fourier series solution of the heat equation in a two-dimensional rectangular space is as follows:

where m and n denote constant frequencies, Bmn is the Fourier coefficient, and L is the length of the sides of the rectangular region. Further simplification of Eq.8:

where gm (x, y, t) = sin ( mπx L ) sin ( nπy L ) ek(m2+n2 )π2 t L2 , βm = Bmn, βm is a constant. From Eq.9, we can conclude that the Fourier series fits the complex heat conduction process by linear combination of basis functions {gm (x, y, t)}Mm=1 at different frequencies, and becomes a more precise approximation as more terms are included. That is, as M increases, it can represent more intricate patterns and capture much finer details.

3.4 Refinement Approximation Layer

In Heat Conduction Layer subsection, we have illustrated the heat conduction of features on a single channel. To effectively extract visual features, it is also necessary to integrate information from multiple channels. We consider each channel as a separate 2-D heat conduction process. Therefore, each channel can be approximated using the Fourier series. According to Eq.9:

where zc (x, y, tn ) denotes the featured heat distribution of the c-th channel of the tn time-step. The accuracy of the approximation can be improved by increasing the number of terms in the Fourier series:

Based on Eq.11, we propose the Refinement Approximation Layer. Mathematically, considering the input as Z n = {zc (x, y, tn )}Cc=1 . Firstly, based on Eq.7, we use {zc (x, y, tn )}Cc=1 to construct more new solutions as candidate terms {zm (x, y, tn )}M-H m=1 . It is worth noting that the simple linear combination cannot generate new frequencies, so we perform a nonlinear transformation of {zm (x, y, tn )}M-H m=1 to introduce new frequencies to obtain new basis functions {gm}M-H m=1 . This can be done by a nonlinear activation function, taking the GELU [Hendrycks and Gimpel, 2023] as an example:

where higher-order terms (z2 , z3 ) will bring new frequencies, take z2 as an example:

New frequencies, such as 2m and 2n, are generated in Eq.13. Although nonlinear transformations can generate new higher-order harmonics, they also introduce undesirable components (e.g., DC component) and make it challenging to precisely control the amplitude of specific Fourier series terms. To address this issue, we introduce convolution operations as filters, enabling the model to adaptively select Fourier series terms with specific frequencies and amplitudes. Specifically, we input the results of the nonlinear transformation into a depth-wise separable convolution to obtain the final basis functions {gm}M-H m=1 . Extensive research has demonstrated that convolutional kernels can be regarded as filters [Mairal et al., 2014; Zoumpourlis et al., 2017]. Then, based on Eq.11, each channel uses a separate set of βm to linearly combine gm to obtain the new term M∑ m=H βmgm (x, y, tn ). Finally, add the new term to the input zc (x, y, tn ) to get ˆ zc (x, y, tn ), where the number of terms increases from H to M. Refinement Approximation Layer:znm = Linear (Z n )

where Z n denotes the feature at the n-th time-step. τ (·) denotes the nonlinear transformation, which consists of a nonlinear activation function and a depth-wise convolutional

the overall flow is shown in Figure 3. The Refinement Approximation Layer increases the number of Fourier series terms in each channel, so as the depth of the model increases, the model's ability to capture complex patterns and dynamics becomes stronger.

From the perspective of the heat equation, the role of the nonlinear activation function is to generate Fourier series terms with new frequencies, while the role of the MLP layer is to linearly combine the basis functions to approximate any function.

3.5 HcNet Model

Equipped with the above Heat Conduction layer and Refinement Approximation layer, the Hc-Ra block is shown in Figure 4 (b). We take advantage of the Hc-Ra block by adopting a hierarchical structure similar to hierarchical transformer variants [Liu et al., 2021; Zhu et al., 2023]. Figure 4 (a) shows the overall framework of HcNet. Specifically, HcNet has 4 stages, each stage contains a patch merging layer and several Hc-Ra blocks. As the network gets deeper, the input features are spatially downsampled by a certain ratio through the patch merging layer. The channel dimension is expanded twice to produce a hierarchical image representation. Specifically, the spatial downsampling ratio is set to 4 in the first stage and 2 in the last three stages. The outputs of the patch merging layer are fed into the subsequent Hc-Ra block, and the number of tokens is kept constant.

Based on this architecture, we build three different variants of HcNet, including HcNet-T, HcNet-S, and HcNet-B, as shown in Table 1.

4 Experiments

To demonstrate the effectiveness of the HcNet, we conduct experiments on ImageNet-1K [Deng et al., 2009]. To further evaluate the generalization and robustness of our backbone, we also conduct experiments on ADE20K [Zhou et al., 2017] for semantic segmentation, and COCO [Lin et al., 2014] for object detection. Finally, we perform comprehensive ablation studies to analyze each component of the HcNet.

4.1 Classification on the ImageNet-1K

Implementation details. This setting mostly follows Swin [Liu et al., 2021]. We use the PyTorch toolbox [Paszke et al., 2019] to implement all our experiments. We employ an AdamW [Kingma and Ba, 2014] optimizer for 300 epochs using a cosine decay learning rate scheduler and 20 epochs of linear warm-up. A batch size of 256, an initial learning rate of 0.001, and a weight decay of 0.05 are used. ViT-B/16 uses an image size 384×384 and others use 224×224. We include most of the augmentation and regularization strategies of Swin transformer [Liu et al., 2021] in training.

Results. Table 2 compares the performance of the proposed HcNet with the state-of-the-art CNN, Vision Transformer and PIM backbones on ImageNet-1K. Surprisingly, despite the simple heat conduction layer, HcNets can still achieve highly competitive performance compared with CNNs, ViTs and PIMs. For example, HcNet-T reaches the top-1 accuracy of 83.0 while only requiring 28M parameters and 4.1G MACs. Compared with ViT-B and Swin-T, HcNet-T outperforms them by 5.1 and 1.7, respectively. With similar parameters and MACs, HcNet performs better than other PIMs. Experimental results show that physics inspired models have the potential to achieve reasonable performance, and their interpretability is unmatched by other models.

4.2 Object Detection and Instance Segmentation

Implementation details. We use the Mask R-CNN [He et al., 2017] framework to evaluate the performance of the proposed HcNet backbone on the COCO benchmark for object detection and instance segmentation. We pretrain the backbones on the ImageNet-1K dataset and apply the finetuning strategy used in Swin Transformer [Liu et al., 2021] on the COCO training set.

Results. Table 3 reports the results of the Mask R-CNN framework with "1×" (12 training epoch) and "3 × +MS" (36 training epoch with multi-scale training) schedule. With a 12epoch fine-tuning schedule, HcNet-T/S/B models achieve object detection mAPs of 46.0%/47.6%/48.8%, outperforming Swin-T/S/B by 2.3%/2.8%/1.9% bAP, respectively. HcNets have similar performance in instance segmentation. Overall, for COCO object detection and instance segmentation, HcNets achieve competitive performance.

4.3 Semantic Segmentation

Implementation details. We further investigate the capability of HcNet for Semantic Segmentation on the ADE20K [Zhou et al., 2017] dataset. Here we employ the widely-used UperNet [Xiao et al., 2018] as the basic framework and followed Swin's [Liu et al., 2021] experimental settings. In Table 4, we report both the single-scale (SS) and multi-scale (MS) mIoU for better comparison. The default input resolution is 512 × 512.

Results. As shown in Table 4, HcNet can attain competitive performance on semantic segmentation. This further indicates the great potential of physics inspired models.

4.4 Visualization Analysis

To comprehensively evaluate the quality of the learned visual representations, we employ the GradCAM technique [Selvaraju et al., 2017] to generate activation maps for visual representations at various stages of efficientnet-b4 [Tan and Le, 2019], Swin-T, and HcNet-T. These activation maps provide insight into the significance of individual pixels in depicting class discrimination for each input image. As shown in Figure 5, the activation maps generated by HcNet are significantly different from the other two types of models, and identifying more semantically meaningful regions in deeper layers.

4.5 Ablation Studies

The experiments of ablation studies are conducted on ImageNet-1K. We discuss the ablation below according to the following aspects.

The way to get k. k = ∆t ∆d2 is the ratio of the time step to the spatial step. To demonstrate the effectiveness of input-dependent, we conduct the following experiments on ImageNet-1K: (1) fix k, (2) treat k as a learnable parameter for each layer, and (3) k in each layer depends on the input. As shown in Table 5, setting k to be input-dependent improves the performance significantly.

Filters. The generated new frequency terms are filtered using depth-wise convolution within the Refinement Approximation Layer. In this subsection, we explore the impact of depth-wise convolution on the model's performance. The model performance while keeping the same number of parameters is shown in Table 6, and the results prove the necessity of filtering the Fourier series terms that introduce new frequencies after nonlinear transformation.

5 Conclusion

In this paper, we draw a relatively comprehensive connection between modern model architectures and the heat conduction equation. Such connection enables us to design new and more interpretable neural network architectures. Benefiting from this connection, we propose the Heat Conduction Layer and the Refinement Approximation Layer inspired by solving the heat conduction equation using Finite Difference Method and Fourier series, respectively. Extensive experiments demonstrate that HcNet achieves competitive performance across various visual tasks, offering new insights for the development of physics inspired model architecture design.

Acknowledgments

This work is partially supported by National Natural Science Foundation of China (62376231), Tangshan Basic Research Science and Technology Program (23130230E), Tangshan Lunan District Science and Technology Program, Fundamental Research Funds for the Central Universities (2682025ZTPY052).