DEV Community: freiberg-roman

A Design Template for Your ML Project

freiberg-roman — Tue, 23 Jul 2024 19:12:35 +0000

The problem is self-evident. Researchers care more about original ideas and less about designing visually appealing graphics. If you pick a random paper from any top ML-conference, most likely you will be faced with figures that would fail any design class.

There are rare exceptions, and these stand out.

On first read of your article, paper, or blog, one typically scrolls quickly through all pictures and tries to build a first impression on whether it is worth spending more time on. If you confuse the reader upfront with a paint-like figure, you diminish your chances of receiving the reader’s full attention.

During my paper writing, I faced the same problem and was searching for some online templates. However, nothing seemed to match my demands. So, I hunted down some well-crafted papers and assembled a package that would have been the solution I was looking for.

The templates now include:

Most common symbols in machine learning architectures
Consistent color scheme for model components
Camera symbols for point cloud methods
Blender shaders and geometric nodes to create presentable point clouds out of any mesh

Have a look at some examples below and feel free to give it a try yourself. All Illustrator templates and Blender file can be downloaded here and used both commercially and non-commercially without any license.

Reach out to me anytime if you have questions or think that something is missing.

Enjoy!

The 1 Principle to Build Regular Equivariant CNNs

freiberg-roman — Wed, 17 Jul 2024 15:32:57 +0000

The one principle is simply stated as 'Let the kernel rotate' and we will focus in this article on how you can apply it in your architectures.

Equivariant architectures allow us to train models which are indifferent to certain group actions.

To understand what this exactly means, let us train this simple CNN model on the MNIST dataset (a dataset of handwritten digits from 0-9).

class SimpleCNN(nn.Module):

    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.cl1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1)
        self.max_1 = nn.MaxPool2d(kernel_size=2)
        self.cl2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, padding=1)
        self.max_2 = nn.MaxPool2d(kernel_size=2)
        self.cl3 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=7)
        self.dense = nn.Linear(in_features=16, out_features=10)

    def forward(self, x: torch.Tensor):
        x = nn.functional.silu(self.cl1(x))
        x = self.max_1(x)
        x = nn.functional.silu(self.cl2(x))
        x = self.max_2(x)
        x = nn.functional.silu(self.cl3(x))
        x = x.view(len(x), -1)
        logits = self.dense(x)
        return logits

Accuracy on test	Accuracy on 90-degree rotated test
97.3%	15.1%

Table 1: Test accuracy of the SimpleCNN model

As expected, we get over 95% accuracy on the testing dataset, but what if we rotate the image by 90 degrees? Without any countermeasures applied, the results drop to just slightly better than guessing. This model would be useless for general applications.

In contrast, let us train a similar equivariant architecture with the same number of parameters, where the group actions are exactly the 90-degree rotations.

Accuracy on test	Accuracy on 90-degree rotated test
96.5%	96.5%

Table 2: Test accuracy of the EqCNN model with the same amount of parameters as the SimpleCNN model

The accuracy remains the same, and we did not even opt for data augmentation.

These models become even more impressive with 3D data, but we will stick with this example to explore the core idea.

In case you want to test it out for yourself, you can access all code written in both PyTorch and JAX for free under Github-Repo, and training with Docker or Podman is possible with just two commands.

Have fun!

So What is Equivariance?

Equivariant architectures guarantee stability of features under certain group actions. Groups are simple structures where group elements can be combined, reversed, or do nothing.

You can look up the formal definition on Wikipedia if you are interested.

For our purposes, you can think of a group of 90-degree rotations acting on square images. We can rotate an image by 90, 180, 270, or 360 degrees. To reverse the action, we apply a 270, 180, 90, or 0-degree rotation respectively. It is straightforward to see that we can combine, reverse, or do nothing with the group denoted as $C_4$ . The image visualizes all actions on an image.

Figure 1: Rotated MNIST image by 90°, 180°, 270°, 360°, respectively

Now, given an input image $x$ , our CNN model classifier $fθf_\theta$ , and an arbitrary 90-degree rotation $g$ , the equivariant property can be expressed as
$f_\theta(\text{rotate } x \text{ by } g) = f_\theta(x)$

Generally speaking, we want our image-based model to have the same outputs when rotated.

As such, equivariant models promise us architectures with baked-in symmetries. In the following section, we will see how our principle can be applied to achieve this property.

How to Make Our CNN Equivariant

The problem is the following: When the image rotates, the features rotate too. But as already hinted, we could also compute the features for each rotation upfront by rotating the kernel.
We could actually rotate the kernel, but it is much easier to rotate the feature map itself, thus avoiding interference with PyTorch's autodifferentiation algorithm altogether.

So, in code, our CNN kernel

x = nn.functional.silu(self.cl1(x))

now acts on all four rotated images:

x_0 = x
x_90 = torch.rot90(x, k=1, dims=(2, 3))
x_180 = torch.rot90(x, k=2, dims=(2, 3))
x_270 = torch.rot90(x, k=3, dims=(2, 3))

x_0 = nn.functional.silu(self.cl1(x_0))
x_90 = nn.functional.silu(self.cl1(x_90))
x_180 = nn.functional.silu(self.cl1(x_180))
x_270 = nn.functional.silu(self.cl1(x_270))

Or more compactly written as a 3D convolution:

self.cl1 = nn.Conv3d(in_channels=1, out_channels=8, kernel_size=(1, 3, 3))
...
x = torch.stack([x_0, x_90, x_180, x_270], dim=-3)
x = nn.functional.silu(self.cl1(x))

The resulting equivariant model has just a few lines more compared to the version above:

class EqCNN(nn.Module):

    def __init__(self):
        super(EqCNN, self).__init__()
        self.cl1 = nn.Conv3d(in_channels=1, out_channels=8, kernel_size=(1, 3, 3))
        self.max_1 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.cl2 = nn.Conv3d(in_channels=8, out_channels=16, kernel_size=(1, 3, 3))
        self.max_2 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.cl3 = nn.Conv3d(in_channels=16, out_channels=16, kernel_size=(1, 5, 5))
        self.dense = nn.Linear(in_features=16, out_features=10)

    def forward(self, x: torch.Tensor):
        x_0 = x
        x_90 = torch.rot90(x, k=1, dims=(2, 3))
        x_180 = torch.rot90(x, k=2, dims=(2, 3))
        x_270 = torch.rot90(x, k=3, dims=(2, 3))

        x = torch.stack([x_0, x_90, x_180, x_270], dim=-3)
        x = nn.functional.silu(self.cl1(x))
        x = self.max_1(x)

        x = nn.functional.silu(self.cl2(x))
        x = self.max_2(x)

        x = nn.functional.silu(self.cl3(x))

        x = x.squeeze()
        x = torch.max(x, dim=-1).values
        logits = self.dense(x)
        return logits

