09 — Complex architectures, and their health

Generated from a notebook

This page is notebooks/09_complex_architectures_and_health.ipynb, executed, with its outputs. Run it yourself with jupyter lab notebooks/ after pip install 'somatize[viz]'. Interactive Plotly figures show here as the PNG the notebook also stores.

Two realistic scenarios. First a multimodal pipeline (a fork and a fan-in), to see how a non-trivial topology reads. Then a complex model inside a single node, into which we inject four training pathologies documented in the literature — and catch all four with gradient_audit(inside=..., channels=...) at the default thresholds.

import warnings

import plotly.io as pio
import torch
import torch.nn as nn
from IPython.display import SVG

import soma
from soma import ChannelConfig, DifferentiableFilter, Filter, Graph

# Interactivo (zoom/hover/brushing) en Jupyter y VS Code, con
# fallback PNG donde no hay JS (GitHub incluido).
pio.renderers.default = "plotly_mimetype+png"
pio.renderers["png"].scale = 2
pio.renderers["png"].width = 950
warnings.filterwarnings("ignore", message="Full backward hook is firing")

9.1 — A multimodal pipeline

Two preprocessing branches, audio and text, converging on a fusion, a shared backbone and two heads. Evaluating g draws the topology.

class LoadAudio(Filter):
    _cache_version = "nb09-la-v1"
    def forward(self, x, state):
        return [v * 1.0 for v in x]

class Mfcc(Filter):
    _cache_version = "nb09-mfcc-v1"
    def fit(self, x, y=None):
        return {"scale": max(abs(v) for v in x) or 1.0}
    def forward(self, x, state):
        return [v / state["scale"] for v in x]

class LoadText(Filter):
    _cache_version = "nb09-lt-v1"
    def forward(self, x, state):
        return [v + 0.5 for v in x]

class Embed(Filter):
    _cache_version = "nb09-emb-v1"
    def forward(self, x, state):
        return [v * 0.1 for v in x]

class Fuse(Filter):
    _cache_version = "nb09-fuse-v1"
    def forward(self, x, state):
        # Fan-in: x = {node_id: rama}
        return [sum(vals) for vals in zip(*x.values())]

class Backbone(Filter):
    _cache_version = "nb09-bb-v1"
    def fit(self, x, y=None):
        return {"mean": sum(x) / len(x)}
    def forward(self, x, state):
        return [v - state["mean"] for v in x]

class ClfHead(Filter):
    _cache_version = "nb09-clf-v1"
    def forward(self, x, state):
        return [1.0 if v > 0 else 0.0 for v in x]

class RegHead(Filter):
    _cache_version = "nb09-reg-v1"
    def forward(self, x, state):
        return [v * 2.0 for v in x]

g = Graph.somatize(
    (LoadAudio() >> Mfcc() | LoadText() >> Embed())
    >> Fuse()
    >> Backbone()
    >> (ClfHead() | RegHead())
)
g

After running it under track_run, the same diagram is annotated with what happened: per-node duration, cache hits, statuses.

with g.track_run("nb09-pipeline", kind="fit", tags=["demo"]) as run:
    g.fit([3.0, -1.0, 2.0, 5.0, -4.0])

pipe_view = soma.RunView(run.dir)
SVG(pipe_view.to_svg())

pipe_view.plot_gantt()

9.2 — A complex model inside a node, with pathologies

MultiModalEncoder has per-modality branches, an observable fusion point (mix), a deep trunk and a context branch. We inject four classic faults:

Pathology	How we inject it	Detector that must fire
Dead channels (dying ReLU)	four biases of the post-fusion layer set to −6, before the ReLU	DEAD_CHANNELS(≥4)
Leakage between modalities	both branches collapsed onto the same weights, so they encode the same thing	LEAKAGE (CKA between groups > 0.95)
A branch the gradient ignores	ctx enters the forward multiplied by 0 — alive in the forward pass, no gradient	IGNORED_CHANNELS(32) + VANISHING
A vanishing trunk	five tanh layers with weights ×0.30	the staircase in plot_module_flow

