Marcus Lewis

It's like a web view, but native

Sun, 10 May 2026 00:00:00 -0700

I think modern operating systems should support something I call the outerframe. An outerframe is like a web view, but it runs compiled machine code and uses the underlying operating system’s APIs to create UI.

A fundamental reason web views are useful is that they let the user experience be driven by something external, rather than confining the app to only use rigid built-in native app logic. The web view serves as a safeguard, preventing the external code from doing something dumb or nefarious.

The problem with the conventional web is its inefficiency, sacrificing the quality of the final product to make life easier for developers. For example, in my benchmark last year, web apps had a ~6x reduction in efficiency compared to their equivalent outerframe app. Up until recently, these costs have actually been fine, because high-end phones and laptops have been getting more powerful without a clear use case for all that power. But now there’s a renewed need for efficiency, both on high-end laptops and low-power devices. At the high end, there is a race to build useful local AI into the device, and this local AI can use every transistor and watt that we can free up. Meanwhile at the low end there’s a race to build a compelling AR / VR headset, which could conceivably move a lot of our work out of desktops and laptops and into the physical space of our rooms. It’s time for conventional software to get out of the way and fit into a smaller footprint, making room for local AI and new classes of devices. And, luckily, we simultaneously have a rise of AI code generation, which makes this practical.

With this blog post, I am open-sourcing an outerframe for macOS. The outerframe is a key part of Outer Loop, which uses it for SSH-based apps. The outerframe repo also contains “Outer Frame”, a simple web browser that you can build and launch from Xcode on a Mac. Try vibe-coding your own outerframe content, or try putting an outerframe in your own app.

Top: The first outerframe app

Because the first app to use the outerframe is an SSH client, the first outerframe app that I’m shipping is a modern top app. I call it Top (with a capital ‘T’).

Here’s a video walkthrough.

The Top backend can be run on any Linux or Mac device, and Outer Loop can install it for you with one click. See more here.

How it works

Broadly: The browser downloads a dynamically loaded library (a .dylib / .dll / .so file) and loads it into a sandboxed process. The sandboxed process renders the window contents. The browser and the sandboxed process send messages back and forth for various user interaction events and APIs.

I designed the outerframe for an ambitious possible future where it becomes part of the open web. The philosophy of the web would shift to being multi-platform, rather than cross-platform, while still using HTML where it makes sense. Webservers could choose to serve multiple implementations of their web apps, one per platform. Here I’ll describe how a browser and webserver negotiate which code to run. Then I’ll describe how the macOS-specific outerframe works.

The HTTP requests and the file formats

We need a cross-platform protocol that hands off to platform-specific behavior. Here’s the current spec for browsers and web views supporting the outerframe.

On navigation, the browser sends the additional HTTP header Outerframe-Accept: application/vnd.outerframe. (Ideally it would instead do this in the Accept header, but Apple’s web view doesn’t expose a way to modify Accept on navigation requests, so I chose to add a new Outerframe-Accept header.) If a server wants to respond with outerframe content, its response HTTP headers should include Content-Type: application/vnd.outerframe. When the browser sees this header, it parses the response body as:

byte 0..3     The ASCII string "OUTR", as a "magic" sanity-check
byte 4..7     UInt32 little-endian format version, currently 1
bytes 8..15   UInt64 little-endian offset to beginning of path
bytes 16..23  UInt64 little-endian length of path
bytes 24..31  UInt64 little-endian offset to beginning of data
bytes 32..39  UInt64 little-endian length of data
remaining     Variable length region, where path and data live

This is the outerframe’s analog of an “.html” file, and I call it the “.outer” file. Its contents are simply a format version, a path, and a data blob. Right up front, you see how this philosophy is different from the conventional web. This is not just plaintext, instead it’s an extremely-fast-to-parse binary format. You’ll generate this blob programmatically (e.g. using a simple Python script), or edit it with a hex editor. Making this format binary is planting a cultural flag: we put users, not developers, first. HTML already exists, so this platform gets to focus on the other extreme.

The browser parses the path and appends a platform string, for example this macOS implementation appends “/macos-arm” (or “/macos-x86” on Intel Macs). Other implementations might append strings like “/windows-x86” or “/linux-wayland-arm”. We then fetch that path to get the compiled code blob. The remaining data is an opaque data blob intended for the compiled code to interpret. (In the “Top” app, I don’t use this blob, but document-based outerframe apps use it heavily.)

This is where we move to the platform-specific part. The implementation gets to choose its own format for this compiled code blob. For this macOS implementation, I made it as Apple-native as possible; the “compiled code” is a “.bundle.aar” file, i.e. a compressed Apple Archive containing a loadable NSBundle. The outerframe code extracts this archive and a heavily sandboxed non-UI process “OuterframeContent” loads the bundle and asks it to start rendering.

For this particular “Top” app, this compressed compiled code download is 356 KB, so it’s fairly small compared to typical web apps. (It’s written in Swift. Anecdotally I’ve found with another app that the Objective-C version is ~100 KB smaller than the Swift version, but I haven’t studied this further.)

Rendering

The philosophy of the outerframe is to lean on the operating system’s UI APIs, but it’s important to note that not all UI frameworks are compatible with rendering in a sandboxed separate process. Neither SwiftUI nor AppKit’s NSViews can be used. Outerframe content still uses the underlying operating system, but it must use lower level primitives.

In my original implementation, the outerframe content code received a framebuffer (a.k.a. graphics context or IOSurface), and populated it. To use an analogy from the current web, this is similar to making the page a single <canvas/> and rendering everything with WebGL.

That approach is perfectly valid, but on macOS your potential is limited if you use a framebuffer. The Firefox blog discusses this. One fundamental issue is that you can’t tell the macOS compositor (WindowServer) about small updates, you can only invalidate the entire framebuffer.

A core UI component of Apple’s platforms is the CALayer, a higher-level abstraction that gives you access to framebuffers when you want them. CALayers form a tree/hierarchy. By embracing the CALayer, you solve the problem above, plus you get a bunch of animation functionality that runs completely in the operating system’s compositor, not even waking up your threads. For example, in the Top video above, when the process rows animate up and down, Outer Loop and OuterframeContent don’t need to do frame-by-frame updates; the threads simply sleep through the entire animation. And, of course, embracing this higher level primitive gives us a lot more functionality for free, which is convenient. There are a few kinds of CALayer, including a CAMetalLayer if you want to use Metal GPU shaders.

To use a remote-rendered CALayer, the outerframe uses the CALayerHost, a private API. I hope Apple someday makes a public API analogous to CALayerHost, as it is vital to having efficient sandboxed UI (Safari and Chrome also use it). For the reasons above, the public API IOSurface is less efficient and much less powerful.

So the macOS outerframe rendering model is: your binary registers a CALayer (which is typically the root of a tree of CALayers) and keeps it up-to-date. You can read more about the binary interface here.

Events, operating system UI

Of course, the rendering layer is just part of an app platform. Other parts include mouse events, keyboard input, interacting with the IME (e.g. the operating system’s overlay UI that converts Pinyin to Chinese characters), text editing hotkeys, accessibility, native blinking carets, Dark Mode, context menus, copy/paste, and so on. For each of these, the Outerframe mirrors the underlying operating system, passing events from the browser down to the content, and letting the content code send messages back. All of the above is implemented, but there are more things in this category that are not implemented yet. The current set of messages is documented here.

This is an area where Apple would be able to create a nicer API than I can, since they could modify the operating system.

Networking

The outerframe content code has access to macOS’s built-in networking APIs, but it is sandboxed to have no direct access to the network. Instead, it connects through a local proxy provided by the outerframe. That proxy enforces same-origin policies, so the content doesn’t have unfettered access to the network. It also gives the hosting app the ability to connect to other things. For example, Outer Loop uses this proxy to let the outerframe connect to servers over SSH, and to connect to local Unix socket files.

Vibe-coding already works

A challenge with launching new platforms nowadays is that you might have to wait for the large language models to pre-train on example code for that platform, which could take years. Conveniently, coding agents are already pretty good at writing for this platform, because the platform is macOS.

There are example projects for creating outerframe content in Swift or in C. Feel free to clone one of these, rename it, and ask your favorite coding agent to build something.

Not compatible with phones or headsets

There’s an elephant in the room. The outerframe is a coherent idea on macOS, Windows, and Linux. But for phone-like operating systems like iOS, Android, and headset operating systems like visionOS, app stores won’t allow apps that load external machine code. So we’re in an ironic situation where an eventual iPhone or Vision Pro version of Outer Loop will have to use a conventional web view to run “Top”, while the MacBook Pro will run the lightweight efficient native version.

I think it’s time to rethink these policies. Compiled machine code is not obviously dangerous, as long as it is sandboxed and only has “safe” APIs available to it.

Maybe the scariest thing about running untrusted machine code is side-channel attacks, like the classic Spectre and Meltdown, and the more recent proof-of-concept Apple M1 GoFetch attack. It’s questionable whether side-channel attacks are prevented by WASM / JavaScript JITs, but let’s suppose they are. Side-channel attacks broadly occur because, for a long time, it has been the job of chip designers to take mediocre code and make it run fast. As AI code generation gets better, I think chip designers can finally plan for a world where code is written with the hardware in mind, and it will actually become faster to use simpler chips with fewer performance hacks. (This is analogous to how GPUs are much simpler than CPUs, because they rely on the code to be better.) I’m betting that side-channel attacks can become a thing of the past, without sacrificing performance, but I confess that this is just a guess.

Apple could still choose to implement restrictions. They could require outerframe content on the iPhone or Vision Pro to be notarized, the same way they require non-App-Store macOS apps to be notarized. Then they would have the ability to instantly block everything from developers flagged as “bad”. Or they could just go all in, treating outerframe content the same way they treat JavaScript (which doesn’t need to be signed).

This isn’t about replacing the web, but letting it grow to fill new niches

I’m not saying that your favorite news website or blog should embrace machine code. The HTML web is good for that kind of thing.

The web has always been awkward as an app platform. It has succeeded anyway, because the browser is such a natural place for certain experiences, but I think not being a first-class app platform has prevented useful tools from arising. For example, with Outer Loop, I’m trying to use the web as the UI layer for remote servers and edge devices, and using actual native binaries makes the experience much better (like with Top, in the video above). With the HTML web, I don’t think I could get people to switch from command line apps to modern GUIs, but with native outerframe, I think I might.

I think some existing web experiences should move to native web apps, but the bigger growth area would be a set of new experiences/apps. The rise of AI brings two broad areas for this growth: the first is conventional software, now that so many more people can build stuff, and the second is dynamic generative UI, with AI creating experiences in real time. Combine all of this with the race to build a compelling AR / VR headset, in addition to the race to repurpose local compute for local LLMs, and I think machine code web content becomes the natural next step.

Maybe native apps will become web-like before the web becomes native

I like the idea of the open web embracing native web apps. It would change the focus of the web from being lowest-common-denominator cross-platform code to being heterogeneous and multi-platform. Each platform would then be free to innovate. Webservers could choose to serve N different native apps to N different platforms, with one of those platforms being the conventional HTML web. All of this would become especially practical because AI could translate code to other platforms and test it.

An alternate path is that a set of native apps will adopt a dynamic browser-like architecture, and the “native web” ends up existing in a collection of apps using outerframes, or something similar.

Either path sounds good to me.

FAQ

Why not WebAssembly? WebAssembly (WASM) is cool, but it’s still a lowest-common-denominator virtual machine, and it just improves the JavaScript part of “the web”. You still have to combine it with one of the web’s UI models (either treating an app as a “document” via the DOM, or using GPU shaders on WebGL / WebGPU). With the outerframe, you have full control over how your code renders different experiences to the screen. This path certainly has greater potential, and my hunch is that the payoff is big enough to be worth it. With the outerframe, you get to really take ownership of your actual assembly code. It’s a lot of fun. Broadly, I think a homogeneous platform will hold us back, and a heterogeneous world where each platform gets to innovate will lead to better things.