But why is this equivariant to rotations?
First, observe that we get four copies of each feature map at each stage. At the end of the pipeline, we combine all of them with a max operation.

This is key, the max operation is indifferent to which place the rotated version of the feature ends up in.

To understand what is happening, let us plot the feature maps after the first convolution stage.

Figure 2: Feature maps for all four rotations

And now the same features after we rotate the input by 90 degrees.

Figure 3: Feature maps for all four rotations after the input image was rotated

I color-coded the corresponding maps. Each feature map is shifted by one. As the final max operator computes the same result for these shifted feature maps, we obtain the same results.

In my code, I did not rotate back after the final convolution, since my kernels condense the image to a one-dimensional array. If you want to expand on this example, you would need to account for this fact.

Accounting for group actions or "kernel rotations" plays a vital role in the design of more sophisticated architectures.

Is it a Free Lunch?

No, we pay in computational speed, inductive bias, and a more complex implementation.

The latter point is somewhat solved with libraries such as E3NN, where most of the heavy math is abstracted away. Nevertheless, one needs to account for a lot during architecture design.

One superficial weakness is the 4x computational cost for computing all rotated feature layers. However, modern hardware with mass parallelization can easily counteract this load. In contrast, training a simple CNN with data augmentation would easily exceed 10x in training time. This gets even worse for 3D rotations where data augmentation would require about 500x the training amount to compensate for all possible rotations.

Overall, equivariance model design is more often than not a price worth paying if one wants stable features.

What is Next?

Equivariant model designs have exploded in recent years, and in this article, we barely scratched the surface. In fact, we did not even exploit the full $C_4$ group yet. We could have used full 3D kernels. However, our model already achieves over 95% accuracy, so there is little reason to go further with this example.

Besides CNNs, researchers have successfully translated these principles to continuous groups, including $SO (2)$ (the group of all rotations in the plane) and $SE (3)$ (the group of all translations and rotations in 3D space).

In my experience, these models are absolutely mind-blowing and achieve performance, when trained from scratch, comparable to the performance of foundation models trained on multiple times larger datasets.

Let me know if you want me to write more on this topic.

Further References

In case you want a formal introduction to this topic, here is an excellent compilation of papers, covering the complete history of equivariance in Machine Learning.
AEN

How the React Pattern Applies Outside of JS: A Case Study

freiberg-roman — Sun, 16 Jun 2024 11:07:23 +0000

Looking from an oblique angle, React solves a peculiar problem: how to make a static markup language act dynamic. Judging by React's popularity, most of us would agree that React does a particularly good job solving this problem, and code using this ideology of component rendering often leads to self-contained reusable modules.

As an ex-web developer, I always prefer React over any other framework. However, the general design patterns have a much broader use case than solely web development.

The Problem I was Facing

For quite an extensive amount of time, I have been working with Deepmind's simulation software Multi-Joint dynamics with Contact, or MuJoCo for short. For our argument, it does not matter too much what exactly MuJoCo is. We will only focus on its scene definition. By design, simulated environments are compiled from XML files, which define the layout of the simulated scene and its objects.

Just to have an idea, you can find below a minimal example with two cubes, a sphere, and a ground plane.

<mujoco model="simple_scene">
    <worldbody>
        <!-- Ground plane -->
        <geom name="ground" type="plane" size="10 10 0.1"/>

        <!-- Red box with 6DOF joint -->
        <body name="box_one" pos="0 0 0.1">
            <freejoint />
            <geom name="box_geom" type="box" size="0.1 0.1 0.1"/>
        </body>

        <!-- Blue box with 6DOF joint -->
        <body name="box_two" pos="0 0 0.5">
            <freejoint />
            <geom name="box_geom" type="box" size="0.1 0.1 0.1"/>
        </body>

        <!-- Small sphere -->
        <body name="small_sphere" pos="0 0.5 0.1">
            <geom name="sphere_geom" type="sphere" size="0.05"/>
        </body>

        <!-- would be nice to add other objects -->
    </worldbody>
</mujoco>

This scene is particularly boring as not much will happen, but the parallel to HTML should be quite obvious. The problem arises when one wants something more dynamic with a variable number of objects. In my case, I had to dynamically exchange robot grippers and random objects throughout various scenes.

So, my goal was to recreate the basics of React for better handling.

It turned out much simpler than I previously imagined.

The Basics of React

At a high level, React provides a simple framework with a custom syntax called JSX to specify components. Components are self-contained units that maintain their own logic and markup. Components can also consist of other components, thus creating a hierarchy.

This creates a tree structure that mimics the actual Document Object Model (DOM) in the browser. To generate the final HTML, React uses a “render” method on the root component, which recursively calls the “render” method on all nested components.

Of course, modern React is much more complex than this, but this covers the basics and is all we need.

Mimicking React in Python

MuJoCo is written in C/C++ and provides native Python bindings for quick prototyping. As I did not plan to recreate JSX inside of Python, I simply opted for the multiline string using placeholders.

In the example MuJoCo scene above, this would translate to:

XML = """
<mujoco model="simple_scene">
    <worldbody>
        {objects}
    </worldbody>
</mujoco>
"""

Fortunately, Python comes with a string formatting convention, allowing us to use named placeholders ({objects} in this example) for future replacement using the format(object="string for replacement") method call on the template string.

This looks promising.

Similarly, we can create an XML template for the objects. Here is an example template string that allows for custom size settings for a box object:

XML = """
<body name="{name}" pos="0 0 0.1">
    <freejoint />
    <geom name="box_geom" type="box" size="{size}"/>
</body>
"""

Next, we need an interface for each component to make them share a method required for rendering each component:

