Understanding Neural Operators: The Next Evolution in Machine Learning

Have you ever wondered how we could teach computers to understand continuous processes, like weather patterns or fluid dynamics, rather than just working with fixed data points? Enter neural operators - a fascinating breakthrough in machine learning that's changing how we approach complex scientific problems. Unlike the traditional neural networks, neural operators learn mappings between function spaces. In some sense, we can view the traditional neural networks as calculators that can only work with specific numbers you input. Neural operators, on the other hand, are more like mathematical wizards that can understand and work with entire functions - continuous patterns that exist across space and time. Imagine being able to predict weather patterns at any location, not just where weather stations exist!

Key Aspects of Neural Operators

Breaking Free from Resolution Constraints

One of the most exciting aspects of neural operators is their "resolution agnostic" nature. A neural operator, with a fixed set of parameters, can be applied to input functions given at any discretization. This means that as the discretization of input functions is refined, the outputs converge to the true solution, differing only by a discretization error. This property is a significant advantage over standard neural networks, as they do not have guarantees of generalizing to other resolutions, and often perform poorly when interpolated to higher resolutions. 

The Secret Sauce: Integral Operators

Neural operators are built using linear integral operators, followed by non-linear pointwise activations. The linear integral operator involves a learnable kernel that maps between input and output domains. 

o The integral operation is given by: ∫ k(x, y)a(y)dy ≈ ∑ k(x, yi)a(yi)∆yi, where a(·) is the input function, and k(x, y) is a learnable kernel between any two points x and y.

o The query point x in the output domain does not need to be limited to the discrete training grid and can be any point in the continuous domain.

Zero-Shot Super-resolution and Super-evaluation

Due to the discretization convergence property, trained neural operators can perform zero-shot super-resolution, where the output is predicted at a higher resolution than what was seen during training, and zero-shot super-evaluation, where the operator can be evaluated on a new, finer discretization than seen during training. 

Fourier Neural Operator (FNO)

An example of neural operators is the Fourier Neural Operator (FNO). FNO consists of one or multiple Fourier layers that learns and emulates the interactions among the variables of interest in Fourier space. These layers are sandwiched by two linear transformation layers that convert the dimensions between inputs, hidden, and output layers. Fourier layers are analogus to convolution layers in the convolutional neural networks (CNN). However, filtering in convolution neural networks are usually local as shown in the figure below. They are good to capture local patterns such as edges and shapes. The filtering by FNOs are global sinusoidal functions. They are better for representing continuous functions and thus for learning mappings between continuous functions. 

CNN filters versus Fourier filters 
Image courtesy of https://zongyi-li.github.io/blog/2020/fourier-pde/

Physics-Informed Neural Operators (PINOs) 

This is another example of popular neural operators. These operators incorporate physics constraints, such as PDEs, into the training process. This leads to improved generalization, extrapolation capabilities and reduced training data requirements compared to purely data-driven neural operators. Please see my earlier post on Physics-informed Neural Networks.

Neural Operators Applications

Neural operators have demonstrated success in various scientific and engineering applications, such as solving PDEs, fluid dynamics, weather forecasting, climate modeling, and inverse problems. As we push the boundaries of scientific discovery, neural operators represent a significant leap forward in our ability to model and understand complex systems. They're not just another machine learning tool - they're a bridge between discrete digital computing and the continuous nature of our physical world. For further exploration, please visit the neural operators library where you can download jupyter notebooks illustrating the their use.

Physics Informed Neural Networks: Bridging Machine Learning and Scientific Computing

Physics Informed Neural Networks (PINNs) represent a groundbreaking approach to solving complex physical problems by combining the power of neural networks with our knowledge of physical laws. In this post, we'll explore what PINNs are, how they work, and implement a simple example to solve a differential equation.

Understanding PINNs