Is this just ActiveX all over again? It is similar to ActiveX, but with two important differences. First, the outerframe embraces heavy sandboxing, so the security concerns are greatly reduced compared to early ActiveX. Sandboxing also means we don’t need the “Do you want to allow this ActiveX control?” UI friction. In fact, ActiveX gradually approached an outerframe-like security model from 2007 - 2012, culminating in IE’s eventual “Enhanced Protected Mode” running ActiveX in a true sandbox, but by that point ActiveX was on its way out, because the web community had chosen HTML5. Second, with AI code generation, the world has changed around us. It is becoming practical to ship different native binaries to different platforms without a lock-in effect. Much of the web community hated ActiveX, even after its security issues were arguably solved, but I think they could come around to this. So I guess I am saying, “ActiveX, circa 2012, actually might be a good idea in 2026.”

(Thanks to Rosanne Liu for reading drafts of this post.)

Tip: Use services, not the terminal, to run local backends

Thu, 09 Apr 2026 00:00:00 -0700

I’ve embraced a fun way of working with Jupyter, Tensorboard, and other local webservers.

Example: Launch Jupyter Lab

In this video, I clicked a new “services” item in my macOS menu bar:

As you can see, the current backends on my laptop are Jupyter Lab, two websites, and a “Top” utility that I built. I only had to create these entries once, and now I have them forever. Now I’m a few clicks away from any of my backends, instead of needing to use a terminal to access them.

How?

I built a lightweight user interface that uses the operating system’s built-in “services” feature.

Most operating systems have built-in service management:

macOS has “launchd”. (You can see the services by running launchctl list.)
Most Linux distributions have “systemd”, which was inspired by launchd. (Run systemctl.)
Windows has offered services longer than either of those. (Open Task Manager.)

Services are heavily used by the operating system and apps, but most people don’t use them directly. Meanwhile, many scientific computing tools make you open a terminal and run a command to launch an app backend. Services are a perfect match for this use case; all that’s missing is good UI.

You could vibe-code your own version of this experience, or use Outer Loop. It’s a pretty lightweight set of features, since the operating system does the service management heavy-lifting.

To me, this flow just feels obviously better. Some overlapping reasons:

Requiring a terminal filters out a subset of people, and it breaks my creative flow. For many people who need use computers for data analysis, it takes a lot of cognitive overhead to find the terminal app, “cd” to the correct folder, and remember which command to run.
In terms of information input, this workflow takes on the same shape as projects. You type the command once into Outer Loop (or your own service creation app), then proceed to reuse it for the next few months, rather than having to repeat the command again and again.
With a terminal, I always seem to have Jupyter / Tensorboard / local webservers hiding in a sea of tabs. I’m finally done fishing through tabs.

Hope you like it! Download the Mac app here.

Footnotes

For Jupyter specifically, VS Code has done a good job with their file-viewer approach, where the backend and frontend are just one thing. But that’s pretty specific to local file browsing. One nice thing about preserving the separation of frontends and backends is that is scales up naturally to running apps like Jupyter on remote servers.

(Thanks to Rosanne Liu for giving feedback on earlier drafts of this post.)

Web apps over SSH can be surprisingly good

Fri, 10 Oct 2025 00:00:00 -0700

When you’re working directly with a server, web apps are useful. A few examples: Jupyter lets you run Python in a visual environment, Tensorboard lets you observe deep learning training runs, and tools like phpMyAdmin help you manage websites. The problem is, these apps typically aren’t directly accessible to web browsers, because exposing them to the web puts your server at risk.

The standard solution to this is SSH port forwarding. You create an SSH connection to the server, and then you tunnel arbitrary TCP connections over that encrypted SSH connection. This works, but it introduces friction and is often janky. Many entire “software-as-a-service” companies are devoted to providing smoother alternatives. For example, some companies give you a convenient way to log data to their servers, and they provide a safe web interface so that you never need to learn how to forward a port.

In this blog post I share two key improvements to SSH port forwarding that make it much more fun and reliable:

Solve scenario: “I closed my laptop, then reopened it, and the connection stopped working”
Incorporate SSH directly into the browser

To demonstrate this, I’m launching Outer Loop, a specialized web browser for Mac.

Change 1: Let me close my laptop

In terms of the information flow, port forwarding and SSH don’t really have the same shape. If we aren’t intentional, we will get a “round-peg-in-square-hole” situation.

When you use SSH more generally, the SSH session is stateful. It has a current working directory, it has running processes, it has environment variables. If you abandon an SSH session, you lose all of this.

When you use SSH for port forwarding, you aren’t relying on any of that. For you, an SSH session can just be a short-lived thing that may even get created and discarded multiple times while you interact with your web app.

So a first-class port forwarding library should treat SSH sessions as transient.

When you perform port forwarding, there doesn’t always need to be a live connection with the server. Rather, the only thing constant is that you have a process that listens to a local port on your system. This process (I’ll call it the “port forwarder”) is responsible for creating an SSH session with the remote server, but that session can be created opportunistically. The port forwarder may wait and create it after a browser connects to the port. Then, after the browser navigation, the port forwarder might disconnect the session immediately, or it might keep it around. The mindset shift here is: reusing an SSH session is not essential, it’s just a performance heuristic.

Of course, it is more efficient to reuse sessions when possible. In fact, it’s worth going further than that: there is a big performance win if we pre-allocate a set of SSH channels, each representing a speculatively-created fresh TCP connection between the remote SSH server and the remote web app backend. This lets us quickly forward HTTP requests when the browser sends them. (Remember: when a page loads, there is one HTTP request per page, image, and script.)

But these performance heuristics make fault tolerance more complicated. When there is some kind of network fault (e.g. “I close my laptop for 10 minutes”), we need to avoid sending browser HTTP requests into a channel of a closed session. Keeping sessions around, and in particular, keeping a pool of speculative channels, increases the chance that this will happen. If we send a request into a failed channel, we will quickly learn that the session has failed, but we will not know with 100% certainty that the HTTP request didn’t reach the web app backend. Perhaps the request reached the backend, and the fault occurred immediately afterward. So we can’t safely just re-send the request over a new session, because we risk doing a double-send, which in rare cases may be destructive. Thus, after sending an HTTP request into a closed session, we are forced to surface the broken connection to the user’s browser.

Fortunately, it is quite easy to solve this with simple pings, sending SSH_MSG_GLOBAL_REQUEST packets to the SSH server. The port forwarder can use a simple rule like, "If N minutes have passed since the last traffic, perform a ping and wait for the reply before using this session again." On my home internet connection to a GCP server, this ping introduces about 70 milliseconds of latency. If the session is found to be invalid, the server is still fast to respond, and we quickly set up a new SSH session and forward the TCP connection over that new session instead.

Independent of SSH, apps with long-running TCP connections will still have their connections broken by long laptop closes. Long-running TCP connections are inherently fragile, and any app that uses them is responsible for building in their own fault tolerance. The point here is to prevent the port-forwarding layer from introducing new faults or hangs, particularly in web app scenarios (which typically involve many very short TCP connections, a.k.a. HTTP requests).

Change 2: Create the obvious intuitive product

Today, using web apps over SSH requires you to open two windows. You have to open a new local terminal window and copy-paste a command like:
ssh -L 24601:localhost:8889 mrcslws@lambda4.mycompany.com

Then you go to your browser window, and type "localhost:24601" in the address bar.

What if, instead, the browser itself is SSH-aware?

This window connects to the server over SSH. Using this window, you can easily navigate to any private web apps via that server. You never have to think about the fact that you’re connecting to your own localhost. In this app, the term “localhost” actually refers to the server’s localhost, which means you can copy-paste URLs from an SSH terminal, for example when Jupyter or Tensorboard prints a launch URL to the terminal. You no longer have to choose your own arbitrary local port numbers; the app remembers the web apps you’ve connected to previously, and it carefully uses the same local port for that app over time, so the web app’s browser cookies are preserved. Port forwarding runs in multiple sandboxed processes, one for each server login. And you are free to use your own preferred web browser, using this app as a standalone port forwarder – just keep the app window open (or minimized).

Feel free to try it yourself.

Surprises

Going into this project, I assumed that connecting over SSH would sacrifice some performance, compared to connecting directly (and insecurely) via the web. I was surprised to see that, in practice, SSH is often faster. There are a couple of interesting discussion points here. First, with SSH, you typically already have an underlying TCP connection to the server, rather than having to create fresh ones, so that gives you a head start. Second, when we connect over SSH, we aren’t simply adding an extra SSH “layer of indirection”. We are replacing HTTPS (or rather, SSL) with SSH. We’re swapping out our authentication / encryption protocol; it’s a replacement, not an addition. If you’re doing big downloads or uploads, going over SSH will likely be slower, but I was surprised to see that in many scenarios it is actually faster.

There is a lot of low-hanging fruit in this space. We can make the experience of “working with a server” a lot more fun and accessible.

(Thanks to Rosanne Liu and Charlie Liu Lewis for reading and listening to drafts of this post.)

Smooth Pursuit, and how Screens Don't Mimic Reality

Mon, 04 Aug 2025 00:00:00 -0700

Your brain’s visual system is economical. Rather than capturing a big high-resolution image, only a small portion of your retina has light receptors packed together closely enough to capture fine detail. To construct a big useful picture of your surroundings, your brain relies heavily on eye movement.

One nice trick your eyes and brain use is smooth pursuit. You can lock your eyes on a moving object and move them continuously, not as a set of discrete changes, but by actually giving your eyes and head an angular velocity. Your visual system uses low-level details to perform quick error-corrections to stay locked on the object. This lets you give the most powerful part of your retina a stable input image for the moving object, allowing you to see fine detail even for fast-moving objects.

This is an elegant design. But this design breaks down if you’re looking at a screen.

Moving text is easy to read, unless it’s on a screen

Watch this animation.

Text movement speed (pixels/second): 1200

The quick brown fox jumps over the lazy dog

Press play and try to read the text.