from abc import ABC, abstractmethod

class MjXML(ABC):

    @abstractmethod
    def to_xml(self) -> str:
        pass

In the final XML assembly step, we will call on each component the to_xml method, which will output the resulting string for each component that will be used in the template.

Next, we implement the corresponding component for the example object above:

XML = """
<body name="{name}" pos="0 0 0.1">
    <freejoint />
    <geom name="box_geom" type="box" size="{size}"/>
</body>
"""

class SimpleBox(MjXML):
    def __init__(self, name: str, size: float):
        self.name = name
        self.size = size

    def to_xml(self):
        box_xml = XML.format(name=self.name, size=self.size)
        return box_xml

Finally, we implement the component for the scene:

XML = """
<mujoco model="simple_scene">
    <worldbody>
        {objects}
    </worldbody>
</mujoco>
"""

class SimpleScene(MjXML):
    def __init__(self, object_list: List[MjXML]):
        self.object_list = object_list

    def to_xml(self):
        obj_xml = " ".join([obj.to_xml() for obj in self.object_list])
        full_xml = XML.format(objects=obj_xml)
        return full_xml

Essentially, this is already a minimal workable example.

To actually simulate, we would need to construct the model and data of MuJoCo using our constructed XML string. As this article is not about MuJoCo, I have excluded the simulation part from this code sample.

Possible Extensions

While this example covers the basic concept, my specific use case required additional features such as dynamic loading of assets, randomizing names, and custom collision logic. Let me know, if you are interested in more complex physical simulations.

However, these additional features might overcomplicate the core message I want to convey: if you are working with static markup that needs dynamic behavior, consider using a component-based rendering mechanism similar to React. It can simplify development and enhance the flexibility of your applications.

How the Transformer Architecture Was Likely Discovered: A Step-by-Step Guide

freiberg-roman — Mon, 08 Apr 2024 16:04:06 +0000

As Transformers, and especially their attention mechanism, are a regular part of my work, I've often wondered how such an innovative architecture was created.

If you're looking for an implementation, I highly recommend checking out fast attention [https://github.com/Dao-AILab/flash-attention]. It's my go-to, and far better than anything we could whip up here using just PyTorch or TensorFlow.

In this article, I want to focus on the thought process that might have led to the attention mechanism. Here's what we'll be exploring:

Quick recap of the attention formulation
Core idea and a step-by-step construction of the weighting mechanism
Technical details
Some possible extensions

I hope this walkthrough sparks some ideas for your own modifications!

Original Formulation

Transformers are the go-to architectures for sequences. The paper "Attention is All You Need" from 2017 became a milestone in neural network architectures. It might seem like a genius stroke of discovery, but tracing the history of ideas can give us insight into how to arrive at such architectures ourselves.

This isn't a typical Transformer tutorial; I assume you have some familiarity with them. Instead, we'll follow the thought chain that likely led A. Vaswani et. al. to their idea, and see how the final attention mechanism emerges:

softmax(\frac{QK^T}{\sqrt{d_k}})V

Discoveries usually happen in the "adjacent possible" – the space of ideas made reachable by prior discoveries. Often, a single minor step creates a tipping point for impressive results.

The idea of attention in ML existed before the Transformer. For example, "Effective Approaches to Attention-based Neural Machine Translation" (published a few months earlier) used attention mechanisms to reweight hidden layers in sequence models. Additionally, Squeeze-and-Excitation networks, which reweighted selected CNN hidden layers, were also considered attention-based architectures.

The core idea was to boost features more relevant to the output. Mathematically, this re-weighting can be expressed as:

\sum_{i} w_i L_i

where $i$ indexes layers $L_i$ , and $w_i$ are the attention weights. It's also common to normalize the weights:

\sum_i w_i = 1

Application to Sequences

The Transformer authors were dealing with text embedded sequences $t_1,...,t_n$ (think of each token as a vector in $Rd\mathbb{R}^d$ , which represents some embedded information).

Before Transformers, machine learning models often had fixed input and output sizes. To process full sequences, we either fed tokens one by one i.e., $NN(t_i)=t_i'$ ) or used recursive architectures like RNNs or LSTMs, where a hidden representation tries to encode the prior sequence

NN(t_i​,NN(t_{i−1​},NN(t_{i−2}​,...)))=h_t​

These architectures were notoriously hard to train, and performance wasn't optimal (although similar ideas re-emerged recently in architectures like RetNet, which claims to be a potential Transformer successor).

The authors sought a way to create an expressive sequence-to-sequence architecture that could handle arbitrary-length input tokens. The goal was for each output token to have knowledge of every prior token in the input sequence:

NN(t_1, ..., t_n) = t'_1, ... ,t'_n

A straightforward idea is to reweight tokens, the essence of attention:

t_j' = \sum_i w_{ij} t_i

However, this limits us to linear combinations of previous tokens. Even with non-linearities like ReLU, we restrict ourselves to linear half-planes, creating a bottleneck.

A simple fix is to introduce a linear layer $V:Rd→RdV:\mathbb{R}^d \to \mathbb{R}^d$ for more flexibility:

t_j' = \sum_{i}w_{ij}V(t_i)

The key question now is how to learn the weights $w_{ij}$ . A naive approach — outputting them from a fixed layer $W:Rn∗d→Rn2W:\mathbb{R}^{n∗d} \to \mathbb{R}^{n^2}$ — fails for two reasons:

Scalability: The weight count of this layer is $d*n^3+biases$ , which doesn't scale well
Input Length Dependence: We'd reintroduce the dependence on $n$ , our original limitation in previous models

An alternative might be pairwise weighting, but this would require the layer to be invariant to input order, a difficult property to achieve with networks.

The solution (which might take some time to arrive at) is to split the weighting terms by indices:

w_{ij} = q_i k_j

These are the query and key terms. This solves our sequence length problem – now, our two layers $Q,K:Rd→RQ,K:\mathbb{R}^d \to \mathbb{R}$ only need to output a single value each.

Normalizing the weights with a softmax gives us a special case of the attention layer:

t_j' = \sum_i w_{ij} V(t_i)

where $w_{ij}'=Q(t_i)∗K(t_j)$ are normalized over index $i$ using the softmax function: $w_{ij}=softmax_i(w_{1j}',...,w_{nj}')$ .

