Investigating Spectral Bias on Functional Autoencoder

[stevengogogo]

Objective

This work investigate the functional autoencoder that encode-decode functional data. This is a neural approach for functional principle component analysis. The data is randomly generated from function that is infinite-diemensional domain. I suspect the functional autoencoder has same issue as regression problem that has spectral bias issue.

Method

The functional autoencoder uses basis fucntions decomposition to encoder function:

Suppose we have a function $u: \mathbb{R}\to \mathbb{R}$, and we want to decompose $u$ into

$$ u \approx \sum_{i}^{n} c_i \phi_n(x) $$

where the basis function $\phi_i$ for $i=1,\cdots, n$ is given by a neural network $NN: \mathbb{R} \times \Omega->\mathbb{R}$

$$ \phi_i(x) = NN_i(x; \theta) $$

The score of coefficient $c_i$ is given by

$$ c_i = \langle \phi_i, u \rangle = \int \phi_i(x) u(x) dx $$

the integration is approximated by trapezoidal rule. The parameter $\theta$ is optimized by

$$ loss = MSE(u, \hat{u} $$

Also , the orthonormal property is induced by loss function.

Related work

Functional encoding on operator learning
The functional encoding is a process to use finite dimensional domain to describe infinite dimensional domain. The Basis Operator Network uses functional encoder to derive the latent vector of target/input function. The advantage of this encoding method is sensor point $x$ can be in arbitrary resolution (not good for DeepONet) and irregular spacing (not good for FNO) once trained. As long as, the inner product is successfully estimated. The complexity of encoding is $O(n)$, and high dimensional space integration can be approximated by Monte Carlo estimation. 

Multiscale

The MultiscaleDNN can be applied to help mitigate the spectral bias. There are two ideas of doing this:

Sequential layer: The autoencoder is in $l$ layer (or sequentially use same encoder), first autoencoder get $\hat{u}_i$, and second autoencoder recover $\hat{u}_2 = u-\hat{u}_1$. The multilevel structure can also use scaling to decrease frequency spectrum or customized activation function for frequency modulation. The drawback of this structure is each layer need to apply sequential.

Horizontal layer: Use scaling series $[x, ..., nx]$ as input, and output multiple basis functions. The tricky part is how to make basis function orthonormal.

Reference:

1. "Basis operator network: A neural network-based model for learning nonlinear operators via neural basis"
2. "Deep Learning for Functional Data Analysis with Adaptive Basis Layers" (arXiv)

ToDo

To complete the implementation

• Find a good problem to check the training works
• Apply orthonormal loss

To research

• add multiscale training or structure

[stevengogogo]

Exp: Functional autoencoder has spectral bias

The data is generated by cosine function with multiple radial speed.

• $u$: is the function to encode and restore
• basis function: $\phi$. The $\phi$ with most and second largest $|c_i|$ is plottted

python adafnn.py

Training: 100%|██████████| 160000/160000 [06:15<00:00, 425.79it/s, loss=36.3]

more example

[stevengogogo]

Copy post on Teams:

The preliminary result shows that the functional encoder has spectral bias. This might not be surprising since the it's still a DNN approach. I believe we can play some tricks like in MultiDNN or multi-step training on this. For example

Cascading functional encoder, and each AE learn to recover the residual function that previous AE produced. This can be same structure (like Residual NN) or different.
Scaling $[x,...,nx]$ like in multiDNN
Multi-step training the cascading functional encoder.

The ultimate goal is to produce fix dimensional latent factor of function to do operator learning with DeepONet.

[stevengogogo]

• Play with simple function

[stevengogogo]

Exp: Simple function

The training can stay in same loss for a long time, and decay for a while

[stevengogogo]

Exp: Train Functional Autoencoder up to a fixed freuqncy

Problem

$$ f(x) = \sum_{i=1}^{n} c_i \sin(i \pi x) $$

where $c_i \sim N(0,1)$.

Set $n=8$, sampling frequency is 100Hz over [0,1] equally.

The functional Autoencoder is able to encode frequency that is seen in the training

[stevengogogo]

[stevengogogo]

Exp: Traing with high frequency range

Functional autoencoder fails to train with high frequency=16

[stevengogogo]

Exp: Fourier Feature embedding

The fourier feature embedding helps training better on training region. However, extrapolation remains bad.

The method is

x->Fourier(x)->encoder-> Fourier(x)->Decoder

$$ B_{m\times D} \sim N(0, \sigma^2) $$

I chose m=100

$$ \gamma(x) = [\cos(Bx), \sin(Bx)]^T $$

[stevengogogo]

Add Orthonalgization and normalization penality

• orthogonal loss: 0.1

[stevengogogo]

Functional autoencoder can not extrapolate different frequencies.