jg/pretty minimal gene exp by james-golden-arcadia · Pull Request #10 · Arcadia-Science/glass-box-umap

james-golden-arcadia · 2026-03-24T22:28:11Z

PR checklist

Describe the changes you've made:
- vmap fast computation of jacobian,
- pca contributions mapped to original features,
- graph filtering improved,
- encoder model improved,
- gene exp embeddings example added,
- mnist example updated with features"
Describe the tests you have conducted to confirm that your changes behave as expected. If you haven't conducted any tests, please explain why.
- pre commit tests

…s for better embeddings

… best ckpt

…t and gene_exp scripts

…al features, graph filtering improved, encoder model improved, gene exp embeddings example added, mnist example updated with features

… pyright conflict

review-notebook-app · 2026-04-03T17:02:07Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

ekiefl

Hey @james-golden-arcadia, nice work overall. Glad to see this is all coming together nicely.

I've branched off your branch and I'm working on "packagizing" your changes (https://github.com/Arcadia-Science/glass-box-umap/tree/ek/pretty-minimal-gene-exp-cleanup). When that's done, we'll merge it into this branch, and then merge this branch into main.

Before I can finish my cleanup, I need some clarity on some things. I've made a few comments below, please help me understand what's going on.

ekiefl · 2026-04-06T22:56:47Z

    X: NDArray[np.floating],
    n_neighbors: int = 10,
-    metric: str = "cosine",
+    metric: str = "euclidean",


Interesting! Flagging this to make sure the change is purposeful and not just experimentation

Should have explained it - euclidean is the default in the original parametric umap and scanpy for gene expression.

But after checking scanpy, cosine distance is preferred for learning the umap embedding directly from the hvg features (with no pca). Maybe that should be automatically set?

ekiefl · 2026-04-07T01:00:11Z

+class ResidualMLPBlock(nn.Module):
+    """
+    Residual MLP block: x -> Linear -> (Norm) -> Act -> Linear -> (Norm) -> +skip
+    Keeps everything bias-free and piecewise-linear (if Norm has no affine).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int | None = None,
+        activation: str = "relu",  # or "leaky_relu"
+        negative_slope: float = 0.01,
+        use_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        hidden_dim = hidden_dim or dim
+
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(hidden_dim, dim, bias=False)
+
+        # Norm that does NOT introduce an additive learned offset
+        self.norm1 = LayerNormDetached(hidden_dim) if use_norm else nn.Identity()
+        self.norm2 = LayerNormDetached(dim) if use_norm else nn.Identity()
+
+        if activation == "leaky_relu":
+            # self.act = nn.LeakyReLU(negative_slope=negative_slope)
+            # a = negative_slope
+
+            self.act = nn.PReLU()  # negative_slope=negative_slope)
+            nonlin = "leaky_relu"
+        else:
+            self.act = nn.ReLU()
+            nonlin = "relu"
+
+        # He/Kaiming init for ReLU-family activations
+        init.kaiming_uniform_(self.fc1.weight, nonlinearity=nonlin)
+        init.kaiming_uniform_(self.fc2.weight, nonlinearity=nonlin)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.fc1(x)
+        h = self.norm1(h)
+        h = self.act(h)
+        h = self.fc2(h)
+        h = self.norm2(h)
+        return x + h
+
+
+class DefaultEncoder(nn.Module):
+    """
+    Bias-free residual MLP encoder for tabular data.
+
+    - Mean-centering/scaling done outside the module (recommended).
+    - All Linear layers use bias=False to preserve your local Jacobian property.
+    - Residual blocks greatly improve trainability for depth.
+    """
+
+    def __init__(
+        self,
+        input_dims: tuple[int, ...],
+        n_components: int = 2,
+        width: int = 128,
+        depth: int = 3,  # number of residual blocks
+        mlp_ratio: float = 2.0,  # inner expansion in each block
+        activation: str = "leaky_relu",  # or "leaky_relu"
+        negative_slope: float = 0.01,
+        use_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        in_dim = math.prod(input_dims)
+        hidden_dim = int(round(width * mlp_ratio))
+
+        self.flatten = nn.Flatten()
+
+        # Input projection (bias-free)
+        self.in_proj = nn.Linear(in_dim, width, bias=False)
+
+        # Init
+        if activation == "leaky_relu":
+            a = negative_slope
+            nonlin = "leaky_relu"
+        else:
+            a = 0.0
+            nonlin = "relu"
+        init.kaiming_uniform_(self.in_proj.weight, a=a, nonlinearity=nonlin)
+
+        # Residual trunk
+        self.blocks = nn.Sequential(
+            *[
+                ResidualMLPBlock(
+                    dim=width,
+                    hidden_dim=hidden_dim,
+                    activation=activation,
+                    negative_slope=negative_slope,
+                    use_norm=use_norm,
+                )
+                for _ in range(depth)
+            ]
+        )


I'm having trouble understanding the intent regarding activations.

The activation="leaky_relu" path creates a PReLU, not a LeakyReLU — was the switch to PReLU intentional?

If PReLU is intentional, should negative_slope be removed since PReLU learns its own slope?

The kaiming init uses negative_slope but PReLU starts at 0.25 — should these match?

Is the ReLU branch (activation="relu") actually used anywhere, or can we simplify to always use PReLU?

The general intent was to switch completely to PReLU (saw better embeddings with it).

Yes, intentional - should have changed it from leaky.

You're right that this should be fixed - we can pass negative slope (set to 0.25) for the PReLU initial slope.

You're absolutely right! Then this negative slope of 0.25 should be passed to kaiming init.

I think we should always use PReLU.

I will clean this up in a commit.

ekiefl · 2026-04-07T01:03:17Z

How does this differ from GlassBoxUMAP.compute_attributions? The code has lots of gene expression language. Does it have to be or is this generalizable computation?

This has the mapping from PCs to gene features, which hadn't been implemented yet. I think compute_attributions works properly with direct gene features.

I think it could be generalized to any expression like data (rows and cols), and probably combined in compute_attributions. Will make a fix.

GlassBoxUMAP.compute_attributions works with both raw features and pca features. If PCA preprocessing was used, then it will project back to raw features.

I think the problem is more that if someone comes with already preprocessed PCA features, then those are the raw features as far as glassbox is concerned

ekiefl · 2026-04-07T01:05:04Z

+    def prelu_to_leaky(self, model: nn.Module) -> nn.Module:
+        """Replace all PReLU modules with LeakyReLU using the learned slopes.
+
+        This is needed for Jacobian computation via ``vmap`` + ``jacrev``, which
+        requires stateless activations.
+
+        Args:
+            model: The model to convert (not modified in-place).
+
+        Returns:
+            A deep copy of the model with PReLU replaced by LeakyReLU.
+        """
+        model = copy.deepcopy(model)
+        for name, module in model.named_modules():
+            if isinstance(module, nn.PReLU):
+                slope = (
+                    module.weight.detach().item()
+                    if module.weight.numel() == 1
+                    else module.weight.detach().mean().item()
+                )
+                parts = name.split(".")
+                parent = model
+                for p in parts[:-1]:
+                    parent = getattr(parent, p)
+                    # parent = getattr(parent, p) if not p.isdigit() else parent[int(p)]
+                # if parts[-1].isdigit():
+                #     parent[int(parts[-1])] = nn.LeakyReLU(negative_slope=slope)
+                # else:
+                setattr(parent, parts[-1], nn.LeakyReLU(negative_slope=slope))
+        return model


Why overwrite PReLU? Why not just construct LeakyReLU from the get-go?

This is a wrinkle of the vmap for instant jacobian. Vmap doesn't work for PReLU, only for leaky relu. But PReLU learns better embeddings. So this hideous chimaera may be a necessary evil.

ekiefl · 2026-05-15T23:09:41Z

Included in #11

james-golden-arcadia added 7 commits February 24, 2026 15:10

minimal implementation of core, graph, loss, models and hyperparamter…

3599c44

…s for better embeddings

train on all graph edges w no repeats, default behavior does not save…

a35cf5a

… best ckpt

Merge main into jg/pretty-minimal and resolve core.py conflicts

8860e87

minor updates and ensuring bias free models, working with pretty mnis…

b691fcc

…t and gene_exp scripts

vmap fast computation of jacobian, pca contributions mapped to origin…

6ea0baa

…al features, graph filtering improved, encoder model improved, gene exp embeddings example added, mnist example updated with features

attribution new file added, type check ignoring vmap due to pytorch /…

aabcee6

… pyright conflict

anndata still optional so ignoring it during type check

ad06104

james-golden-arcadia marked this pull request as ready for review March 25, 2026 20:29

james-golden-arcadia requested a review from ekiefl March 25, 2026 20:29

james-golden-arcadia added 2 commits April 2, 2026 14:01

pre-commit-config remove uv run

f5571b3

vis gene exp nb added to docs/examples, fixed type error in graph

d591f7c

james-golden-arcadia added 2 commits April 3, 2026 10:40

gene exp example improved vis

b5793f3

gene exp example side by side embedding and gene features plot

a917887

ekiefl mentioned this pull request Apr 6, 2026

jg/pretty minimal #9

Closed

1 task

ekiefl reviewed Apr 7, 2026

View reviewed changes

ekiefl closed this May 15, 2026

ekiefl deleted the jg/pretty-minimal-gene-exp branch May 15, 2026 23:09

Uh oh!

Conversation

james-golden-arcadia commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR checklist

Uh oh!

review-notebook-app Bot commented Apr 3, 2026

Uh oh!

ekiefl left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ekiefl commented May 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

james-golden-arcadia commented Mar 24, 2026 •

edited

Loading