Softmax is common in ML for making outputs resemble probability distributions.

Technical Details

We essentially learn three things for each token:

Value: Embeds the token to a new linear space
Query: Learns how relevant a token is
Key: Learns how relevant a token is for others

Mapping to a single dimension often creates performance bottlenecks. Therefore, making queries and keys multidimensional makes sense. The simple multiplication now becomes a dot product:

q_i,k_j \in \mathbb{R}^{d_k}, w_{ij}' = dot(q_i, k_j) = \sum_{l=1}^{d_k} q^{l}_{i} k^{l}_j

This extension provides an intuitive way to measure alignment between vectors. If two vectors align, their dot product is maximized compared to any other vector with the same norm. Essentially, the query now offers multiple directions for relevance and allows 'choosing' the contribution in each dimension. The key selects the direction in which a token wants to aggregate information for the next layer of tokens.

Now we can assemble the original formulation using queries, keys, and values written as matrices:

\in \mathbb{R}^{n * d_k} \text{ and } V \in \mathbb{R}^{n * d_v}

The dot product is calculated in one step as $QK^T$ .

You might notice the missing square root term. This was likely an experimental discovery. A useful rule of thumb is:

Neural networks perform best for inputs and outputs with a value range between -1 and 1.

The authors realized that if keys and values have a standard normal distribution, their sum has a variance of $d_k$ .

q_i, k_i \sim \mathcal{N}(0,1),\ \text{then } \sum_{i=1}^{d_k}q_ik_i \sim \mathcal{N}(0, d_k)

Normalizing by the square root of the variance returns the distribution to a standard normal one:

\frac{1}{\sqrt{d_k}}\sum_{i=1}^{d_k}q_ik_i \sim \mathcal{N}(0, 1)

Possible Extensions