Traditional neural networks learn patterns from data alone. PINNs go a step further by incorporating physical laws directly into the learning process. They do this by adding physics-based constraints to the loss function, ensuring that the network's predictions not only fit the data but also satisfy known physical equations. Thus, the loss function used in PINNs consists of two components. The first component is the commonly used data loss measure that measures how well the network fits the available data. The second component consists of the physics loss measuring how well the network satisfies the governing physical equations. For example, if we're solving a differential equation du/dt = f(u,t), the physics loss would include terms that measure how far our predicted solution is from satisfying this equation.

Key Advantages of PINNs:

1. They require fewer training data points compared to traditional neural networks

2. Solutions automatically satisfy physical constraints

3. They can handle both forward and inverse problems

4. Capable of solving complex partial differential equations (PDEs)

Implementing a Simple PINN

Let's implement a PINN to solve a basic ordinary differential equation (ODE):

du/dt = -u, u(0) = 1

This is the equation for exponential decay, with the analytical solution u(t) = exp(-t). The initial condition is specified as u(0) = 1. The code is shown below. We use a simple feedforward neural network with tanh activation functions. The input is time t, and the output is our solution u(t). Our loss combines two terms:

   - Physics loss: Measures how well our solution satisfies du/dt = -u

   - Initial condition loss: Ensures u(0) = 1

We use PyTorch's autograd to compute du/dt, which is needed for the physics loss. The network is trained using the Adam optimizer to minimize the combined loss. We also include code segments for visualization.

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

class PINN(nn.Module):
    def __init__(self):
        super().__init__()
        # Neural network architecture
        self.net = nn.Sequential(
            nn.Linear(1, 20),
            nn.Tanh(),
            nn.Linear(20, 20),
            nn.Tanh(),
            nn.Linear(20, 1)
        )
    
    def forward(self, t):
        return self.net(t)
    
    def loss_function(self, t, u):
        # Compute du/dt using autograd
        u_pred = self.forward(t)
        u_t = torch.autograd.grad(
            u_pred, t,
            grad_outputs=torch.ones_like(u_pred),
            create_graph=True
        )[0]
        
        # Physics loss: du/dt + u = 0
        physics_loss = torch.mean((u_t + u_pred)**2)
        
        # Initial condition loss: u(0) = 1
        ic_loss = torch.mean((self.forward(torch.zeros_like(t)) - 1.0)**2)
        
        return physics_loss + ic_loss, physics_loss.item(), ic_loss.item()

# Training setup
model = PINN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
t = torch.linspace(0, 5, 100, requires_grad=True).reshape(-1, 1)

# Lists to store loss history
total_losses = []
physics_losses = []
ic_losses = []

# Training loop
n_epochs = 5000
for epoch in range(n_epochs):
    optimizer.zero_grad()
    total_loss, physics_loss, ic_loss = model.loss_function(t, None)
    total_loss.backward()
    optimizer.step()
    
    # Store losses
    total_losses.append(total_loss.item())
    physics_losses.append(physics_loss)
    ic_losses.append(ic_loss)
    
    if (epoch + 1) % 1000 == 0:
        print(f'Epoch {epoch+1}, Total Loss: {total_loss.item():.6f}, '
              f'Physics Loss: {physics_loss:.6f}, IC Loss: {ic_loss:.6f}')

# Create subplots for solutions and loss convergence
plt.figure(figsize=(15, 6))

# Plot 1: Solution comparison
plt.subplot(1, 2, 1)
with torch.no_grad():
    t_plot = torch.linspace(0, 5, 100).reshape(-1, 1)
    u_pred = model(t_plot)
    u_true = torch.exp(-t_plot)
    
    plt.plot(t_plot, u_pred, 'b-', label='PINN prediction')
    plt.plot(t_plot, u_true, 'r--', label='True solution')
    plt.xlabel('t')
    plt.ylabel('u(t)')
    plt.legend()
    plt.title('PINN Solution vs True Solution')
    plt.grid(True)

