07 — A real architecture, and three things wrong with it

Three modalities, three kinds of network, one head. Convolutions with residuals, a transformer stack, a recurrent cell, and a bottleneck where they meet — some of it inside one node and some of it across several, which is the choice a graph exists to let you make.

Every section has the same four parts, and the fourth is the one the other notebooks were missing:

the problematic architecture and configuration
the symptoms, read off what the framework measured
the solution, and why that one
the healthy architecture and configuration, run and shown to be healthy

import torch

import somatize.torch  # noqa: F401
from somatize import Graph, Node, Opaque, Recorder, Store
from somatize.health import alerts, diagnose, overlaid, profile, seen
from somatize.record import gantt, progress
from somatize.torch import Audit, Trainer, architecture, parameters

torch.manual_seed(0)
store = Store(__import__("tempfile").mkdtemp())

DIM, STEPS, VOCAB, ROWS, OUT = 24, 16, 60, 64, 4

The architecture

Inside one node: the convolutional trunk is a stem plus three residual blocks, all of it audio’s business. Nobody outside it needs to know.

Across nodes: the three modalities are three nodes, because they are three things that can be placed, cached, sent to another machine and diagnosed apart. That is what a node is for.

class Residual(torch.nn.Module):
    """Norm, non-linearity, convolution, and the input added back on."""

    def __init__(self, channels, norm="batch"):
        super().__init__()
        self.norm = (
            torch.nn.BatchNorm1d(channels) if norm == "batch" else torch.nn.Identity()
        )
        self.act = torch.nn.GELU()
        self.conv = torch.nn.Conv1d(channels, channels, 3, padding=1)

    def forward(self, x):
        return x + self.conv(self.act(self.norm(x)))


class Audio(Node):
    """A convolutional trunk — a whole architecture inside one node.

    `keeps` is how the trunk turns a sequence into a vector, which is the one
    architectural decision in here that throws information away.
    """

    def __init__(self, channels=DIM, blocks=3, norm="batch", keeps="mean"):
        self.stem = torch.nn.Conv1d(1, channels, 5, padding=2)
        self.body = torch.nn.Sequential(*[Residual(channels, norm) for _ in range(blocks)])
        self.keeps = keeps

    def forward(self, said, ctx):
        h = self.body(self.stem(said["audio"]))
        if self.keeps == "mean":
            return Opaque(h.mean(-1))
        return Opaque(torch.cat([h.mean(-1), h.max(-1).values], dim=1))

    def parameters(self):
        return list(self.stem.parameters()) + list(self.body.parameters())


class Text(Node):
    """An embedding and a stack of transformer layers."""

    def __init__(self, dim=DIM, layers=4):
        self.emb = torch.nn.Embedding(VOCAB, dim)
        self.body = torch.nn.TransformerEncoder(
            torch.nn.TransformerEncoderLayer(dim, 4, 48, batch_first=True), num_layers=layers
        )

    def forward(self, said, ctx):
        return Opaque(self.body(self.emb(said["text"])).mean(1))

    def parameters(self):
        return list(self.emb.parameters()) + list(self.body.parameters())


class Vitals(Node):
    """A recurrent cell over a short series."""

    def __init__(self, dim=DIM, kind="gru"):
        self.cell = (torch.nn.GRU if kind == "gru" else torch.nn.RNN)(
            4, dim, batch_first=True
        )

    def forward(self, said, ctx):
        return Opaque(self.cell(said["vitals"])[0][:, -1])

    def parameters(self):
        return list(self.cell.parameters())


class Fuse(Node):
    """Where the three meet, through a bottleneck."""

    def __init__(self, wide=3 * DIM, squeeze=16):
        self.down = torch.nn.Linear(wide, squeeze)
        self.act = torch.nn.GELU()
        self.up = torch.nn.Linear(squeeze, OUT)

    def forward(self, said, ctx):
        return Opaque(self.up(self.act(self.down(torch.cat(list(said.values()), dim=1)))))

    def parameters(self):
        return list(self.down.parameters()) + list(self.up.parameters())


def built(*, blocks=3, layers=4, squeeze=16, keeps="mean+max"):
    wide = (2 * DIM if keeps != "mean" else DIM) + 2 * DIM
    return Graph.somatize(
        (
            Audio(blocks=blocks, keeps=keeps).named("audio")
            | Text(layers=layers).named("text")
            | Vitals().named("vitals")
        )
        >> Fuse(wide=wide, squeeze=squeeze).named("fuse")
    )

The data

Three channels about a person and four things to predict. Two of the four are in the shape of the audio — its first half against its second — and not in its level. That distinction is the whole of Problem 1: a trunk that reduces a sequence to its mean has thrown those two away before anything downstream ever sees them.