While the dot product is a natural similarity measure, if your data has spatial information, you might consider rotary attention (as introduced in Act3: [https://arxiv.org/abs/2306.17817] ), which uses spatial distances within a 3D point cloud.

Even non-standard metrics like Manhattan or Mahalanobis distances could be used for attention relationships.

Hopefully, this inspires you to explore your own extensions! Remember, complexity should always have a reason.

Let me know in the comments if you'd like a deeper dive into these topics :)

How I Work Fully Remote through SSH without Sudo Rights

freiberg-roman — Sat, 09 Mar 2024 18:24:24 +0000

The pinnacle of my dream setup wishlist was a fully functional terminal workflow.

Throughout my academic journey, I dabbled with various Linux distributions, investing countless hours into configuring my tiling window manager and experimenting with suckless software. Although enjoyable, it wasn't particularly productive. Post-graduation, with the pressing need for efficiency, I reluctantly transitioned to an Ubuntu workstation — a tough but necessary adjustment.

Yet, the allure of a terminal-only setup remained.

I eventually discovered a method to tick all the boxes for a terminal-only setup, giving me compelling reasons to pursue this path. Primarily, it allowed for a fully remote development setup on the same system where my projects would be deployed. Whether it's executing torch code on a cluster entry node or deploying services in a Docker container, this approach proved invaluable.

The constraints imposed on the remote machines were strict:

No sudo rights
No option to install my preferred development software
No GUI
Restricted internet access behind a firewall

Moreover, I sought a development setup that could be ready in under an hour without any configuration requirements.

Even if you opt not to adopt this method fully, having such a system partially ready for on-the-fly production system edits could be beneficial.

Here's a rundown of my software stack:

junest
AstroNvim + LSP setup
nnn
Tmux
htop ...

Besides a stable internet connection, a proxy for mostly unrestricted internet access is essential. While not mandatory, Git is assumed to be a ubiquitous resource on most server Linux installations.

I'll detail each component below. Feel free to tailor anything to your needs.

Junest

This project was a game-changer, enabling my full terminal transition earlier than expected.

At its core, Junest is an Arch Linux subsystem that leverages the pacman package manager for installations, encompassing AUR repositories to meet the demands of Linux power users.

Explore the Junest Git project and follow the installation instructions. Activate the environment with junest setup and enable the running-from-host options.

That's all!

You can now install virtually anything with:

sudoj pacman -S <package>

For example, to install the software listed above:

sudoj pacman -S neovim nnn tmux htop

AstroNvim with LSP Setup

My Python development needs are straightforward:

Quick file navigation
Proper completion and search by symbols
Basic debugging
Minimal configuration with immediate functionality
Vim bindings

AstroNvim could be your entry point into a terminal editor that incorporates essential VSCode functionalities, especially if you're accustomed to Vim bindings.

Previously a PyCharm and VSCode user, I switched to Nvim after discovering AstroNvim, a pre-configured Nvim installation ready to use.

For Python debugging, follow the easy two-step installation process detailed on the AstroNvim website.

For a similar Python setup to mine, refer to my simple user configuration, which includes installation instructions.

Alternatively, Emacs is worth exploring for those not satisfied with Vim, though I haven't used it myself. For minimal editing, Nano may suffice.

NNN

This CLI file manager, compatible with Vim bindings, facilitates rapid directory browsing. Press ? within the program to access the manual.

NNN's learning curve is surprisingly gentle, even for non-Vim users.

Tmux

This widely used terminal multiplexer is invaluable, especially in server environments lacking a job scheduler. It ensures ssh sessions persist after disconnection.

For common workflows, consult this handy tmux cheat sheet.

To start and connect to a new session:

tmux new -s <mysession>
tmux at -t <mysession>

Disconnect without closing using Ctrl-b.

Final Thoughts

While there's much more to explore, these tools likely cover the majority of command-line software used in my daily development routine.

I still rely on GUI software, including web browsers. The primary motivation behind this development setup is to bridge the gap between my development and production environments, facilitating immediate debugging when necessary.

Compared to my idealistic setup during my studies, this system has significantly boosted my productivity, even under the constraints often imposed by corporate environments including custom hardware and non-existent Linux sudo rights.

Speed Up Your PyTorch Development Using Types

freiberg-roman — Sat, 20 Jan 2024 20:28:00 +0000

Types in Python are optional, but for any PyTorch project exceeding 1,000 lines of code, implementing them can be highly beneficial.

You've probably encountered a scenario where, after coding a model with several networks and initiating training, you realize there's a mismatch in data dimensions or type. Types help clarify the information flow. Moreover, they enable code completion tools to suggest more appropriate options.

While types won't eliminate debugging, they do make the process more focused.

Python Types

First, check if your editor or IDE supports static type checking. If not, a quick online search will reveal popular options like Mypy or Pyright.

Python types are akin to a simplified version of the typing systems in Java or C++.

Here's how the syntax looks:

variable: int = 10

For functions, specify the types of parameters and the return type:

def function(x: float, y: float) -> float:
    return x ** y

Below is a list of the most fundamental Python types you'll encounter:

int    # Whole number integers
float  # Floating point numbers
bool   # Boolean values (True or False)
str    # Strings
list   # Ordered mutable lists
tuple  # Ordered immutable lists
dict   # Dictionaries of key-value pairs

For the latter three, use the Python typing package for more precise type specifications:

from typing import Dict, List, Tuple

Example usages:

a_dict: Dict[str, float] = {'a': 1.0, ...}
a_list: List[int] = [1, 2, 3]
a_tuple: Tuple[int, bool, float] = (1, True, 3.141)

def forward(input: List[float]) -> Dict[str, float]:
    ...

For dealing with uncertain types, Any is a versatile type that accepts any value. Use it sparingly to maintain the effectiveness of your type checker.

The Self type represents the current class:

import numpy as np
from typing import Self, List, Tuple

@dataclass
class PointCloud:
    points: List[Tuple[float, float, float]]

    def from_numpy(arr: np.ndarray) -> Self:
        ...

Frequently recurring patterns in your codebase may warrant the creation of custom types from basic types:

from typing import Dict, List

SomeType = Dict[str, List[int]]
var: SomeType = ...

However, using expressive classes as types is often more sensible:

pcd: PointCloud = PointCloud([(1.0, 1.0, 1.0), ...])  # example from above

Be aware that your static type checker can often automatically identify the correct class upon assignment. However, use explicit types in scenarios where automatic deduction is not possible.

Types for PyTorch

A clean architecture combined with proper types elevates the professionalism of your PyTorch codebases significantly.

Key base classes in PyTorch you should be aware of:

torch.Tensor
torch.optim.Optimizer
torch.utils.data.Dataset
torch.nn.Module

For more details, refer to the official PyTorch documentation or explore the PyTorch codebase using your code editor's symbol search feature.

Dataclasses are extremely useful for defining types. They allow for quick class definitions intended for data holding and validation, usually without an __init__ method. Validation can be implemented using the __post_init__ method.

Example:

import torch
from dataclasses import dataclass

@dataclass
class PointCloud:
    points: torch.Tensor
    colors: torch.Tensor

    def __post_init__(self):
        assert self.points.shape == torch.Size([...])
        ...

Employing these types for network inputs and outputs significantly reduced my debugging time. I hope they will prove just as beneficial for you.

Common Troubleshooting for DataLoaders

If you adopt a strict dataclass typing strategy in PyTorch, you might face a unique challenge. The PyTorch DataLoader typically expects data to consist of Tensors or iterable collections of Tensors. To resolve this, reimplement the default collate_fn and provide it to the DataLoader. Here's an example:

@dataclass
class PointCloud:
    position: torch.Tensor
    color: torch.Tensor

def custom_collate_fn(batch: List[PointCloud]):
    positions = torch.stack([pcd.position for pcd in batch])
    colors = torch.stack([pcd.color for pcd in batch])
    return PointCloud(positions, colors)
...
loader = DataLoader(..., collate_fn=custom_collate_fn)

For distinguishing between batched and unbatched types, I usually rely on simple shape assertions.

With these minor adjustments, you'll find that your types work harmoniously with the PyTorch framework.

Let me know if you use types in your code and how they've benefited your projects.

Don't Know How to Structure Your Data Science Project? Try the Extended Clean Architecture Method

freiberg-roman — Sun, 14 Jan 2024 08:50:51 +0000

This project includes a PyTorch template. You just need to understand the basic concept behind it to start using it.

Data Researchers aim to solve problems and contribute meaningfully to science.

Nevertheless, most research projects - sometimes even those submitted to top conferences - lack any real structure. After the hacked-together prototype is done, the project often gets abandoned. Two months later, the original author might not even remember how to navigate their own codebase.

This is bad news when the project turns out to be useful and needs continuation.

The architectures commonly used in commerce do not meet the needs of modern research scientists.

Projects need to be submitted in an extremely short amount of time. Opting for Micro-services or Server-Client architectures does not make sense, as radical code changes might happen just three days before submission. Long-term planning becomes impossible. Therefore, guiding principles should be used, which guarantee the most maintainable code for minimal overhead.

Fortunately, creating a template structure that suits most of the everyday needs of a data scientist is simple to build.

One just needs to know a few principles and a bit of discipline. Once learned, placing a class in the right file becomes second nature. Furthermore, these principles are meant to save you time in the long run and shouldn't be seen as an additional overhead for achieving perfection.

This writing is highly opinionated. Use it as an idea marketplace on how it could be done, not as a definitive must.

Introduction to Clean Architecture

The concept is straightforward: in the lifespan of a project, let modules less likely to change not depend on those more prone to alteration.

Bob Martin's original Clean Architecture model delineates an application into four distinct layers. Dependencies are allowed in only one direction, towards the core. Consider the coupling of core business logic to a database as a prime example. Without a properly built adapter, you essentially vendor-lock your application by design. While the specific layers Martin discusses may not hold the same importance for data scientists, the overarching principles certainly do.

For a more in-depth understanding, I recommend checking out Bob Martin's original blog post.

For our purposes, we will adapt this model to include four layers, plus an additional core layer (which we will discuss later). Layer four forms the outer shell, encompassing data sources and data source adapters. Layer three involves adapter interfaces and gateways, while layers two and one focus on our core method. Layer zero concerns the language itself.

Below, you can see the original structure proposed by Bob Martin, which resembles what we need closely.

How do we achieve independence between logic and data source?

The answer lies in dependency inversion. This approach might take some getting used to, but it is quite straightforward once you get the hang of it. Essentially, you use an interface to break the dependence. Let us examine a minimal example of database dependency:

class BillingInformation:
    def __init__(self, db: VendorxyzDB):
        self.info = db.query("some SQL statement")
        ...

If BillingInformation is a part of your core logic, tying it directly to a specific database sets you up for trouble should that database become obsolete.

Now, let us apply dependency inversion to the same example:

class BillingInfoQuery(ABC):
    @abstractmethod
    def query_billing_info():
    ...

class VendorxyzDBAdapter(BillingInfoQuery, ...):
    def query_billing_info():
        self.db.query("some SQL statement")

class BillingInformation:
    def __init__(self, db: BillingInfoQuery):
        self.info = db.query_billing_info()

Yes, this results in more code. However, consider a scenario where the data source changes. The only component that needs reimplementing is the Adapter, leaving your core logic untouched. In this setup, BillingInformation and BillingInfoQuery would be in either layer one or two, while VendorxyzDBAdapter belongs in layer four, where our database resides.

Next, we will tackle a commonly overlooked issue in ML frameworks.

ML Frameworks and Dependencies

So ML frameworks are external dependencies, they are layer four dependencies, aren't they?

This is obviously impractical. Imagine you would have to call an adapter for every torch.Tensorcall and you would end up in an unmaintainable mess. Personally, I would argue that frameworks such as PyTorch, Tensorflow, NumPy etc. are part of the language definition. This means, their operations can be seen as a core part of the programming language such as Pythonor C++ and therefore can be used throughout the project without dependency inversion.

To be fair, a conventional software architect would not bother mentioning a specific language such as Java in a software architecture.

They reason beyond such concepts, but for a researcher, most programming language syntax is simply not effective enough. They require operations such as matrix multiplication or tensor products. Think of R or MatLab, where these concepts are already integrated by default. This also means, you could write your own language extensions of common operations under utility functions and use them throughout your project, without violating any principles.

It is obvious, but choosing a language extension means locking yourself into a framework.

But be careful with ML libraries.

Most frameworks like PyTorch provide also higher concepts like Dataloaders and complete built-in modules. Using them without architectural principles could result in the same mess we try to avoid in the first place. Think it through once and the work will pay off multiple times later on.

We will walk through common components you might encounter in your project.

Common Modules in an ML-Research Project

You can find a template for PyTorch here.
It showcases a small example that adheres to the principles we have discussed. I will be using the code from this template as a running example to demonstrate these principles in action.

Let us assume you are set to implement an experimental method and train a network eventually. Typically, you would need a data source (like AWS S3, an external server, a folder, or a Replay Buffer), several network modules, a training loop, helper functions, and a way to manage configuration. As we have discussed, our helper functions, such as data type converters between libraries and common operations in our utils, are integral to our language, placing them in layer zero.

Data Sources, a typical layer four component, require an Adapter for integration. PyTorch offers the Dataset class precisely for this purpose. Here’s an example:

@dataclass
class SourceCoordinates(CircleCoordinates):
    coords: torch.Tensor

class RandomCoordinatesDS(Dataset): 
    ...
    def __getitem__(self, _) -> SourceCoordinates:
        # Retrieves an item from the dataset
        # ...
        return SourceCoordinates(coords)

    @classmethod
    def collate_fn(cls, batch: List[SourceCoordinates]):
        # Function for batching items, required by torch
        stacked_coords = torch.stack([s.coords for s in batch])
        return SourceCoordinates(stacked_coords)

    ...

Here, CircleCoordinates is a dataclass coming from our method data interface. The collate_fn is a requirement in torch for batching custom elements. Here is the method outline for the full picture.

@dataclass
class CircleCoordinates:
    coords: torch.Tensor

class CircleMethod(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # Network definitions go here
        ...

        def forward(self, circle_data: CircleCoordinates) -> ...:
            # Processing 
            ...
            return out

CircleMethod encapsulates the project's method, including high-level interfaces for our pipeline, like computing bounding boxes for images, generating images, classifications, etc. The dependency inversion principle allows us to define our data layout without being constrained by the source. The dataset interface empowers us to load from any source and transform it into the format our network requires.

In general, your method should reside in layer two, and the networks and procedures used within this method in layer one.

Consequently, our training loop, which also serves as our command-line interface in layer four, now resembles a basic PyTorch tutorial script.

dataset = RandomCoordinatesDS(...)
dataloader = DataLoader(dataset, ...)
method = CircleMethod(...)
optimizer = AdamW(...)

for epoch in range(cfg.train.num_epochs):
    for data in dataloader:

        target = method.get_target(data)
        out = method(data)
        loss = method.compute_loss(out, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # Logging and model saving operations go here
        ...

This example is somewhat simplified, but it does not take much to extend this principle to any method. A more flexible toy example can be found in the PyTorch-Starter project.

Regarding configurations, they are a necessary evil. From a software architecture standpoint, configurations are part of the UI, residing in layer four. However, they often define hyper-parameters crucial to the core logic. Ideally, we would avoid this dependency, and thankfully, frameworks like hydra provide a solution with their hierarchical configuration structure, helping to decouple this dependency.

It mostly works seamlessly. However, if it does not, the need for quick prototyping might require some flexibility. It remains a topic where some artistic license could be warranted.

If you think your project can not be structured using these principles, I would love to hear from you. I am always on the lookout for counterexamples and thinking about how one could work with them.

Closing Thoughts

Sometimes, implementing radical changes quickly is necessary, and strictly following coding guidelines might only add additional overhead.

In such cases, it is advisable to focus on the implementation first and tidy up later. Everything presented in this discussion is intended as a toolkit to enhance efficiency in the long term. However, these tools should not overshadow the primary goal if the success of the project is at risk.

I invite you to try these methods and judge their effectiveness based on your own experiences. Adaptation and modification are key. Reflect on how these principles can be integrated into your current or upcoming projects, and make decisions that best suit your needs. Ultimately, the value of these guidelines is realized through their practical application and adaptability to your unique scenarios.

This article references specific tools like PyTorch and hydra primarily for illustrative purposes.

It is crucial to understand that the core principles of structured project management and clean architecture are not confined to these tools alone. They are broadly applicable across various technologies and frameworks in research and development. The fundamental goal is to promote a mindset of structured yet adaptable project development, applicable irrespective of the specific technologies employed.

Learn the Math of Diffusion Models (Enough to Reason about Ideas from Top-ML Conferences)

freiberg-roman — Sun, 07 Jan 2024 12:49:13 +0000

Diffusion models took off in recent machine learning conferences.

This year, NeurIPS had around 400 publications that employ diffusion techniques. At first glance the math might seem familiar, but some results might catch you off guard and require more clarification. This article tries to fill the gap without going the painful formal math way.

Here you will get the necessary basics covering the foundational building blocks.

Stochastic Processes
Brownian motion
Stochastic Integral
Ito Differential- Stochastic Differential Equations
Ornstein-Uhlembeck Process

A quick disclaimer.

The results presented here have been significantly condensed to keep things simple to not go into formal math territory. This comes at a significant cost in correctness. However, if you require a more formally accurate foundation, there will be no other way than literature. Make sure to check out the references below to have some starting points.

That being said, let us begin.

Stochastic Processes

In the domain of stochastic calculus one reasons about probabilities that a certain path will take.

For our sake, stochastic processes are noisy continuous paths in $Rd\mathbb{R}^d$ indexed by time. One could think of a function $Xt:R+→RdX_t: \mathbb{R}_+ \to \mathbb{R}^d$ . Since we want to keep track of a lot of them, we need a way to distinguish each by an additional variable $ω\omega$ - hence we denote $Xt(ω)X_t(\omega)$ - defining the total function space as $Ω\Omega$ . In Haskell-like notation it would result in $ω→(t→x)\omega \to (t \to x)$ . So we end up with a function outputting paths.

This allows us to assign probabilities to paths.

Brownian Motion

Often referred to as the most important object of discussion in probability theory.

In essence, Brownian Motions are continuous paths (Kolmogorov's continuity theorem) starting from zero $B_0 = 0$ and having an independent Gaussian increment

For\ s' < s < t\ B_t - B_s \text{ are independent from } B_{s'}

B_t - B_s \text{ follows a normal distribution}, \mathcal{N}(0, t-s = \sigma^2).

This means, at each time point, the tracked particle position

Bt(ω)B_t(\omega)

will completely forget its movement history and choose the position around the current one with a normal distribution. Extending this concept to

Rd\mathbb{R}^d

involves using multiple independent Brownian Motions in conjunction.

Following, we will use the Brownian Motion to define an integral, using the independent increments as our measurement delta.

Stochastic Integral (Itô)

In this chapter, we will explore the concept of the stochastic integral,

\int_{t_{start}}^{t_{end}}X_tdB_t

which is essential for defining the stochastic integral

The key idea is similar to the definition of the Riemann Integral through the limit of sums $limπ→0∑if(ti)∗Δt=∫f(x)dxlim_{\pi \to 0} \sum_i f(t_i) * \Delta_t = \int f(x)dx$ , which goes over the partitioned space $π=tstart=t0<t1<...<tN=tend\pi = t_{start} = t_0 < t_1 <...<t_N = t_{end}$ with each partition becoming finer. First we require the process to behave without hindsight (adapted process). This simply means that the process, for any given time point, does not look into future values. Here lies the crux of the stochastic integral definition,

lim_{\pi \to 0} \sum_{\pi}X_{t_i}(B_{t_{i+1}} - B_{t_i}) = \int_{t_{start}}^{t_{end}}X_tdB_t

where the integrand is required to be evaluated at the beginning of the interval to not peek into the future. Hence the resulting integral is an adapted process itself.

Note: There are other definitions of stochastic integrals like Stratonovich Integral, where evaluation happens in the middle of the interval, but also comes also with other properties. For diffusion models the Itô integral remains the goto choice.

As this material can be quite dense, we will go through some key results and properties of this integral.

We have learned previously that the increments, appearing now in the sum, are independent Gaussians. In some cases the distribution remains in the limiting cases. For this we require the integrand to be square integrable i.e. $∫∣∣f(s)∣∣2ds<∞\int ||f(s)||^2 ds < \infty$ then the process

I_t = \int_0^tf(t)dB_t

is Gaussian itself.

Some way to think about it, we integrate over paths, which spread unbiased in space, hence the integral is "centered" by design and sums up the "spread" of the integrand.

This neat property allows us to compute in some cases the integral in closed form.

As we will be concerned with the distribution of the integral, we will often encounter its square expectation to compute variances. There is a neat property.

E(\int X_t dB_t)^2 = E(\int X_t^2 dt)

For this property, the process needs to be in

M^2

. Think of it as the analogue of

L^2

space for stochastic processes.

The next theorem is less concerned with computation, but gives us some key insights of the nature of stochastic integrals

First, observe the bare-bone integral, $∫0tdBt=lim∑(Bti+1−Bti)=Bt\int_0^tdB_t = lim \sum (B_{t_{i+1}} - B_{t_i}) = B_t$ which is the Brownian motion. Now, consider plugging in a process that 'remaps' time, taking a path from zero to infinity. If this process is well-behaved to some extent i.e. in $M^2$ the result aligns with the Brownian motion. This statement holds if one defines,
$\tau(t) = min [u; \int_0^u X_s^2ds > t]$ for the integral,

\int_0^{\tau(t)}X_tdB_t

which is our recovered Brownian motion. There are a lot of other worth mentioning properties concerning martingales. However, these few hand picked examples should already convey the type of structure we are dealing with

We will continue now with the core topics.

Itô Differential

The previous topic was concerned about integration, this one will be with differentiation.

Differential equations are typically denoted by $dydt=g(y,t)\frac{dy}{dt}=g(y,t)$ . For stochastic differential equations (SDE) it is more convenient to look at the differential version $d y = g (y, t) d t$ as we can use the integral to define these. Thus the Ito stochastic differential is denoted as

X_{t_2} - X_{t_1} = \int_{t_1}^{t_2}F_tdt + \int_{t_1}^{t_2}G_tdB_t

for each

t_2 > t_1

, which emits (think of the limit) the differential

dX_t = F_tdt + G_tdB_t

We end up with a deterministic part and a stochastic.

$F_t$ is referred to as the drift term, it influences the overall direction of the process. One could think of the expectation of $dX_t$ . The term $G_t$ plays the role of a valve controlling the randomness of the process – how much the process oscillates – around the expected value. This is in line with our intuition on how the stochastic integral operates.

So how do we compute the derivative of a stochastic function?

This is where Itô's lemma gives us the solution.

df(X_t, t) = \frac{\partial f}{\partial t}dt + \frac{\partial f}{\partial x}dX_t + \frac{1}{2}\frac{\partial^2 f}{\partial x^2}G_t^2dt

For example, this result in

dBt2=BtdBt+12dtdB_t^2 = B_tdB_t + \frac{1}{2}dt

, which is contrary to the standard way derivatives work. An intuition on where the last part comes from lies in the Taylor series expansion.

df(X_{t})=f^{\prime}(X_{t})dX_{t}+\frac{1}{2}f^{\prime\prime}(X_{t})(dX_{t})^{2}+...

=f^{\prime}(X)(F_tdt + GdB_t) + f^{\prime\prime} (X)(F_tdt + G_tdB_t)^2 + ...

=f^{\prime}(X_{t})(F_{t}dt+G_{t}dB_{t})+\frac{1}{2}f^{\prime\prime}(X_{t})(F_{t}^{2}dt^{2}+2F_{t}G_{t}dt~dB_{t}+G_{t}^{2}(dB_{t})^{2})+...

The standard calculus derivative, is the best linear approximation as the other higher terms tent to zero in the limiting case. However, for the stochastic differential, the term

dB_t)^2

behaves as

d t

. A sort of hand wavy explanation why that is provides the following equation.

\mathbf{E}(B_{t_2} - B_{t_1})^2 = t_2 - t_1

As such, in the limit – in

L_2

– this lets

dB_t)^2

behave as

d t

However, some aspects of standard calculus do translate.

For example the partial differentiation,

\int_0^Tf(s)dB_s = f(T)B_T - \int_0^T B_s f(s)ds

which allows us to compute the stochastic integral for any deterministic function. This is a direct consequence of the Itô lemma applied to the function $f(t)B_t$ .

Note: We keep things simple here and define all terms only in one dimension. The extension to higher dimension can get quite hairy, but fortunately, in most research publications the dimensions are assumed to be independent. Hence, we can compute each dimension separate.

Stochastic Differential Equations

As we might now expect, Stochastic Differential Equations (SDEs) are defined as follows

dX_t = b(X_t, t)dt + \sigma(X_t, t)dB_t, t \in [u, T]

with an initial condition

X_u = x

. (could also be a random variable). The solution has to satisfy the following:

X_t = x + \int_{u}^Tb(X_t, t)dt + \int_u^T\sigma(X_t, t)dB_t

This forms the basic structure of SDEs.

Next, let us explore one of the most predominant forms of SDEs and its applications.

The Ornstein-Uhlenbeck Process

This process is arguably the most used example that can be encountered in every second publication on diffusion models.

dX_t = -\lambda X_tdt + \sigma dB_t,X_0 = x, where\ \lambda > 0

Some prior orientation. The negative sign pushes the process in the opposite direction. We should expect the process to settle in around zero. The constant before the Brownian motion lets us expect the overall stationary (for time going to infinity) to look like a Gaussian distribution.

As a first step we look at the deterministic part of the equation.

dX_t = -\lambda X_tdt

The solution here is of course

exp(−λt)xexp(-\lambda t)x

. Next we try to guess the solution by computing the differential of

exp(−λt)Ztexp(-\lambda t)Z_t

using Itô's lemma.

d(exp(-\lambda t)Z_t) = -\lambda exp(-\lambda t)Z_t dt + exp(-\lambda t)dZ_t

and compare it to the original SDE term by term.

exp(-\lambda t)dZ_t = \sigma dB_t

Finally solve for

Z_t

by integrating the differential.

Z_t = \int_{0}^{t}exp(\lambda s)\sigma dB_t

obtaining the solution

X_t = exp(-\lambda t)x + exp(-\lambda t)\int_0^{t}exp(\lambda s) \sigma dB_s

Recall the note before, that under certain conditions the stochastic integral results in a Gaussian distribution. In this case, conditions are met and we can even compute the distribution of the process at any time in close form.

Thus, our goal is to find the expectation and variance.

$E(Xt)\mathbf{E}(X_t)$ is simply the drive term $exp(−λt)xexp(-\lambda t)x$ and for the variance we recall the property on squared stochastic integrals and apply it to $E(Xt−exp(−λt)x)2E(X_t - exp(-\lambda t)x)^2$ .

E(exp(-\lambda t)\int_0^t exp(\lambda s)\sigma dB_s)^2 = exp(-\lambda t 2)\sigma^2\int_{0}^{t}(exp(\lambda s 2))ds = \frac{\sigma^2}{2 \lambda}(1 - exp(-2\lambda t))

Hence we end up with the distribution

N(exp(−λt)x,σ22λ(1−exp(−2λt)))\mathcal{N}(exp(-\lambda t)x, \frac{\sigma^2}{2\lambda}(1 - exp(-2\lambda t)))

This is basically it. Let us have a final discussion on the result.

As time progresses the distribution settles in the normal distribution $N(0,σ22λ)\mathcal{N}(0, \frac{\sigma^2}{2\lambda})$ for the infinite case. This is the reason why this process is of such use for diffusion models. First any data diffused will end up in the stationary distribution. As a byproduct the entropy of the distribution is kept in check and does not diverge to infinity. Second the distribution at any time point is known in closed form and can be step wise computed. This all allows us to learn the step wise diffusion updates, which diffusion models are all about.

Final notes

This article is getting a bit too long.

There are some topics left to discuss such as reverse processes and the SDE-ODE relationships. Let me know if there is interest in a continuation.

References

Øksendal, Bernt. “Stochastic Differential Equations.” In Stochastic Differential Equations, Springer, 2003.

Baldi, Paolo. Stochastic Calculus: An Introduction Through Theory and Exercises. Springer, 2010.

Song, Yang; Sohl-Dickstein, Jascha; Kingma, Diederik P.; Kumar, Abhishek; Ermon, Stefano; Poole, Ben. “Score-Based Generative Modeling Through Stochastic Differential Equations.” arXiv preprint arXiv:2011.13456, 2020.