# Plot 2: Loss convergence
plt.subplot(1, 2, 2)
epochs = range(1, n_epochs + 1)
plt.semilogy(epochs, total_losses, 'b-', label='Total Loss')
plt.semilogy(epochs, physics_losses, 'r--', label='Physics Loss')
plt.semilogy(epochs, ic_losses, 'g-.', label='IC Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss (log scale)')
plt.legend()
plt.title('Loss Convergence')
plt.grid(True)

plt.tight_layout()
plt.show()
       
 

When run, the code above produces a plot comparing the PINN's solution to the analytical solution exp(-t). As seen, the PINN typically learns to approximate the true solution very well, even though we never explicitly told it the analytical solution. 

The loss convergence plot reveals several interesting aspects of the training process:

Initial Phase

- The total loss starts relatively high as the network's predictions are far from satisfying both the physics and initial conditions

- Both physics and initial condition (IC) losses contribute significantly to the total loss

Middle Phase

- We observe a rapid decrease in all loss components as the network learns to satisfy both constraints

- The physics loss typically takes longer to converge than the IC loss, as it needs to satisfy the differential equation across the entire domain

Final Phase

- The losses stabilize as the network finds a solution that satisfies both the physics and initial conditions

- Small fluctuations may persist due to the optimization process and the precision limits of our network.

Conclusion

Physics Informed Neural Networks represent a powerful fusion of machine learning and scientific computing. They allow us to solve complex physical problems while respecting underlying physical laws, often with less data than traditional approaches would require.

As the field continues to develop, we're seeing PINNs being applied to increasingly complex problems, from turbulent flows to quantum systems. Their ability to incorporate physical knowledge into the learning process makes them a valuable tool in scientific computing and engineering.

ModernBERT: The Evolution of Language Understanding in AI

The world of artificial intelligence has taken another leap forward with ModernBERT, an advanced evolution of the revolutionary BERT (Bidirectional Encoder Representations from Transformers) language model that Google AI Language introduced in 2018. Building on BERT's groundbreaking ability to understand context in human language, ModernBERT brings powerful new capabilities to the table.

What Makes ModernBERT Special?

ModernBERT isn't just a simple upgrade - it's a significant advancement in how AI understands and processes language. The model comes in two sizes: a base version with 149 million parameters and a larger version with 395 million parameters. But what really sets it apart is its ability to handle much longer pieces of text - up to 8,192 tokens at once!

Key Innovations and Improvements

ModernBERT introduces several game-changing features:

- Extended context length for better understanding of longer texts

- Rotary positional embeddings (RoPE) for improved word placement awareness

- Enhanced activation functions through GeGLU layers. GEGLU is a novel activation function which is a    variant of the Gated Linear Unit (GLU) and Generalized Linear Unit (GELU) activations designed    to    address some of their limitations

- Flexible, modular design that can be customized for specific needs

Real-World Applications of Modern BERT

ModernBERT shines in several key areas:

Code Search and Development

Developers can use ModernBERT to quickly find relevant code snippets and integrate them into their work. It's the first encoder-only model specifically trained on large amounts of code data, making it especially valuable for software development.

Text Analysis and Understanding

Whether it's analyzing sentiment in social media posts or moderating content, ModernBERT processes text faster and more accurately than its predecessors. It excels at tasks like spam detection and identifying different types of information in text.

Smart Recommendations

From streaming services to social media, ModernBERT helps create more personalized recommendations by better understanding user preferences and content.

Challenges to Overcome

Despite its impressive capabilities, ModernBERT faces some important challenges:

- Like all AI models, it doesn't truly "understand" language the way humans do

- It can sometimes produce inappropriate content or reflect biases from its training data

- The model requires significant computing power to run effectively

- Its decision-making process isn't always easy to explain or interpret

While the development of ModernBERT represents an exciting step forward, but it's just the beginning. Researchers are working on:

- Improving the model's ability to work with multiple languages

- Enhancing its reasoning capabilities

- Making it more efficient and accessible

- Ensuring it operates ethically and fairly

As technology continues to advance, ModernBERT stands as a testament to the rapid progress in AI language understanding, while pointing the way toward even more impressive developments to come.