(Turn your phone, it's better in landscape.)

Painful, right?

Now pick up a piece of paper, a book, or some other device. Move it in front of your screen at the same speed as this animation. You’ll see that it looks amazing compared to the animation on this screen. You’ll see that “1200 pixels/second” is not actually very fast.

The text isn’t moving, it’s teleporting

This may be obvious to some of you.

On a screen refreshing at 120 frames per second (fps), the text isn’t moving smoothly—it’s teleporting between positions every 8.3 milliseconds. Here are 12 frames of the animation:

Text movement speed (pixels/second): 1200

Simulated display refresh rate:

At 120 frames per second, the text shifts 10 logical pixels per frame.

What your moving eyes see

When your eyes track moving text on a screen, the actual image passing into your eyes is the text oscillating repeatedly.

Text movement speed (pixels/second): 1200

Simulated display refresh rate:

Actual time-varying input to your retina during "smooth pursuit" (slowed 8x)

The quick brown fox

Actual "static" input to retina during "smooth pursuit" (average pixel at each position)

This phenomenon is called display motion blur. Motion is much more pleasant in the real world than on screens.

Designing for this

In the age of 120 frames-per-second ProMotion™ Retina™ displays, you may have the impression that screens are now good at showing motion. Nope.

As we try to build “buttery-smooth” animations, we need to be aware that it’s a fool’s errand. Every motion animation that is much faster than 1 pixel per frame is going to look terrible if you actually pay attention to it. The upgrade from 60 fps to 120 fps doesn’t make motion animations good, it just makes them less bad. For designers, this is an important constraint; we should avoid creating situations where the user tracks a moving object. These animations can happen, but users shouldn’t be paying too much attention to the moving objects, or they will become unhappy. Until we invent screens that move with your eyes, motion is going to look ugly on screens.

Motion in movies

In these animations, your computer is generating individual discrete frames. In video recordings, the camera stores an aggregate blurred image for each frame. So, videos have pre-blurred images; they essentially render the blurry version of “The quick brown fox” above. Like software designers, filmmakers have to design around this; they need to avoid situations where viewers are tracking fast-moving objects.

Pre-blurring the images obviously doesn’t solve the problem of making the image look realistic and readable. It does solve a different problem: making the animation look smooth and blurry, rather than choppy. Game developers intentionally include “motion blur” in games, rendering multiple positions of an object into a single frame to simulate the film effect. This blurs the image, but it also makes it more obvious what motion occurred, which is important in fast-moving video games, where it’s more important to see where an object went than to show its fine detail in individual frames.

An interesting observation of this blog post is: in video games, even without motion blur, you’ll still perceive moving objects as blurry, if you’re viewing them on a screen.

Attaching a screen to your face

Screens are incompatible with smooth pursuit, and this has implications for VR / AR / “spatial computing”, since anything that gets rendered into the world will be blurry while you move.

What would happen if you wear an Apple Vision Pro while driving? Let’s assume it works perfectly, with a perfect resolution screen and perfect passthrough.

Each of these frames is a clear picture of a sign, and yet you perceive a blurry one, because the screen is not behaving like reality.

Actual input images to your retina when you track the sign

Actual time-varying input to your retina (slowed 8x)

Average pixel at each position

What you would see if you take the VR headset off

Sign distance:
Simulated display refresh rate:

This might mean that the soldier or police officer of the future will not use VR with video passthrough.

What about AR glasses? Passthrough VR can create amazing experiences, but it has the above problem. Glasses let the real world pass directly through, then add their own additional photons.

(Source: Meta)

As you’re working in the kitchen, these text overlays will be blurry. You’ll be forced to pause more often. AR devices have to render motion that otherwise isn’t there, because you are always moving. A screen in the kitchen can actually be useful, because it sits still. But with AR glasses, everything that “sits still” is actually constantly moving on your AR screen.

I wonder what the AR / VR people will come up with. Do their hardware labs have secret rotating screens that move with your eyes? Or will they treat this as a design constraint, and carefully design experiences around the limitations of screens?

(Thanks to Rosanne Liu and Charlie Liu Lewis for reading and listening to drafts of this post.)

The web could use machine code

Sun, 08 Jun 2025 00:00:00 -0700

Think of all the client-side code that runs on your devices. Most technical people would say that it falls into two categories:

Native apps, which are written for a specific platform and compiled to machine code.
The web, which is written in cross-platform interpreted code.

This mental model is a misconception. These categories are real, but nothing I mentioned about them is fundamental. The actual thing that distinguishes between these categories is whether the code can render a top-level window. Native apps can do this, while the web is confined to a “visually sandboxed” area. All those other details are orthogonal.

Web-like platforms will remain relevant for the foreseeable future, because we will always need a safe space to run other people’s code. Curated app stores are good for certain classes of code, like apps or games, but they aren’t appropriate for other use cases like scientific visualizations.

Today, a core feature of the web is that it lets you “write once, run everywhere”, but this could change. As AI code generation becomes cheaper and better, future web-like platforms could be based on actual machine code, compiled from any language. Each visualization would include compiled code for multiple device platforms. This would free us from the current “lowest common denominator” web, and instead would let visualizations use the full power of the device, which will be essential for making the leap to low-power devices like AR glasses.

Rather than waiting for the web to change, today we can create new analogs of the web view that support machine code. I call these outerframes. They could show up in lots of apps, new and old:

Four example apps that could use outerframes. The blue dotted line shows which part of the app could be an outerframe. The first two apps are Jupyter and Claude, which use the web. The third is ChatGPT, which is native and uses only static visualizations. The fourth is an actual app I've built using outerframes.

An outerframe is a safe space for running untrusted code within an app. As you would expect, the untrusted code is sandboxed in its own process, unable to access your data. Importantly, the untrusted code is also restricted from taking over the screen, and instead is “visually sandboxed” by the app, similar to how a website is visually sandboxed to a browser frame. Unlike in a conventional web browser, this code can include any compiled machine code, targeting both the CPU and the GPU, with full support for threading. Thus, the outerframe eliminates the “lowest common denominator” overhead of the web.

I built an outerframe for macOS. To demo it, I built a toy browser that, instead of using HTML and JavaScript, uses a cross-platform binary file format as its “page” and platform-specific machine code as its “script”.

This app is open-source, you can build and run it yourself.

In this post, I try to quantify the efficiency gains of moving to a machine code web on today’s devices.

Benchmark

In a previous post, I used a set of pragmatic model visualizations:

See old post for full version

When I wrote that post, I found that these visualizations were far too expensive to animate, at least when using a standard 2D canvas. Now, I use this task as my benchmark. I plot large distributions of machine learning model parameters at 120 frames per second. (120 hz is the refresh rate of modern phone and laptop screens.)

For the web, I tested using JavaScript and the browser’s GPU shader language, testing both WebGL and WebGPU. For the outerframe, I used a mix of Swift, C, and Metal Shading Language.

This benchmark plays to the web’s performance strengths in a couple of ways. First, I measure the visualization’s render loop, not its initial loading and display, where I think the outerframe will have a larger advantage. Second, this task is GPU-heavy, so a lot of the work is being done by compiled shader code, whereas in CPU-heavy tasks I think the outerframe would again have a larger advantage. For these reasons, I think it’s okay that I relied on JavaScript and didn’t test WebAssembly (Wasm), which would be more important for benchmarks of initialization or of CPU-heavy tasks.

I measure CPU utilization, GPU utilization, and macOS Activity Monitor’s “Energy Impact”. I use this Energy Impact as the measure of efficiency. (CPU utilization is a very noisy measure of efficiency, because “100% utilization” only sometimes means high power consumption. For example, when coordinating with the GPU, a CPU will often run a special “spin loop” that uses little energy but still uses the CPU at 100% for a non-negligible fraction of every frame.)

Result 1: Machine code can be more web-like than the web

The web is built on documents, and I first approached this visualization as a document. I implemented it using text with inline canvas elements. In the outerframe with machine code, this worked great.

In a conventional web browser, it was terribly inefficient.

WebGL

View in browser ↗

WebGPU

View in browser ↗

Platform	Frames per second	CPU	GPU	Energy Impact (relative)
Firefox (WebGL)	5	67%	10%	1550 (86x)
Chrome (WebGL)	61	101%	37%	1610 (90x)
Safari (WebGL)	120	76%	56%	1005 (56x)
Safari (WebGPU)	120	55%	16%	710 (39x)
Chrome (WebGPU)	120	55%	20%	435 (24x)
outerframe (macOS)	120	16%	14%	18 (baseline)

Rendering on the web consumed roughly 24x more power, at best, than rendering with machine code.

The issue here is that WebGL and WebGPU are not designed to be used as figures within a document. When you use these, you are supposed to use very few canvases, maybe 1-2. This was particularly true in WebGL, where using more than 16 WebGL canvases isn’t even possible, but it continues to be very expensive with WebGL’s successor, WebGPU. Meanwhile, in an outerframe I embrace the native platform, and macOS will happily let you draw to as many canvases (a.k.a. CALayers) as you want.

So my first result was: I learned that building this visualization for the web requires moving away from the natural approach, and instead using tricks like treating the canvas as a background behind the text, and carefully drawing to the correct locations. To be fair, when building native document apps, this can be a performant way of doing things, and I also see some efficiency wins when I adopt this approach in an outerframe. But it was striking how machine code and the native platform allowed me to build something in the way that felt natural and web-like, whereas the web did not.

Result 2: Machine code was about 6x more efficient

I changed to a single-canvas approach, and the web was finally able to handle it.

WebGL

View in browser ↗

WebGPU

View in browser ↗

Platform	Frames per second	CPU	GPU	Energy Impact (relative)
Safari (WebGL)	120	55%	18%	63 8x
Firefox (WebGL)	120	44%	15%	53 7x
Chrome (WebGL)	120	52%	22%	57 7x
Safari (WebGPU)	120	35%	13%	49 6x
Chrome (WebGPU)	120	51%	13%	47 6x
outerframe	120	13%	13%	8 (baseline)

With this approach, the outerframe was 6x more efficient than the the web. The outerframe version improved on its Result 1 numbers because rendering one canvas saves us from lots of copying, and macOS likes it when you draw to one CALayer, rather than drawing to 20 CALayers.

I suspect Results 1 and 2 are a representative story. The story is: there’s a set of scenarios that the web does a pretty good job of supporting, and sometimes you need to tweak your design to fit within that set of scenarios. But even then, it’s about 6x less efficient. (With error bars, of course. The actual number might be 3x, or it might be 10x.)

Result 3: The operating system’s native text view is not good at fast dynamic text updates

From that same previous blog post, here’s a variant of the visualization which plots parameters in one dimension:

See old post for full version

In that visualization, the min / max number labels are actual text in the page, not text rendered into the canvas. When I implemented this visualization in an outerframe, I used the operating system’s native text view and I batch-updated each subrange of the text in the text view, rather than rendering it myself. I found something disappointing: the performance was abysmal. In fact, even just this is terrible:

Platform	CPU utilization	GPU utilization	Energy Impact
outerframe	60%	0%	1550

Not surprisingly, the operating system’s text view was not designed for dynamically updating text many times per second. I’m sure this is one reason Safari’s WebKit has a different text stack.

So, I moved the text into the canvas, rendering it myself. With this change, the text view’s text is now static, with all changes occurring inside of the canvas. Now it’s much better:

Platform	CPU utilization	GPU utilization	Energy Impact
outerframe	21%	5%	17

Takeaway: If you’re going to create an outerframe, you probably shouldn’t base it on the operating system’s native text view.

How outerframes work

The point of the outerframe (and the point of the web) is to enable playful sharing. To support this, outerframes must perform “visual sandboxing”. Without that, a rogue actor could take over your whole screen, show you a pixel-perfect replica of any app, and trick you into doing anything.

There is no single correct way to implement an outerframe. There are two natural ways of building it:

Create a viewer for a data structure. Let the untrusted code modify that data structure, while your trusted viewer displays it.
Let the untrusted code write directly to a framebuffer / canvas / graphics context that the parent app controls.

The web uses both of these approaches, using the DOM and canvas elements, respectively.

This blog post’s demo outerframe is very web-like, in that it offers untrusted code analogous ways to display stuff to the user:

Create and update a rich text “attributed string” document, which the app displays.
Render to a “canvas” that is inline in the document, editing pixels directly.

Using this trusted document viewer approach, I got basic functionality like text selection for free. I hosted it in Apple’s native text view, using text attachments to host the canvases. You could imagine hosting it in another text stack, e.g. Zed’s GPUI, or a custom stack.

Since the native text stack is inefficient at dynamic text updates, I might simplify the interface and make the outerframe hold nothing but one big canvas. For outerframe content that wants to be document-like, it can implement the document affordances itself, via the canvas. This means implementing basic features like text selection, but it might not be a crazy idea, especially when you consider that it could become a reusable component, analogous to how JavaScript libraries are reused today.

Operating system vendors could choose to provide their own frameworks for hosting “visually sandboxed” UI in background processes. That would be great.

Classically, browsers allowed machine code plugins to write to canvas-like elements using a scheme called NPAPI. Browsers moved away from that, with the Chrome team describing it as a “90s-era architecture”. I like that. Now that we have LLM code generation, I think a few 90s-era architectures are worth trying again. This time around, from the user’s perspective it wouldn’t be a plugin model, but would be much lower friction, like JavaScript.

The whole point here is to lean in and embrace the machine, so I’m not trying to provide any standard interface. The outerframe is the opposite of a platform; I’m trying to remove layers, not add them. The platform is the platform.

How would an “outerdoc” work?

The outerframe solves the problem of running untrusted code on a particular platform. We still need a good shareable document file format. How can we write documents, presentations, blog posts, etc., that can be viewed from any device, while using the full power of each device?

I designed a format that enables this, but I’m still working on it. Two useful tricks:

Let documents provide preference-ordered implementations. This would be similar to how the HTML <video> element lets you provide preference-ordered <source> tags. Documents would provide a list of platform-implementation pairs, and they would typically include a HTML/JavaScript implementation as a fallback.
Provide “content types” like com.probablymarcus.ScalarDistributionList, letting viewer apps provide implementations for any content types that they want, either directly or via extensions. This would let apps provide first-party “trusted” implementations of certain content types, which would be useful if you want to bridge to the operating system’s native UI controls (e.g. NSTableView on macOS), and it could enable native apps to render a limited set of outerframe visualizations on platforms like iOS that don’t currently allow unvetted machine code.

Rebalancing trade-offs

With AI code generation, I think table stakes are changing. The “optimal way to build something” is always a function of development cost, and we’re experiencing a step change downward in development cost. The “write once, run everywhere” web made sense in the previous era, but now it seems extravagant.

Chip designers and chip manufacturers are pushing against the laws of physics to eke out every last bit of efficiency, and there’s a wave of new low-power devices that are within a couple years of being viable. A machine code web would give us an immediate ~6x software efficiency boost, and that’s on today’s hardware that has compromised its design to be good at JavaScript. Larger gains would come when hardware designers are freed from this constraint. I doubt we’re going to leave all of these gains on the table.

There is a good path forward for a machine code web. Today, anyone can design their own outerframe and use it in their apps. If you do this, you won’t be able to publish your app in most app stores, since they don’t allow untrusted machine code. But that’s okay! On those platforms, just fall back to a conventional web view. The experience won’t be as good, but that’s the platform’s choice, and they will have to compete with other platforms that do have outerframes. They’ll come around eventually.

(Thanks to Rosanne Liu, Jason Yosinski, Adam Zethraeus, and Mirko Klukas for reading drafts of this post.)

Extended material for 'Expressions are Pragmatic Model Visualizations'

Wed, 10 Jan 2024 01:00:00 -0800

This is the extended material for Expressions are Pragmatic Model Visualizations.

Followup on Example 1

Let’s further loosen the prior on the parameters.

Model Visualization

Gaussian Process with the following kernel.

Mean: constant

Covariance: Use distance between points as follows:

 * sum([
   
   # Kernel: Factorized scalar vs choice parameters
   * sum([
     
     # Scalar parameters
     * matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     
     # Choice parameters
     * exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))]),
   
   # Kernel: Joint scalar and choice parameters
   * prod([
     matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))])])

