Spoiler: Category theory has applications in machine learning engineering

Note: The code lives in ianqs - applied_category_theory/2_functors

Recap

You’re an engineer on a ML team working on distributed hyperparameter search. You have a collection of worker machines, each with a copy of the dataset. Each worker trains a collection of configurations,recording a validation loss for each one. In the previous post on monoids, from July *2022, you implemented a monoid structure that cleaned up a lot of messy code. The structure is listed below to jog your memory (it has been 4 years, after all):

@dataclass(frozen=True)
class Summary:
    """Summary statistics over a set of validation losses.

    This is a monoid:
      - Identity: `Summary()` an empty summary
      - Associative: we've defined + as an associative binary op
      - Closed: `+` over Summaries is always a Summary
    """
    count: int = 0
    total: float = 0.0
    total_sq: float = 0.0
    minimum: float = math.inf

    @classmethod
    def of(cls, loss: float) -> Summary:
        """Lift a single loss into a one-element Summary."""
        return cls(count=1, total=loss, total_sq=loss * loss, minimum=loss)

    def __add__(self, other: Summary) -> Summary:
        return Summary(
            count=self.count + other.count,
            total=self.total + other.total,
            total_sq=self.total_sq + other.total_sq,
            minimum=min(self.minimum, other.minimum),
        )

    # Derived quantities, read off once at the very end.
    @property
    def mean(self) -> float:
        return self.total / self.count if self.count else math.nan

    @property
    def variance(self) -> float:
        # E[x^2] - E[x]^2
        return self.total_sq / self.count - self.mean ** 2 if self.count else math.nan

we motivated the monoid by showing that it is cleaner and more extensible than storing raw dictionaries that we pass around.

New Requirements

As part of increased train job monitoring, the big boss has come back with a few requirements:

• get the best model(s) (not just the best, but maybe the top-N).
• transform the data, like mark if some of them have been reviewed, if some of them had a certain flag on them, etc.
• filter the data e.g. by a loss range, or filter the models that have not yet been reviewed

A simple list would work for smaller runs, but we work at #LARGE_CORPORATION, and we want to make sure our code is web scale. Clearly, other data structures are a better alternative, especially if we want to grab the models that are within some range of values; consider a model which is within 1\% of loss of the best model, but has 10x faster inference, where speed is critical - obviously we choose that model.

A first pass: just use a list

In the monoid post every run got reduced into a final Summary, where the individual results were thrown away. But you can’t retrieve the top-N models from a Summary. For this post, we maintain the individual results$^1$:

@dataclass(frozen=True)
class Run:
    loss: float
    config: dict


def run_config(config: dict) -> Run:
    # No failures here. We assume runs succeed and revisit that in the monad post.
    return Run(loss=round(random.random(), 4), config=config)

The container we might reach for is a list:

def best(runs: list[Run]) -> Run:
    return min(runs, key=lambda run: run.loss)

def between(runs: list[Run], low: float, high: float) -> list[Run]:
    return [run for run in runs if low <= run.loss <= high]    # scans everything - O(N)

The two other requests: transform (map), and filter, come for free - a list already gives us these things:

penalized = list(map(lambda run: replace(run, loss=run.loss + 1.0), runs))
good      = list(filter(lambda run: run.loss < 0.5, runs))

With that said, hold on to that thought, because map and filter “already being there” is the whole point of this post.

So what’s wrong? Nothing, until we hit #WebScale. Both best and between scan the entire container every single time you call them. At twenty runs, the computation is trivial - my 2011 x220 can easily handle this$^2$. At a few hundred thousand runs, with a dashboard requesting best and between on every refresh, you’re doing a full O(N) operation every time. We can do a lot better if the container keeps itself sorted.

$^1$: Note, we could also have a max/min-heap, and instead of just using the validation loss as our key, we could “weight” the loss and inference time.

$^2$: That 15 year old X220 is probably still my favorite laptop of all time - if I could get it to x64 I’d daily drive that bad boy again.

Swapping the container

Being at #LARGE_CORPORATION, and working at #WebScale, we know a list won’t scale. We reach for every CS-101 student’s best friend - the binary search tree. Suddenly, best becomes “walk left until you can’t” in O(log N), and between gets to prune whole subtrees that fall outside the range instead of scanning everything 😎 (a range query lands k results, so it’s O(log N + k) rather than O(N)):

Lots of notes in the post but another note: our LossTree is a plain, unbalanced BST, so those bounds assume it stays roughly balanced. I’d use some existing data structure either from a python library or build an AVL or Red-Black tree in prod; the code and core concepts will work the same anyways, which is the whole point.