def people(how_many=6, seed=0):
    torch.manual_seed(seed)
    made = []
    for _ in range(how_many):
        said = {
            "audio": torch.randn(ROWS, 1, 32),
            "text": torch.randint(0, VOCAB, (ROWS, STEPS)),
            "vitals": torch.randn(ROWS, STEPS, 4),
        }
        shape = said["audio"][:, 0, :16].mean(1) - said["audio"][:, 0, 16:].mean(1)
        made.append(
            (
                {one: Opaque(what) for one, what in said.items()},
                torch.stack(
                    [
                        3.0 * shape,
                        3.0 * shape.abs(),
                        0.5 * said["vitals"].mean((1, 2)),
                        0.5 * (said["text"].float().mean(1) / VOCAB),
                    ],
                    dim=1,
                ),
            )
        )
    return made


DATA = people()
ONE = DATA[0][0]


def scored(g, data=None):
    with torch.no_grad():
        data = data or DATA
        return sum(
            float(torch.nn.functional.mse_loss(g.forward(one), y)) for one, y in data
        ) / len(data)


def fit(g, run, *, steps=200, lr=3e-4, auditing=True, data=None):
    t = Trainer(
        g,
        objective=torch.nn.functional.mse_loss,
        optimizer=torch.optim.Adam(parameters(g), lr=lr),
        auditing=Audit(inside=True, every=4) if auditing else None,
        watching=Recorder(store, run=run, summarising=["loss"]),
    )
    data = data or DATA
    for which in range(steps):
        t.step(data[which % len(data)])
    return g

What it looks like

The architecture inside each node, traced by running the graph once. A skip connection is an edge, a repeated block is a ×N, a composite everybody recognises is one box, and the shape on every layer is what makes the bottleneck visible.

healthy = built()
architecture(healthy, ONE)

{'audio': Inside(6 layers, 6 edges, traced),
 'text': Inside(2 layers, 1 edges, traced),
 'vitals': Inside(1 layers, 0 edges, traced),
 'fuse': Inside(3 layers, 2 edges, traced)}

healthy.figure(inside=architecture(healthy, ONE))

Figure from cell 8

Three branches side by side inside a wave — they run at once — and the bottleneck at the join.

Every kind has its own silhouette, because a Linear, a convolution, a recurrent cell and a non-linearity are four different kinds of thing and four identical rectangles with different words in them make the reader do the sorting a picture was supposed to have done:

	drawn as
a convolution	a parallelogram — a window sliding along
a recurrent cell	a box with a tab: it feeds itself
an attention block	a box with its corners cut, and a line saying what is in it
a normalisation	a capsule: no capacity, and thin
a non-linearity	pointed at both ends: nothing lives in it
anything that changes the width	tapered the way it really goes

That last one is what makes a bottleneck look like one instead of like three boxes with different numbers written on them.

And every number says what it is. 4 batch · 16 steps · 24 dim is a shape somebody can read; 4×16×24 is three numbers and a puzzle. The batch is the one that can be checked rather than assumed — the caller knows how many rows went in — and a layer that did not change the shape keeps the names of the one that did, so a BatchNorm1d in a convolutional trunk says ch and len and not steps and dim.

Opening a composite

A TransformerEncoderLayer is one box because read as its fourteen leaves it is fourteen things. When the inside of one is the question, depth= opens it — and depth counts composites opened rather than names, because that layer sits three names deep inside a TransformerEncoder and asking for one level of detail should not have to know that.

healthy.figure(inside=architecture(healthy, ONE, depth=1))

Figure from cell 11

Problem 1 — the join is a bottleneck

The problematic architecture and configuration

Three branches of twenty-four each, meeting through two numbers.

pinched = built(squeeze=2)
fit(pinched, "1-problem")
print("loss :", round(scored(pinched), 4))
print("flags:", diagnose(store, run="1-problem") or "nothing tripped")

loss : 0.3539
flags: nothing tripped

The symptoms

No flag at all. Every gradient arrives, nothing is dead, nothing saturates. The network is in perfect health and it is costing about a quarter of the loss, which is exactly the kind of thing no health check is looking for. The diagnosis is in the picture:

pinched.figure(inside=architecture(pinched, ONE))

Figure from cell 15

72 → 2 → 4. Everything the convolutional trunk, the transformer and the recurrent cell computed has to pass through two numbers. The shapes on the layers are what say so, which is why they are written on the box.

The solution

Widen the join to something the three branches can get through. Nothing else changes — not the trunk, not the stack, not the cell.

The healthy architecture and configuration

opened = built(squeeze=16)
fit(opened, "1-healthy")

print(f"squeezed to two     : {scored(pinched):.4f}")
print(f"squeezed to sixteen : {scored(opened):.4f}")
print("flags:", diagnose(store, run="1-healthy") or "nothing tripped")

squeezed to two     : 0.3542
squeezed to sixteen : 0.3079
flags: nothing tripped

opened.figure(inside=architecture(opened, ONE))

Figure from cell 18

Problem 2 — the rate is too high

The problematic architecture and configuration

The architecture from above, trained a thousand times faster.

hot = built()
fit(hot, "2-problem", lr=0.3)
print("loss:", round(scored(hot), 4))

loss: 0.4227

The symptoms

alerts(store, run="2-problem")

2-problem — 35 finding(s)

⚠ VANISHING

audio