When comparing a point to itself, add noise value: (log scale)

Visualization 3: Model parameters now that the priors on the "lengthscales" have been loosened even further.

Compared to Visualization 2 in the original post, the parameters are more free. The discrete change in behavior as the dataset grows as been further reduced. Parameters are now reaching higher absolute values than before.

How were results impacted?

Chart 1: Hold-one-out cross-validation results, now adding a third configuration.

There is a significant benefit for very small datasets. However, for most dataset sizes, this leads to worse results. It seems that having priors is important, especially as the dataset gets larger.

Followup on Example 2

Here’s a plot of the training run, using batch training again. This time I overlay the mean of all of the models’ losses.

Chart 2: Each of the 60 models’ negative losses after each step of training.

We find something even more interesting. Yes, the vast majority of the models converge hundreds of steps before the final model converges. But we also see that individual models often get worse on single training steps; only the average of all scores is monotonically improving. This may seem normal if you are accustomed to optimization in neural networks, but BoTorch’s optimization (scipy.optimize with L-BFGS-B) is only ever supposed to take a step if it improves the model.

BoTorch trains multiple models in parallel by adding all of their loss functions to create a single scalar loss. This type of composability is a valid strategy with optimizers like SGD or Adam, since those follow the gradient wherever it leads. But other optimizers, even if they are gradient-based, decide whether to accept a parameter update depending on the loss at the destination point. For these optimizers, you can’t simply add loss functions, or you will get a change in the optimizer’s behavior. After adding losses, only the sum of those losses is guaranteed to improve monotonically, while updates can harm the individual losses. For some model types like neural networks, we might describe this as “regularization” and treat it as a good thing, but in those cases we should just use a different optimizer.

So, not only are we evaluating converged models unnecessarily; we are taking indirect paths to the optimum. When I count model evaluations, 18,817 total evaluations happen when training sequentially, while 92,820 happen when training in parallel, so we are doing approximately 5 times too many operations. In my experiments, on GPUs it is still worth training in batch, rather than sequential, but on CPUs it is better to do simply loop over models and optimize them independently.

The effect of all of this is that cross-validation in BoTorch is much slower than it needs to be. As you increase the size of the dataset, the following slowdowns occur:

More models are trained.
Each of those models is more expensive, since it has a larger dataset.
The optimization trajectory becomes longer and longer as more models are trained in parallel. (Unnecessary)
More pointless evaluations of converged models occur, due to the increased likelihood of a single random slow-to-converge model. (Unnecessary)

These four factors multiply to create a slow experience.

I am inclined to implement batch training differently, maybe by implementing a single training run and then using something like JAX’s vmap. This would eliminate the factor 3, and maybe it could be used in conjunction with JAX’s while_loop to also solve factor 4.

(That’s it for the appendix! Here’s a link back to the main post.)

Expressions are Pragmatic Model Visualizations

Wed, 10 Jan 2024 01:00:00 -0800

Most machine learning models are never visualized. Visualizing a model and its parameters often leads to immediate insights or bugfixes, but getting a good visual requires a lot of one-off work.

How do we get useful visualizations without requiring too much human overhead? I think code is underrated as a visualization. In this blog post I show a family of pragmatic visualizations that are each created by simply printing code expressions and rendering their parameters inline. Often, the printed code is not runnable, but is instead a visually optimized version of the model’s code.

I start with two real-life examples where these visualizations provided valuable insights, using these examples to demonstrate the visualizations. Then I turn to my main focus: eliminating the one-off work. I share rows2prose, a Python library that generates these visuals from a dataframe of model parameters and any styled text (not just code). This leaves only the problems of tracing model parameters to a dataframe and printing visually optimized model code. I show how scientific computing frameworks can help with this, in this case using a Lisp-macro-like approach which I demonstrate with Vexpr.

Example 1: The Model That Was Good, Then Bad, Then Good

Last year I ran a machine learning experiment and noticed that the baseline model from BoTorch had a weird curve.

Chart 1: Scaling curve for a Gaussian Process performing hold-one-out cross-validation. Given a dataset, the model is trained on all but one point in the dataset, then it predicts the output for the held-out point. I repeat the experiment on 50 different random subsets of a larger dataset and plot the 10th percentile, 90th percentile, and geometric mean.

The model is a Gaussian Process, but that detail really doesn’t matter for this blog post. This is a strange scaling curve for any machine learning model. We would expect a model to continually get better at prediction as it receives larger datasets, with diminishing returns only at the end. Instead, this one has a big lull in the middle, giving it an almost staircase shape. What could this mean? (Consider pausing here and guessing the reason. I had a guess, and my guess turned out to be wrong.)

An obvious next step is to look at the model itself, not just the model’s results. How do the final trained model parameters change as we scale up the size of the dataset? I visualize each individual parameter in diagrams like this:

The vertical axis separates different experiments—in this case, different dataset sizes—from top to bottom. The horizontal axis shows parameter values across different repetitions of that experiment. To visualize the entire model, I print a simplified version of the the model’s code, rendering each parameter in-place using one of these diagrams for each parameter.

Here’s what the model looked like for each point in Chart 1. Feel free to zoom in, but don’t get bogged down in the details, just observe that each parameter had interesting changes about halfway down, coinciding with the interesting changes from the experiment.

Model Visualization

Predict using a Gaussian Process with the following covariance and mean.

Covariance kernel: Use distance between points as follows:

 * sum([
   
   # Kernel: Factorized scalar vs choice parameters
   * sum([
     
     # Scalar parameters
     * matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     
     # Choice parameters
     * exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))]),
   
   # Kernel: Joint scalar and choice parameters
   * prod([
     matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))])])

When comparing a point to itself, add noise value: (log scale)

Mean: constant

Visualization 1: All of the parameters of a Gaussian Process model rendered in context. The covariance kernel at the top contains a few parameter types: a multiplicative positive scale (top-left), four multiplicative mixing weights (the other four parameters along the left side), while the rest are "lengthscale" parameters that are used to scale distances. Additionally there are noise and mean parameters (bottom). The noise parameter always ends up being very low, due to the default prior pushing it toward zero, and because all the points in the dataset are spaced apart by pseudorandom Sobol generation so the model is never forced to incorporate variance at a single location.

The sudden jump in accuracy in Chart 1 corresponds to a number of jumps in the parameters. Perhaps the most dramatic change was in the top half of the covariance kernel, where we see that a number of parameters stay fixed at about 0.3, until they suddenly jump up to large values.

This strongly suggests a theory: the model’s priors are too strong. With small dataset sizes, the gradients from better predicting the dataset are not powerful enough to overpower the gradient from the priors. After the dataset size crosses some threshold, the parameters are able to break free.

Let’s loosen the prior on the parameters and see if the issue is solved.

Chart 2: Hold-one-out cross-validation results, comparing the BoTorch baseline to a new configuration with a weaker lengthscale prior.

The model improved significantly. How do its parameters look?

Model Visualization

Predict using a Gaussian Process with the following covariance and mean.

Covariance kernel: Use distance between points as follows:

 * sum([
   
   # Kernel: Factorized scalar vs choice parameters
   * sum([
     
     # Scalar parameters
     * matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     
     # Choice parameters
     * exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))]),
   
   # Kernel: Joint scalar and choice parameters
   * prod([
     matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))])])

When comparing a point to itself, add noise value: (log scale)

Mean: constant

Visualization 2: Model parameters now that the priors on the "lengthscales" have been loosened. Specifically, I changed the prior on the lengthscales—the parameters along the right side—from its default value Gamma(3.0, 6.0) to Gamma(1.125, 0.375), a distribution with the same mode but higher variance, so during training there is a weaker gradient pushing each parameter toward the mode. (The model tunes its mixing weights to favor the top half of the kernel, so the parameters in the top half are the ones that increase, while those in the bottom are still pulled toward 0.3.)

The discrete change in parameters is now mostly gone. I ran further experiments with extra weak priors to make the discrete change disappear more, and I found that it worked, further improving results for small datasets, however it began harming results for large datasets.

So I’ve learned that with this model, I’ll get best results if I use weaker priors, at least for small-to-medium datasets. BoTorch / Ax’s built-in priors did not serve me well. That doesn’t mean the default priors are wrong, rather it suggests that users need to be willing to look closely at their machine learning models if they want to get good results. If using a machine learning model always gave users visualizations like these, I think many more people would use them well.

Example 2: The Pitfalls of Parallel Cross-Validation

I think people ought to always see their model, including while it trains. I built this experience for myself, and it quickly provided an insight. Here is a visualization I watched in realtime as my model above trained. This is a batch cross-validation task with 60 datapoints, so I am training 60 models in parallel and visualizing their parameters (hence, multiple dots per parameter). Click the button below to watch the models train.

Model Visualization

Predict using a Gaussian Process with the following covariance and mean.

Covariance kernel: Use distance between points as follows:

 * sum([
   
   # Kernel: Factorized scalar vs choice parameters
   * sum([
     
     # Scalar parameters
     * matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     
     # Choice parameters
     * exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))]),
   
   # Kernel: Joint scalar and choice parameters
   * prod([
     matern_25(
      norm_l2([
         compare('log_epochs') / ,
         compare('log_batch_size') / ,
         compare('log_conv1_weight_decay') / ,
         compare('log_conv2_weight_decay') / ,
         compare('log_conv3_weight_decay') / ,
         compare('log_dense1_weight_decay') / ,
         compare('log_dense2_weight_decay') / ,
         compare('log_1cycle_initial_lr_pct') / ,
         compare('log_1cycle_final_lr_pct') / ,
         compare('log_1cycle_pct_warmup') / ,
         compare('log_1cycle_max_lr') / ,
         compare('log_1cycle_momentum_max_damping_factor') / ,
         compare('log_1cycle_momentum_min_damping_factor_pct') / ,
         compare('log_1cycle_beta1_max_damping_factor') / ,
         compare('log_1cycle_beta1_min_damping_factor_pct') / ,
         compare('log_beta2_damping_factor') / ,
         compare('log_conv1_channels') / ,
         compare('log_conv2_channels') / ,
         compare('log_conv3_channels') / ,
         compare('log_dense1_units') / ])),
     exp(
      -norm_l1([
         compare('choice_nhot0') / ,
         compare('choice_nhot1') / ,
         compare('choice_nhot2') / ,
         compare('choice_nhot3') / ]))])])