class MultiModalEncoder(DifferentiableFilter):
    _cache_version = "nb09-mm-v1"

    def build_module(self, input_shape):
        branches = nn.ModuleDict({
            "audio": nn.Sequential(nn.Linear(16, 16), nn.ReLU()),
            "text": nn.Sequential(nn.Linear(16, 16), nn.ReLU()),
        })
        with torch.no_grad():
            # leakage: branches collapsed onto the same weights
            branches["text"][0].weight.copy_(branches["audio"][0].weight)
            branches["text"][0].bias.copy_(branches["audio"][0].bias)
        post = nn.Sequential(nn.Linear(32, 32), nn.ReLU())
        with torch.no_grad():
            post[0].bias[-4:] = -6.0   # dying ReLU after the fusion
        trunk = []
        for _ in range(5):
            lin = nn.Linear(32, 32)
            with torch.no_grad():
                lin.weight.mul_(0.30)   # contractivo → vanishing
            trunk += [lin, nn.Tanh()]
        return nn.ModuleDict({
            "branches": branches,
            "mix": nn.Identity(),        # observes the fused vector [audio|text]
            "post": post,
            "trunk": nn.Sequential(*trunk),
            "ctx": nn.Linear(32, 32),    # will be multiplied by 0 in the forward
            "head": nn.Linear(32, 8),
        })

    def output_shape(self, input_shape):
        return (8,)

    def forward(self, x, state=None):
        import numpy as np
        x_t = x if isinstance(x, torch.Tensor) else torch.as_tensor(np.asarray(x), dtype=torch.float32)
        self.materialize(tuple(x_t.shape[1:]))
        m = self._module
        a = m["branches"]["audio"](x_t[:, :16])
        t = m["branches"]["text"](x_t[:, 16:])
        mixed = m["mix"](torch.cat([a, t], dim=1))
        deep = m["trunk"](m["post"](mixed))
        out = m["head"](deep + 0.0 * m["ctx"](mixed))
        if self.training:
            return out, {}
        return out.detach().tolist(), {}


torch.manual_seed(0)
g2 = Graph()
enc = MultiModalEncoder()
g2.node("encoder", enc)

audio = torch.randn(96, 16)
x = torch.cat([audio, audio * 0.97 + 0.05 * torch.randn(96, 16)], dim=1)
y = torch.randn(96, 8)
g2.materialize(x)
g2.train()
g2.make_optimizer(torch.optim.Adam, lr=1e-3)
enc   # the model draws itself

MultiModalEncoder (8.2k θ en 5 submódulos)

Auditing inside, with channel groups

inside= selects which submodules to watch (fnmatch patterns), and ChannelConfig.groups declares which channels of mix belong to which modality — the basis of the CKA leakage detector.

cfg = ChannelConfig(
    snapshot_every=2,
    groups={"encoder/mix": {"audio": range(0, 16), "text": range(16, 32)}},
)

with g2.track_run("nb09-health", tags=["demo"]) as run2:
    with g2.gradient_audit(
        inside={"encoder": ["branches.audio", "branches.text", "mix", "post", "ctx", "trunk.*"]},
        channels=cfg,
    ) as audit:
        for epoch in range(6):
            with g2.context() as ctx:
                g2.zero_grad()
                out, _ = g2.forward(x)
                g2.backward(ctx, nn.functional.mse_loss(out, y))
            g2.step(ctx)

for f in audit.report().filters:
    if f.flags:
        print(f"{f.filter_id:26s} {f.flags}")

encoder/mix                ['LEAKAGE']
encoder/post               ['DEAD_CHANNELS(5)']
encoder/ctx                ['VANISHING', 'IGNORED_CHANNELS(32)']

All four pathologies, caught at the default thresholds. Now the same conclusions, visually:

view = soma.RunView(run2.dir)
view.plot_health()   # what failed, where, and at which step

SVG(view.to_svg(node="encoder"))   # the inner architecture, layer by layer

view.plot_module_flow("encoder")   # the trunk's vanishing staircase

view.plot_channels("encoder/mix")  # audio↔texto correlacionados + 4 canales muertos (†)

view.plot_channel_evolution("encoder/mix")   # effective rank and CKA across traininging

And on the outer DAG, the rollup marks the encoder node with the detail of which submodules failed — the graph tells you where to look, the inner views tell you what happened.

parent = [f for f in view.health_flags() if f["node_id"] == "encoder"]
for f in parent:
    print(f"{f['flag']:18s} {f['detail']}")
SVG(view.to_svg())

DEAD_CHANNELS      in: post
IGNORED_CHANNELS   in: ctx
LEAKAGE            in: mix
VANISHING          in: ctx

soma report <run_id> packages every one of these views — the Module flow section included — into a shareable HTML file. The full guide is docs → guides/gradient-audit.md, section Looking inside a node.