class LossTree:
    """A binary search tree of Runs keyed by validation loss."""

    ...

    def best(self) -> Run:
        node = self._root
        if node is None:
            raise ValueError("no runs")
        while node.left is not None:      # leftmost node holds the lowest loss
            node = node.left
        return node.run

    def between(self, low: float, high: float) -> LossTree:
        keep: list[Run] = []
        self._between(self._root, low, high, keep)
        return LossTree.from_runs(iter(keep))

However, notice what we lost: list implements map and filter, so we lost those and need to implement them for our LossTree:

    def map(self, f: Callable[[Run], Run]) -> LossTree:
        return LossTree.from_runs(f(run) for run in self.items())

    def filter(self, keep: Callable[[Run], bool]) -> LossTree:
        return LossTree.from_runs(run for run in self.items() if keep(run))

Lets look at an example of using filter for both the list and LossTree:

list_impl = filter(lambda run: run.loss < 0.5, runs)
losstree_impl = runs.filter(lambda run: run.loss < 0.5)

SHEESH - basically the same (okay, list_impl needs to call a list() on the result to convert it back into a list but they’re basically the same). What’s the takeaway from this? Our downstream consumers - customers, other teams, CXOs who are visualizing the results - don’t need to change anything. The map and filter contract/signature did not change AT ALL. We could have swapped out the BST for an AVL, Red-Black, etc. and nothing would change.

So what is a functor?

With that all out of the way, what is a functor? As with the monoid, it comprises a few ideas:

• a container that holds values of some type e.g. list, LossTree, …
• a map that operates over the value inside, while returning the same container type e.g. we can’t go from list to AVLTree

And that’s really it - map operates on the contents of the functor while maintaining the container type. That’s all we care about - the “contract” is obeyed so the downstream consumers don’t have to worry about what’s happening under the hood (sorry to really beat this point into the ground, but I think it’s a good way to remember WHY it’s structured this way.)

NOTE: filter is a nice-to-have but it is not part of the definition of a functor - I just added it in because, other than map, a filter is one of the most-used operations over a container.

The Haskell signature

Haskell writes the functor interface as a single line:

fmap :: (a -> b) -> f a -> f b

Read it left to right. fmap takes a function (a -> b), an ordinary function from some type a to some type b. It takes an f a, a container f full of as (our LossTree full of Runs). It gives back an f b, the same container now full of bs. f is the part that stays put (our second condition from before, we can’t go from a LossTree to a list); a turning into b is the part that moves. One thing worth calling out: in our example we’re not actually changing the type, we do Run -> Run rather than Run -> X, but fmap is happy to let b be a different type from a.

The functor laws

So, what are the actual requirements? We discussed the signature of map earlier, but a matching signature isn’t enough. For something to really be a functor, its map has to obey two laws, and these are properties of map itself that have to hold for every function you pass in:

1. identity: the no-op does “nothing”

tree.map(lambda run: run) == tree

1. composition: given two functions, f and g:

map(f . g) == map(f) . map(g)

said differently, mapping g and then mapping f gives the same result as mapping “g then f” in a single pass. We illustrate this directly in functors.py:

def test_composition(tree: LossTree):
    """
    We don't bother to test the:

        - Functors must preserve identity morphisms

    because that's pretty trivial.
    """
    def shift(run: Run) -> Run:
        """Transform the underlying data"""
        return replace(run, loss=round(run.loss + 1.0, 4))

    def tag(run: Run) -> Run:
        """Tag the run as reviewed"""
        return replace(run, config={**run.config, "reviewed": True})

    # f . g
    composed = tree.map(lambda run: tag(shift(run)))

    # g .f
    chained = tree.map(shift).map(tag)
    assert list(composed.items()) == list(chained.items())
    print("Composition law + Chaining Holds")

The assertion passes! Why is this significant? Let’s open up functors.py:

    def map(self, f: Callable[[Run], Run]) -> LossTree:
        """(fmap . f) may change the loss (key), so we rebuild and re-key by the new value."""
        return LossTree.from_runs(f(run) for run in self.items())

    def filter(self, keep: Callable[[Run], bool]) -> LossTree:
        """(filter . f) may change the loss (key), so we rebuild and re-key by the new value."""
        return LossTree.from_runs(run for run in self.items() if keep(run))

See how chained walks and rebuilds the entire tree twice (once for shift, once for tag), whereas composed fuses both functions and rebuilds it once. Depending on the cost of the tree rebuilding, this equates to huge savings.

Retrospective

We did something here that SOUNDS simple, but was actually quite significant. We implemented an interface, functor, that allows us to abstract over containers. In all honesty, you’ve probably already implemented something similar, even if you weren’t familiar with the name of it.

In the next post we will handle the final portion, elegantly handling failures. Stay tuned!