When comparing a point to itself, add noise value: (log scale)

Mean: constant

Visualization 3: Cross-validation with dataset size 60. I used the looser lengthscale prior from above, and I used a different noise prior that doesn't endlessly push noise toward 0.

As I watched this in my notebook, I got a sudden impression: many of these points are converging much faster than others. In fact, I think there are hundreds of steps where 59 of the 60 models have converged, and we’re just waiting for the last one. This is particularly evident if you watch the parameters in the lower half of the kernel, where one single faint blue dot slowly approaches the cluster of overlapping dots. This is concerning because all 60 models are being evaluated on every step, even though the last few hundred steps are unnecessary for most of the models.

I tested this theory by comparing the two training approaches.

Chart 3: Training trajectory for 60 models trained in batch, compared to that of training each of them separately.

The problem is worse than I thought. Not only do some optimizations finish well before others, but every optimization takes many more steps when trained in batch. When I count model evaluations, 18,817 total evaluations happen when training models one at a time, while 92,820 happen when training in parallel, so we are doing approximately 5 times too many operations. I describe this in more depth in this post’s extended material. I am inclined to implement batch training differently, maybe by implementing a single training run and then using something like JAX’s vmap in conjunction with JAX’s while_loop.

I wouldn’t have noticed this problem if I hadn’t been able to see my model during training. Of course, other standard visualizations could have revealed this issue; Chart 3 is fairly standard, and it would have also done the job. But I didn’t have Chart 3, I didn’t know I should be building it, and bulding it is actually difficult and inefficient with BoTorch. I think a visualized expression is a useful jumping off point, and maybe we should try to always have it available to us.

A Recipe for Pragmatic Model Visualization

How do we give ourselves a playful environment where our models are always visualized by default?

I solve part of the problem with a new Python library called rows2prose.

You give rows2prose two things:

a dataframe containing scalars that should be visualized
a visually-useful string that describes the model (for example, the model’s code), including placeholder text for visualized values

You can use rows2prose in Jupyter notebooks, and it can output HTML visualization files from arbitrary Python scripts. It generated every visual in this blog post, and you can use it today.

The remaining gap has multiple solutions

How do you take your model and get a visually-useful summary of it? There are, of course, many ways to do this, including crazy new approaches like asking an LLM to generate one for you.

But here’s how I did it.

I created Vexpr, a Python library that takes inspiration from Lisp. In Vexpr, you build up expression data structures (“Vexprs”) similar to Lisp S-expressions. The expressions you see in these visualizations are simply printed Vexprs.

(JAX fans may be familiar with “Jaxprs”; Vexprs are similar, but they are more user-facing. A user of Vexpr is intentionally building up an elegant Vexpr, while a user of JAX doesn’t really care what their Jaxpr looks like. Vexprs and Jaxprs two solve different problems and I plan to use them together.)

Just like Lisp, Vexpr lets you use macros to modify these expressions. For my previous post, I used macros to vectorize expressions, and for this post, I used them to visually optimize the expressions. For example, the actual Vexpr program for my model uses an elementwise division between two arrays to divide N different distances by N different lengthscales, but I wanted to visualize this as N different divisions, with each parameter rendered next to its corresponding “compare” feature. I changed the code using a macro, thus I actually unvectorized the division operations to make them prettier. Here’s another example: my printed Vexpr was more verbose than I wanted to be, so I used macros to convert it into pseudocode! Each visualization above contains the function call compare('log_epochs'), which doesn’t actually exist. In the runnable expression, this compare term is replaced with a larger subexpression that extracts a "log_epochs" feature from two vectors (x1 and x2) and computes the distance. I wanted a succinct visual, so I “visually optimized” the expression, removing those details, and I never implemented compare. Code can be much more succinct when it doesn’t actually need to run.

In addition to macros, another useful idea that filled this gap was partial evaluation. I take a Vexpr, plug in its parameters, then evaluate all parts of the expression that are ready to be evaluated. This computes the “unvectorize” from the previous paragraph, taking arrays of divisors and indexing into them. This is also useful when machine learning models put constraints on parameters; often they implement constraints by storing “raw” versions of the parameters and passing them through an exp or sigmoid to move them into a constrained interval. Partial evaluation moves the values into the constrained interval so that they are ready to be visualized. One thing that made me laugh, putting the ideas of these two paragraphs together: even after converting my runnable Vexpr into pseudocode, I still ran partial evaluation on it, evaluating all expressions that could be evaluated. It feels funny when you tell your computer to “evaluate the parts of this code that are not pseudocode”.

Putting all of this together, here is the final architecture underlying these visualizations.

The “Macros + Partial Evaluation” functionality I used is all present in Vexpr, but the ideas are still baking. In a future post I might try to convince you to use them.

But what about Deep Learning?

This blog post featured human-comprehensible machine learning models like Gaussian Processes. In these models each parameter has a very clear meaning. Is this blog post applicable to Deep Learning?

First, let me appeal to you that comprehensible models are important, and I think people playing with Deep Learning ought to be among the most enthusiastic users of comprehensible models. Suppose I grant you the extreme position that a deep learning model is a black box that isn’t worth looking into. In that extreme, you have a great use case for comprehensible models: exploring the space of Deep Learning architectures and training regimes. You get to take the giant space of models and regimes, design your own hand-engineered features of that space like “learning rate” or “number of attention heads”, generate your own datasets of experiment results, and conduct symphonies of computers to explore the space. Deep Learning system design is what got me into these models in the first place.

Regardless, I think it’s possible to build useful, pragmatic visuals for Deep Networks. My main design goals would be: (1.) enable the user to detect when something in the network is broken / not being used, and (2.) put an expression in front of the user to encourage playful tweaking of the architecture.

Conclusion

These visualizations were immediately useful, and they are pragmatic because they are not specific to any model type. If you can extract a text description of your model, and if you can trace a set of useful-to-visualize scalars, then you can visualize your model.

I think that somehow we should give all users of machine learning models access to visuals like these. Using some combination of our shared frameworks, our example code, and our crazy LLM tools, we should take on the responsibility of not only performing our desired computation, but also rendering a useful expression of it.

(This post has an appendix. This project is supported by a GCP cloud compute grant from ML Collective, which has been super helpful. Thanks, also, to Rosanne Liu for useful feedback on drafts of this post.)

What happens when you vectorize wide PyTorch expressions?

Thu, 19 Oct 2023 07:00:00 -0700

In scientific computing, code is often naturally expressed as wide, tree-like expressions. Often different branches of that tree contain similar chunks of logic, so there is potential to run many different branches together in parallel vectorized operations. What happens when you take your nice tree-like code and mangle it into hard-to-read vectorized code? How would a person do that?

I created Vexpr and used it to take real experiment code and convert its readable expressions into vectorized expressions at runtime. In this post, I present the results. Topics include:

What is the immediate impact?
How does this relate to torch.compile?
What is the more detailed breakdown of the impact on the GPU and CPU?

Introduction: Wide expressions

Mathematical expressions naturally form trees. Here’s a toy example.

\[\sqrt{a^2 + b^2} + \sqrt{c^2 + d^2}\]

Here’s Python code implementing this expression and highlighting its wide tree-like structure:

sum([math.sqrt(sum([a ** 2,
                    b ** 2])),
     math.sqrt(sum([c ** 2,
                    d ** 2]))])

For this simple function, we can write vectorized PyTorch code by hand:

torch.tensor([a, b, c, d]).pow(2).view((2, 2)).sum(dim=0).sqrt().sum()

The former code calls pow 4 times, calls sum twice, calls sqrt twice, then calls sum one last time. The latter code flattens the tree into a single pipeline so that one operation occurs for each level of the tree. Previously we ran one Python function call per node of the tree, and afterward we run one call per level of the tree – a number that is exponentially smaller.

PyTorch is often used in a pipelined way as shown above, but that’s usually because the expression is inherently deep; neural networks are the obvious example. Here we are concerned with expressions that are wide.

Imagine scaling up this toy example: use actual, larger expressions; let each variable represent a list of vectors, not just a number; within the expression, call functions like SciPy’s cdist which compute pairwise distances and return large matrices; do all of this on batches of inputs. With these changes, we now have a wide expression that is worth running on a GPU.

This scenario comes up often in scientific computing. For example, the kernel of a Gaussian Process (GP) takes in two lists of vectors and returns pairwise similarities. We can encode some of our intuition into these kernels by composing them via weighted sums and products. This leads to giant tree-like expressions, and because those expressions have many similar operations in each parallel branch, there is a lot of potential for vectorization.

Writing vectorized code for giant expressions is hard, so I created Vexpr to make it easy. Vexpr takes readable-but-slow PyTorch, NumPy, and JAX expressions and compiles them into fast, delightfully ugly vectorized expressions.

How to vectorize one level of an expression tree

Wide expressions naturally form a tree. When we vectorize that tree, we create a new narrower expression that invokes a series of Python functions, one for each level of the original tree.

How do we collapse a tree and reduce operations to be one per level? Each of the operations now will have to take in a tensor that contains every input to that level and return a tensor that contains all of the outputs.

We can see some operations already support such a change.

Elementwise operations like pow and sqrt can run on multiple inputs as-is.
Parallel sums can be implemented as .view(...).sum(dim=0)

What about more difficult cases? For example, what if the parallel sums are of different lengths? On GPUs, fast parallel reductions only work when inputs all have the same length. Would providing a single-input single-output interface still have a benefit? The answer is yes, even if the operation internally just loops over the sums. Collapsing one level of the tree allows subsequent levels to be collapsed, and low-hanging fruit tends to appear in the deeper levels. Moreover, we can do better than looping over the sums, even when the lengths are different. Vexpr’s vectorizer groups the inputs by length and performs a reduced number of operations—one for each unique length. For example, this decreases the number of torch.cdist operations in my hands-on kernel from 49 to 5. These 5 calls happen invisibly inside of a cdist_multi function that the code calls once.

These collapsed operations often introduce some overhead. Batch operations often require permuting tensors (e.g. for batched pairwise distances) or reordering values (e.g. putting equal-length sums next to each other). For operations that are grouped by size, we must split input tensors then re-concatenate the output tensors. Thus, while vectorizing has many benefits, it also introduces extra work for the GPU that is not necessary when using a non-vectorized expression. Is this overhead worth it? Let’s turn to experiments to find out.

Vectorizing leads to 4x-7x speed-up on one set of benchmarks

Here I run a pair of benchmarks on an NVIDIA V100 GPU. I describe the benchmarks within the context of real Gaussian Process use cases, but you don’t need to understand Gaussian Processes to understand these results; I am simply running this big vectorized expression on inputs of different shapes.

When you use Gaussian Processes for Bayesian Optimization, you first fit a GP’s parameters to the training data, then you optimize a set of candidate points to maximize expected improvement. Both of these steps include backpropagating gradients through the GP kernel.

Benchmark 1: Fit the kernel’s parameters. I run the forward and backward pass of the kernel, using x1.shape == x2.shape == (379, 26), i.e. a training set of 379 26-dimensional vectors. I repeat 100 times, then wait for a final CUDA synchronize. I test with and without torch.compile.

	Don’t compile	Compile
Baseline	6.43s	3.0s	Compile speed-up: 2.1x
Vectorized	0.99s	0.76s	Compile speed-up: 1.3x
	Vectorize speed-up: 6.5x	Vectorize speed-up: 3.9x	Combined speed-up: 8.5x

Benchmark 2: The optimization loop. In this experiment, x1.shape == (60, 2, 26) and x2.shape == (60, 381, 26). That means we’re searching for 60 single candidate points, with an additional point included in each (hence the 2, not 1) because the Noisy Expected Improvement algorithm also gets predictions for a set of previous potentially best points. These two points are appended to 379 training points to give us 381 points. Again I run the kernel’s forward and backward pass 100 times.

	Don’t compile	Compile
Baseline	6.03s	2.57s	Compile speed-up: 2.3x
Vectorized	0.83s	0.63s	Compile speed-up: 1.3x
	Vectorize speed-up: 7.2x	Vectorize speed-up: 4.1x	Combined speed-up: 9.6x

In these experiments, vectorizing my kernel caused a 4x speed-up when I used torch.compile, and a ~7x speed-up with straight pytorch code.

…but this speed-up doesn’t happen in all benchmarks

Now I test the kernel’s performance in another scenario: hold-one-out cross-validation. During cross-validation, we fit \(N\) models on \(N\) slightly different datasets. We can test this by just rerunning Benchmark 1 on \(N\) models in parallel. To demonstrate a surprising phenomenon, I set \(N=20\).

Benchmark 3: Cross-validation. Repeat Benchmark 1, but train 20 models in parallel rather than 1. So we have x1.shape == x2.shape == (20, 379, 26), and the shapes of every parameter, e.g. lengthscale, have (20,) prepended to them.

	Don’t compile	Compile
Baseline	15.6s	7.50s	Compile speed-up: 2.1x
Vectorized	17.7s	6.15s	Compile speed-up: 2.9x
	Vectorize slow-down: 1.1x	Vectorize speed-up: 1.2x	Combined speed-up: 2.5x

This result surprised me. The vectorized version is sometimes slower than the baseline, at least when you examine only the GP kernel in an isolated benchmark. The end-to-end performance of the larger system is still often faster due to the freed up CPU (see next section), but it is noteworthy that vectorizing code doesn’t always directly speed up that code.

To understand where the change occurs, I rerun this benchmark for different \(N\). Note that \(N=1\) and \(N=20\) correspond to the original Benchmark 1 and Benchmark 3 results, respectively. I used the same hardware, but I enabled some additional profiling, hence the slower times compared to above.

The vectorized kernel initially has a huge advantage over the baseline, but this advantage diminishes as we give it more parallel work, and without torch.compile the vectorized kernel is eventually slower when viewed in isolation. Why does this happen?

The CPU is a bottleneck. Vectorizing removes that bottleneck.

To really understand the performance impact of vectorization, we need to understand the CPU and GPU usage before and after. First, note that CUDA / PyTorch are built on an asynchronous relationship between the CPU and GPU, where the CPU ideally should always run ahead, always queueing the GPU’s future work, while the GPU is always working through its queue. However, there are two events that cause the CPU to wait for the GPU:

When the CPU chooses to read a value from the GPU. This is a decision made by your code.
When the CPU gets too far ahead of the GPU. This is a decision made by CUDA.

I profiled the kernel using NVIDIA’s nsys, and I used that trace to obtain GPU and CPU active time. My so-called CPU “active” time is actually an inferred value; PyTorch uses CUDA in a way that spins the CPU 100% constantly, even when the CPU is just waiting for the GPU, so I use heuristics to detect these waits and subtract it them from the actual active time. (2023-10-31 update: Thanks to gregjm for pointing out that CUDA offers a way to avoid this CPU-spinning, and that this is actually a PyTorch issue. The initial version of this blog post blamed CUDA for this unnecessary power consumption.)

Here is the same plot from above, but with the CPU and GPU time overlaid.

First, to understand the result of Benchmark 1 above, look at the left side of both plots. Both the baseline and vectorized models have low total active GPU time. But the baseline model puts a much larger workload on the CPU, which is responsible for orchestrating the set of operations that are sent to the GPU. Thus, the GPU spends the vast majority of the time idle, waiting for the CPU to give it more work. For the vectorized model the amount of CPU work is almost always less than the amount of GPU work, so the GPU is almost never idle; the benchmark time is roughly equal to the GPU active time.

Now we focus on the surprising result from Benchmark 3, where the baseline model did slightly better than the vectorized model. As we scale up the number of models, we see an interesting phenomenon. When training many models simultaneously, even the baseline model is able to keep the GPU busy. Once the GPU active time exceeds CPU active time. one of the key selling points of vectorized code is eliminated, because the nonvectorized code becomes good enough. This was a fun, surprising fact; even unvectorized code can outrun the GPU if you pass in large enough tensors. I expect this phenomenon to occur in other large-batch scenarios like training on very large datasets or doing Bayesian Optimizations with very large sets of candidate points. Of course, the vectorized code is still superior when using torch.compile, and in all cases its CPU usage is far superior.

Finally, let’s look at the GPU workload. Independent of the effect on CPU, what is the impact of vectorization on the total amount of work that the GPU has to do? Looking at slopes of the “GPU time” lines, we see that for some models, e.g. my non-compiled model, vectorization increases the total workload, and for other models it decreases the workload. I studied the CUDA traces closely and found that vectorization does indeed reduce many aspects of the GPU workload, greatly reducing the number of operations and decreasing the total amount of time spent on the fundamental computations of the algorithm. However it also introduces overhead (mentioned above) by interspersing operations that permute and reorder the tensors, or splitting them into groups then concatenating results. Sometimes the reduced “fundamental” time outweighs the additional overhead, while other times the overhead outweighs the reduction in fundamental time.

So we see that vectorization has three effects:

It lets us keep the GPU busy even when inputs are small.
It frees up the CPU to do other work.
It can slightly change the total amount of GPU work, sometimes for the worse.

The speed-up from vectorization can be great, but it can be underwhelming in scenarios where none of the benefits are needed.

The benefits of vectorization increase as GPU speed increases

Here is a point that follows naturally from everything above, but it might not be immediately obvious. It is quite striking when you experience it.

As I built Vexpr, I tested on an NVIDIA T4. Then for this experiment, I upgraded to a much-faster NVIDIA V100, and the benefits of vectorization greatly improved.

You always want your CPU to stay ahead of the GPU so that the GPU is never idle. Code that is good enough to stay ahead of this year’s GPU might not be good enough for next year’s GPU. With each upgrade, you need the dotted CPU line from these charts to be lower and lower.

This means that vectorizing your code is good strategy for future-proofing it.

GPyTorch’s “structure” kernels show similar results

I also tested GPyTorch’s limited support for vectorization. GPyTorch lets you take sets of identically-shaped kernels and run them as a single vectorized kernel. This capability is easy to use when summing single-feature kernels, so I created a partially vectorized kernel by replacing sums of single-feature Matern kernels with single Additive Structure kernels.

(Don’t focus too much on comparing absolute speed of the Vexpr and the GPyTorch kernels. There are many small differences between the two that have nothing to do with vectorization. For example, unlike my Vexpr kernel, GPyTorch has a nice optimization of putting this line before this line.)

Vexpr has been optimized much more for this use case than GPyTorch, but we see that the same fundamental phenomena occur. Vectorization leads to wins at small batch sizes, but the advantage diminishes at large batch sizes. In this experiment, vectorization increased the total GPU workload, so with large batch sizes the only advantage of vectorization is a freed up CPU.

Interestingly, neither GPyTorch kernel ever reaches a point where GPU time is equal to the benchmark time. The GPU always spends at least 1-2 seconds idle. This happens because GPyTorch’s kernels do a synchronous equality check on every call, forcing a GPU synchronize, which causes the CPU to fall behind the GPU immediately afterward. So when you use GPyTorch kernels on GPUs, you don’t get the ideal fully asynchronous execution that you’re supposed to get with CUDA / PyTorch. (There is good news: PyTorch has recently introduced torch.cuda.set_debug_mode(2) which detects these unwanted synchronize events. Like others, I think every PyTorch library developer should become friends with this API.)

Closing thoughts on compilation

Everybody agrees that vectorization is good. Is it worth doing even in cases where it requires extra work, like with wide expressions? The results above suggest: most of the time, yes, but with interesting nuances. Vectorization is especially great for making the most of a GPU when your task is running many iterations on smaller batches of input data. For larger-batch scenarios, maybe you’ll only benefit from vectorization after you add JIT compilation, or after you upgrade your GPU, or after you’ve found some use for all the newfound idle CPU time. But, to a first approximation, vectorization is good.

The real open question for the field remains: for “wide” computation graphs, what is a practical strategy for getting vectorized code?

My experiments above show that this is not solved by JIT tools like torch.compile, nor do I expect it to be. I think a compiler would need to become a huge slow unreliable hairball to solve this type of auto-vectorization problem. Vexpr’s vectorizer is a compiler that solves this, but it does so by making the programmer meet the compiler in the middle. The programmer gives Vexpr a tree-like expression and tells Vexpr, “You should try vectorizing this.” Both of those pieces of information are valuable: the programmer structures the logic in a certain way, and the programmer indicates that there is opportunity for vectorization there. The programmer doesn’t need to do anything too difficult, and neither does the compiler.

I think this “meet-in-the-middle” approach between programmer and compiler is a good design principle that leads to good systems. And I think this will become even more true in the age of Large Language Models; rather than cramming too much magic into our compilers, let’s rely on humans-with-LLMs to meet the compiler in the middle.

(This project is supported by a GCP cloud compute grant from ML Collective, which has been super helpful. Thanks, also, to Rosanne Liu for useful feedback on drafts of this post.)

Gaussian Processes Extrapolate, Sometimes in Goofy Ways

Tue, 28 Mar 2023 02:00:00 -0700

Here is a toy function. (To see the code and more plots, check out this notebook.)

Figure 1: 80 random observations of a deterministic function (black) and the predicted maximal point in that function (orange), according to a Gaussian process trained on those 80 observations.

Intuitively, it seems clear that this function’s highest value probably occurs when x is in the center region. But a Gaussian Process (GP) thinks the highest value is out in a more mediocre region. This isn’t just a explore-exploit trade-off; the GP thinks the expected value is high at that orange point – higher than any value the GP has ever seen before!

I ran into this issue while exploring a real loss function, tuning hyperparameters with Bayesian Optimization. I found that my GP kept insisting there would be excellent results at unpromising points like the one above. The model sent me through hundreds of iterations of whack-a-mole, so I investigated the issue and came up with this toy example that makes the issue obvious.

This issue occurs because:

Gaussian Processes extrapolate, and they do it over an inferred distance, or “lengthscale”.
Sudden drop-offs in the function cause the GP to choose a low lengthscale. (This issue is most pronounced when testing deterministic functions. With noisy functions, these sudden drop-offs can be characterized as noise.)

So the GP makes local extrapolations that appear goofy given more global context.

Intuition: Why Gaussian Processes extrapolate

Consider a simple scenario with two observations and one prediction.

Figure 2: If you were a Gaussian process, what would you predict? You are given three random variables, A, B, and C, and the observed values for the first two. What is your prediction for C?

Gaussian Processes work by modeling this as a 3D multivariate Gaussian, using distance to determine correlation between variables. This clever trick transports us from thinking about functions over arbitrary spaces to thinking about Gaussian distributions with a finite number of variables.

Imagine another multivariate Gaussian that matches this correlation structure. Consider the performance of athletes in three different sports: Swimming (A), Cycling (B), and Running (C). Suppose correlations match distances in the chart above, with cov(A, B) = high, cov(B, C) = high, cov(A, C) = medium-to-low. In other words, elite swimmers are usually good cyclists but aren’t always great at running. Elite cyclists are usually good swimmers and runners. Elite runners are usually good cyclists but might not be great at swimming. A large set of unobserved latent causal factors underlie A, B, and C, and we don’t attempt to model these explicitly.

Suppose an athlete is a good cyclist. If we are given the extra context that she is a bad swimmer, does that increase the probability that she is a good runner? It depends on the actual covariance values, but there are definitely possible realities where that answer is yes. She doesn’t have the latent causal factors that make a good swimmer, and these latent factors also tend to make a good cyclist, yet she is still a good cyclist, so we expect she especially has the latent causal factors that make a good runner, to compensate.

If B is a medium value and A is low, we expect C to compensate for A. That's why Gaussian Processes extrapolate. (You can also show this by inverting a covariance matrix, but I like this playful explanation. See the appendix for something more rigorous.)

This is a good thing, usually.

Why extrapolation sometimes causes goofy predictions

Let’s zoom in on the section with the goofy prediction.

Figure 3: Zooming into the previous chart, we can see that the high prediction can be viewed locally as an extrapolation. Moreover, it is extrapolating from both left and right, which explains why the predicted value is extra high.

The GP is essentially making predictions using only local data. It is doing this because fitting the GP on this dataset caused the GP to treat everything extremely locally, i.e. it selected a very small “lengthscale”. It chose a small lengthscale because the function contains discrete drop-off points. Tiny changes in x sometimes yield large changes in f(x), and the GP is trying to accommodate this.

This leads to something that is undesirable in hyperparameter tuning: if ever you discover a sudden drop-off in performance from mediocre results to very-bad results, the GP will get excited that maybe there are very-good results immediately on the opposite side of the mediocre results. In my experience, it ends up spending nearly all of its time testing these crazy theories.

What does this mean? What can I do?

I’m not sure yet what is the best way to handle this.

One solution is to switch to using a Matern kernel with \(\nu = 0.5\), rather than the often-used value \(\nu = 2.5\). This specifies that the underlying function is not differentiable, allowing for sharp changes in direction on predictions.

Figure 4: Modifying the Matern kernel allows for sharp changes in the prediction’s slope.

Here’s a nice blog post on that subject. However, making this change runs the risk of hurting predictions – after all, these sudden dropoffs only occur in some parts of the space, and often those parts of the space are irrelevant.

It might be worth adopting the following view: a deterministic function with sudden discrete changes in output are bad news for a Gaussian Process. A broad solution is: figure out how to get rid of them.

You could do this using clamping / winsorizing tricks on your data; imagine taking the whole bottom row of points from Figure 1 and clamping them to be equal to the middle row.

Alternately, you could change your function so that it is not deterministic, for example by randomly initializing parameters or introducing something analogous to stochastic gradient descent. Then, presumably, these random dropoffs will occur randomly across a whole range, not at specific points, and the GP will capture this as observation noise. Moving away from a deterministic function has other benefits anyway, e.g. greatly reducing the risk of overfitting your hyperparameters to the validation set.

One takeaway is that you really need to pay attention to what your Gaussian Process is doing. It will not always automatically do what you consider intuitive.

Appendix

In case you’re interested, here are the GP predictions at other points, with standard deviations:

To get more rigorous with the simple A, B, C example, suppose we assign constant prior mean 0, as is common in Gaussian Processes. Denoting \(\Sigma_{XY} : cov(X,Y)\), the expected value for C is:

\[E\left[C \mid A=a, B=b\right] = \frac{a(\Sigma_{AC}\Sigma_{BB} - \Sigma_{AB}\Sigma_{BC}) + b(\Sigma_{BC}\Sigma_{AA} - \Sigma_{BA}\Sigma_{AC})}{\Sigma_{AA}\Sigma_{BB} - \Sigma_{AB}^2}\]

This can be derived from these equations, specifically the one for \(\boldsymbol \mu_*\).

It is difficult to understand this function at a glance, but it is interesting to simply ask, “What happens as \(a\) increases?” We see that: if the correlative chain from \(A \to B \to C\) is strong, but the direct correlation between \(A \to C\) is weak, i.e. \(\Sigma_{AB}\Sigma_{BC} > \Sigma_{AC}\Sigma_{BB}\), then increasing \(a\) will decrease \(E[C \mid ...]\). Thus, reducing \(a\) in isolation will increase \(E[C \mid ...]\), which is an example of extrapolation.

(This work was supported by ML Collective via their donated GCP compute resources which have been super helpful. Thanks, also, to Rosanne Liu for useful feedback on early drafts of this post.)

Maybe Bayesian Optimization Should Be Harder, Not Easier

Wed, 30 Nov 2022 01:00:00 -0800

(2023-11-01 updates: Refreshed the charts, other small tweaks, posted reproducible experiments.)

This post has two parts:

A fun point of view that I find compelling
A project that follows from that point of view, with some early results

Point of view

When you tune an AI model’s design and its training regime, you are exploring a search space. Bayesian Optimization is a framework that naturally arises when you try to automate your intuitive process of exploring a search space. The promise of Bayesian Optimization is: “If I can tell a computer my beliefs about this search space, then the computer can perform the search better than I can.”

From what I’ve seen, AI / ML people tend to feel burned by their experiences with Bayesian Optimization. Often I hear stories where they once invested a bunch of time into trying it on toy projects, but it never managed to cross over and be useful in their actual work where they continue to use manual search and non-sophisticated brute force searches. Often they feel guilty about all the time they wasted trying Bayesian Optimization.

Here is a slightly contrarian view on why Bayesian Optimization has not yet become ubiquitous.

Lots of people have tried to increase adoption of Bayesian Optimization by making it easier, creating products and libraries that hide the messy details. My hunch is that existing Bayesian Optimization solutions block / discourage the user from taking ownership of the process. In particular, there is a lot of unrealized potential in giving users more powerful ways to state their beliefs about a search space. I think what is needed is a set of tools / documentation / cookbooks designed for people who are willing to put in the time to master their tools. Existing tools could grow to fill this role.

I think Bayesian Optimization ought to be less one-size-fits-all. I think it should be more hands-on.

Taking ownership of your search

If you use a one-size-fits-all Bayesian model, your search will be data-intensive and will often be much less efficient than a manual search. Results will be much better if you invest some time telling the model your priors about your parameters. This is more fun than you might expect.

Consider: when you manually explore a search space, you have some intuition about which details of the model are “orthogonal”. For example, hyperparameters describing a neural network’s architecture seem like they are roughly independent of hyperparameters that describe the training regime, and we can afford to optimize these groups of parameters almost independently. Meanwhile, hyperparameters like “momentum” and “learning rate” are intertwined, and it is important to explore their joint space. A one-size-fits-all Bayesian Optimization model totally lacks this intuition. Fortunately, it is easy to input these beliefs to a Gaussian Process (GP). If you have two groups of parameters that are independent, simply dedicate a kernel to each one, then add the kernels together. If they are dependent, multiply the kernels together. If it’s somewhere in between, do both, and allow the model to adapt to the data. If you run your own Bayesian Optimization code locally, this is all trivial to implement, but it is impossible in any Bayesian Optimization products that I’ve seen. (Some products use a fully AutoML approach to infer these dependencies, but this is a data-intensive process.)

In owning your search, another important degree of freedom is choosing useful features for the GP kernel. Your hyperparameters might not be a good basis for a GP. You might want to configure your GP to transform the hyperparameter vector into a basis that enables the GP to make better predictions with fewer data points, maybe simply by adding a few redundant features to the hyperparameter vectors that may prove useful. This insight can be combined with the “orthogonal” trick above to greatly reduce the effective size of the search space.

Opportunities also lie in having fine-grained control over the optimization loop. For example, I have found it useful to create a “tick-tock” optimization loop that alternates between multiple objectives (accuracy and training time).

The project: Is hands-on Bayesian Optimization worthwhile?

This question has multiple facets, each involving wearing a different hat:

Is it usable? What is the right user interface? What is the overall on-ramp / learning process for a person motivated to learn hands-on Bayesian Optimization?
Is it fast? I am proposing we use more complicated kernels. Is that tractible, from a performance standpoint?
Does it work? Are there considerable benefits in hands-on Bayesian Optimization, relative to a one-size-fits-all approach? Is it much better than random search?

Is it usable?

There is a user interface problem here that needs solving. The first UI to build is a first-class programming interface. Later we can consider whether a graphical UI is needed.

I draw inspiration from D3, a library for data visualization. In a world of people trying to build “easy” visualization tools, D3 took a different approach. D3 doesn’t generate visualizations; D3 helps you program your own visualizations. Rather than exposing some magic function that outputs a customized pre-built visualization, D3 provides a useful set of primitives, then it gets out of the way and lets you write code. It takes time to learn, but once you have learned it, you are powerful. It’s usually a tortoise-and-the-hare phenomenon; the person who chooses D3 ends up winning the race (and winning with style).

I think BoTorch (with gpytorch) is a good start at filling this role for Bayesian Optimization, but it’s not finished. It provides the core algorithmic tools of Bayesian Optimization, but doesn’t provide a way for describing a search space (for example). When you read botorch tutorials, they essentially recommend that you don’t use botorch directly, but instead use a wrapper like Ax, but Ax is more of a one-size-fits-all framework, adding a lot of unnecessary friction to hands-on Bayesian Optimization. I would rather see botorch grow outward into a usable library, giving you all the pieces you need to write your own Bayesian Optimization. They could aim to be the D3 of Bayesian Optimization. I found myself needing to “finish” botorch with my own makeshift library that adds search spaces with conditional parameters, an expanded set of input transform utilities, and random search utilities.

Is it fast?

I am proposing that we take one-size-fits-all GP kernels and replace them with customized compositions of kernels. I was worried that these compositions of kernels would be slow, and, yes, by default, they are. But if you write them efficiently, they are very fast. I ended up writing a WeightedSPSMaternKernel, i.e. a “weighted sum of products of sums Matern kernel” which runs many kernels in parallel. This satisfied my needs.

Otherwise, I’m also proposing we start performing transformations on feature vectors. This is already a supported use case and is fast.

Optional reading: In case you want details, here is my code that describes a 22-dimensonal hyperparameter space and how it is transformed into a search space (for random search / optimization) and then into a feature space (for the GP). Vexpr and outerloop are libraries I created. Vexpr makes it possible to write fast, readable compositional GP kernels, and outerloop attempts to “finish” botorch.

Expand code view

import botorch
import gpytorch
import outerloop as ol
import outerloop.vexpr.torch as ovt
import torch
import vexpr as vp
import vexpr.torch as vtorch
import vexpr.custom.torch as vctorch

parameter_space = [
    ol.Choice("optimizer", ["adam", "sgd"]),
    ol.Choice("nesterov", [True, False],
              condition=lambda choices: choices["optimizer"] == "sgd"),
    ol.Int("epochs", 2, 60),
    ol.Int("batch_size", 16, 4096),

    ol.Scalar("conv1_weight_decay", 1e-7, 3e-1),
    ol.Scalar("conv2_weight_decay", 1e-7, 3e-1),
    ol.Scalar("conv3_weight_decay", 1e-7, 3e-1),
    ol.Scalar("dense1_weight_decay", 1e-7, 3e-1),
    ol.Scalar("dense2_weight_decay", 1e-7, 3e-1),

    ol.Scalar("1cycle_initial_lr_pct", 1/80, 1/2),
    ol.Scalar("1cycle_final_lr_pct", 1/30000, 1/100),
    ol.Scalar("1cycle_pct_warmup", 0.01, 0.5),
    ol.Scalar("1cycle_max_lr", 0.01, 20.0),
    ol.Scalar("1cycle_max_momentum", 0, 0.9999,
              condition=lambda choices: choices["optimizer"] == "sgd"),
    ol.Scalar("1cycle_min_momentum_pct", 0.0, 1.0,
              condition=lambda choices: choices["optimizer"] == "sgd"),

    ol.Int("conv1_channels", 4, 64),
    ol.Int("conv2_channels", 8, 128),
    ol.Int("conv3_channels", 16, 256),
    ol.Int("dense1_units", 8, 256),
]

# Transform to log-space
xform = ol.transforms.ToScalarSpace(
    space,
    ol.transforms.log({
        "epochs": "log_epochs",
        "batch_size": "log_batch_size",
        "1cycle_initial_lr_pct": "log_1cycle_initial_lr_pct",
        "1cycle_final_lr_pct": "log_1cycle_final_lr_pct",
        "1cycle_max_lr": "log_1cycle_max_lr",
        "1cycle_pct_warmup": "log_1cycle_pct_warmup",
        "1cycle_momentum_max_damping_factor":
        "log_1cycle_momentum_max_damping_factor",
        "1cycle_momentum_min_damping_factor_pct":
        "log_1cycle_momentum_min_damping_factor_pct",
        "1cycle_beta1_max_damping_factor":
        "log_1cycle_beta1_max_damping_factor",
        "1cycle_beta1_min_damping_factor_pct":
        "log_1cycle_beta1_min_damping_factor_pct",
        "beta2_damping_factor": "log_beta2_damping_factor",
        "conv1_weight_decay": "log_conv1_weight_decay",
        "conv2_weight_decay": "log_conv2_weight_decay",
        "conv3_weight_decay": "log_conv3_weight_decay",
        "dense1_weight_decay": "log_dense1_weight_decay",
        "dense2_weight_decay": "log_dense2_weight_decay",
        "conv1_channels": "log_conv1_channels",
        "conv2_channels": "log_conv2_channels",
        "conv3_channels": "log_conv3_channels",
        "dense1_units": "log_dense1_units",
    })
)


class VexprHandsOnLossModel(botorch.models.SingleTaskGP):
    def __init__(self, ...)

       # Transforms for the kernel

        xforms += [
            ol.transforms.append_mean(
                ["log_conv1_channels", "log_conv2_channels",
                 "log_conv3_channels", "log_dense1_units"],
                "log_gmean_channels_and_units"),
            ol.transforms.subtract(
                {"log_conv1_channels": "log_conv1_channels_div_gmean",
                 "log_conv2_channels": "log_conv2_channels_div_gmean",
                 "log_conv3_channels": "log_conv3_channels_div_gmean",
                 "log_dense1_units": "log_dense1_units_div_gmean"},
                "log_gmean_channels_and_units"),
            ol.transforms.add(
                {"log_1cycle_initial_lr_pct": "log_1cycle_initial_lr",
                 "log_1cycle_final_lr_pct": "log_1cycle_final_lr"},
                "log_1cycle_max_lr"),
            ol.transforms.add(
                {"log_1cycle_momentum_min_damping_factor_pct":
                 "log_1cycle_momentum_min_damping_factor"},
                "log_1cycle_momentum_max_damping_factor"),
            ol.transforms.add(
                {"log_1cycle_beta1_min_damping_factor_pct":
                 "log_1cycle_beta1_min_damping_factor"},
                "log_1cycle_beta1_max_damping_factor"),
            ol.transforms.append_mean(
                ["log_conv1_weight_decay", "log_conv2_weight_decay",
                 "log_conv3_weight_decay", "log_dense1_weight_decay",
                 "log_dense2_weight_decay"],
                "log_gmean_weight_decay"),
            ol.transforms.subtract(
                {"log_conv1_weight_decay": "log_conv1_wd_div_gmean",
                 "log_conv2_weight_decay": "log_conv2_wd_div_gmean",
                 "log_conv3_weight_decay": "log_conv3_wd_div_gmean",
                 "log_dense1_weight_decay": "log_dense1_wd_div_gmean",
                 "log_dense2_weight_decay": "log_dense2_wd_div_gmean"},
                "log_gmean_weight_decay"),
            partial(ol.transforms.ChoiceNHotProjection,
                    out_name=N_HOT_PREFIX)
        ]

        # ...


def make_handson_kernel(space, batch_shape=()):
    """
    This kernel attempts to group parameters into orthogonal groups, while
    also always allowing for the model to learn to use the joint space.
    """
    zero_one_exclusive = partial(gpytorch.constraints.Interval,
                                 1e-6,
                                 1 - 1e-6)

    state = State(batch_shape)

    ialloc = IndexAllocator()

    lengthscale = vp.symbol("lengthscale")
    x1 = vp.symbol("x1")
    x2 = vp.symbol("x2")

    def index_for_name(name):
        return next(i for i, p in enumerate(space) if p.name == name)

    def scalar_kernel(names):
        ls_indices = ialloc.allocate(len(names))
        indices = [index_for_name(name) for name in names]
        return ovt.matern(
            vtorch.cdist(x1[..., indices] / lengthscale[ls_indices],
                         x2[..., indices] / lengthscale[ls_indices],
                         p=2),
            nu=2.5)

    def choice_kernel(names):
        ls_indices = ialloc.allocate(len(names))
        indices = [index_for_name(name) for name in names]
        return ovt.matern(
            vtorch.cdist(x1[..., indices] / lengthscale[ls_indices],
                         x2[..., indices] / lengthscale[ls_indices],
                         p=1),
            nu=2.5)


    def scalar_factorized_and_joint(names, suffix):
        w_additive = vp.symbol("w_additive" + suffix)
        alpha_factorized_or_joint = vp.symbol("alpha_factorized_or_joint"
                                              + suffix)
        state.allocate(w_additive, (len(names),),
                       zero_one_exclusive(),
                       ol.priors.DirichletPrior(torch.full((len(names),), 2.0)))
        state.allocate(alpha_factorized_or_joint, (),
                       zero_one_exclusive(),
                       ol.priors.BetaPrior(4.0, 1.0))
        return vtorch.sum(
            vctorch.heads_tails(alpha_factorized_or_joint)
            * vtorch.stack([
                vtorch.sum(
                    w_additive
                    * vtorch.stack([scalar_kernel([name])
                                    for name in names],
                                   dim=-1),
                    dim=-1),
                scalar_kernel(names),
            ], dim=-1),
            dim=-1
        )

    def regime_kernels():
        return [
            # kernel: regime choice parameters
            choice_kernel([f"{N_HOT_PREFIX}{i}"
                           for i in range(4)]),

            # kernel: lr schedule
            scalar_factorized_and_joint(
                ["log_1cycle_initial_lr", "log_1cycle_final_lr",
                 "log_1cycle_max_lr", "log_1cycle_pct_warmup"],
                "_lr"),

            # kernel: momentum schedule
            scalar_factorized_and_joint(
                ["log_1cycle_momentum_max_damping_factor",
                 "log_1cycle_momentum_min_damping_factor",
                 "log_1cycle_beta1_max_damping_factor",
                 "log_1cycle_beta1_min_damping_factor",
                 "log_beta2_damping_factor"],
                "_momentum"),

            # kernel: relative weight decay
            scalar_factorized_and_joint(
                ["log_conv1_wd_div_gmean", "log_conv2_wd_div_gmean",
                 "log_conv3_wd_div_gmean", "log_dense1_wd_div_gmean",
                 "log_dense2_wd_div_gmean"],
                "_wd"),
        ]

    regime_joint_names = ["log_epochs", "log_batch_size",
                          "log_gmean_weight_decay"]

    def architecture_kernels():
        return [
            # kernel: lr schedule
            scalar_factorized_and_joint(["log_conv1_channels_div_gmean",
                                         "log_conv2_channels_div_gmean",
                                         "log_conv3_channels_div_gmean",
                                         "log_dense1_units_div_gmean"],
                                        "_units_channels"),
        ]

    architecture_joint_names = ["log_gmean_channels_and_units"]

    regime_kernel = vctorch.fast_prod_positive(
        vtorch.stack(([scalar_kernel(regime_joint_names)]
                      + regime_kernels()),
                     dim=-1),
        dim=-1)
    architecture_kernel = vctorch.fast_prod_positive(
        vtorch.stack(([scalar_kernel(architecture_joint_names)]
                      + architecture_kernels()),
                     dim=-1),
        dim=-1)
    joint_kernel = vctorch.fast_prod_positive(
        vtorch.stack(([scalar_kernel(regime_joint_names
                                     + architecture_joint_names)]
                      + regime_kernels()
                      + architecture_kernels()),
                     dim=-1),
        dim=-1)

    alpha_regime_vs_architecture = vp.symbol("alpha_regime_vs_architecture")
    alpha_factorized_vs_joint = vp.symbol("alpha_factorized_vs_joint")
    scale = vp.symbol("scale")

    state.allocate(alpha_regime_vs_architecture, (),
                   zero_one_exclusive(),
                   ol.priors.BetaPrior(2.0, 2.0))
    state.allocate(alpha_factorized_vs_joint, (),
                   zero_one_exclusive(),
                   ol.priors.BetaPrior(4.0, 1.0))
    state.allocate(scale, (),
                   gpytorch.constraints.GreaterThan(1e-4),
                   gpytorch.priors.GammaPrior(2.0, 0.15))

    kernel = (scale
              * vtorch.sum(vctorch.heads_tails(alpha_factorized_vs_joint)
                           * vtorch.stack(
                               [vtorch.sum(
                                   vctorch.heads_tails(
                                       alpha_regime_vs_architecture)
                                   * vtorch.stack([
                                       regime_kernel,
                                       architecture_kernel
                                   ], dim=-1),
                                   dim=-1),
                                joint_kernel],
                               dim=-1),
                           dim=-1))

    state.allocate(lengthscale, (ialloc.count,),
                   gpytorch.constraints.GreaterThan(1e-4),
                   gpytorch.priors.GammaPrior(3.0, 6.0))

    return kernel, state.modules

You can read the rest of the code here.

Does it work?

My initial results are promising, but my baseline (the default BoTorch model) has some issues that might need to be ironed out.

I ran 160 semi-random MNIST LeNet training experiments on a 22-dimensional hyperparameter space. I fed the results to a Gaussian Process and checked its predictive power for held out experiments, performing leave-one-out cross-validation.

For a baseline one-size-fits-all model, I was hoping to just use Botorch’s default MixedSingleTaskGP off the shelf, but it does not use priors on its parameters and is hence prone to overfit and output extreme low log probabilities for held-out results. I added priors from one of their other models, and this partially solved the issue.

For the hands-on model, I designed a kernel off the top of my head, trying to capture some of my intuition about my hyperparameters. This was without any iteration, I just guessed at what kernel might work, chose priors that seemed sensible, and now I’m sharing my first results.

As a sanity-check, here is the raw data for cross-validation 60 of the 160 experiments. To generate each point, the model is trained on 59 configurations and outputs, then must predict the 60th output, given its configuration. (Click or zoom to see details.)

Here are the results in aggregate, using two different success metrics, and plotted for different subsets of the 160 experiments. I gathered this on 50 different shuffles of the data and I plot the mean.

On average, for small-to-medium datasets the “hands-on” model assigned higher probability density to the “correct” output for the held-out data. (It is important to use geometric mean, i.e. take the mean of the log probability, so that we are sure to penalize very low probabilities.) When using the model in MLE mode (maximum likelihood estimator) to output a single prediction, the hands-on model on average has lower prediction error.

Toward the end, the hands-on model falls behind the baseline. The MLE predictions aren’t obviously worse, but the mean log probability is, suggesting that hands-on model gives more confident predictions than it ought to. Thus, there is room for improvement. Still, it is promising that my first try was this good. I suspect a little bit of iteration on the kernel would lead to one that is superior for all dataset sizes.

Early in the chart, the baseline displays a strange phenomenon where it fails to improve as it receives more data. It is worth doing some due diligence here and try to understand what is going on. Maybe the baseline has a low-hanging possible improvement, and maybe that improvement could also be brought to the hands-on model.

My conclusions from this experiment are:

The initial results are promising, but not conclusive.
The weird baseline phenomenon illustrates my point that you should be careful treating a Bayesian Optimization as a black box. I suspect Bayesian Optimization is only worthwhile if we figure out how to make users take ownership of their search.

(This project is supported by a GCP cloud compute grant from ML Collective, which has been super helpful.)