KSimplex Geometric Prior for Stable Diffusion: Complete Mathematical Reference

A Cayley-Menger validated pentachoron attention skeleton for rectified flow diffusion.

Author: AbstractPhil
License: MIT
Weights: AbstractPhil/sd15-rectified-geometric-matching

---

Abstract

We present KSimplex, a geometric cross-attention prior that modulates CLIP conditioning before it enters the frozen UNet of a Stable Diffusion 1.5 rectified flow pipeline. The prior operates in a 4-dimensional simplex (pentachoron) coordinate space, using Cayley-Menger determinants to validate geometric configurations and enforce structural coherence. With only 4.8M trainable parameters (0.56% of the 859M frozen backbone), the prior produces measurable improvements in spatial coherence and anatomical correctness after a single epoch on 10,000 synthetic images.

This document catalogs every mathematical formula implemented in the system, from simplex construction through training loss to ODE sampling.

---

1. Architecture

CLIP(prompt) → [B, 77, 768]
    → KSimplexCrossAttentionPrior  (4.8M params, fp32)
        → StackedKSimplexAttention (4 layers)
            → KSimplexAttentionLayer × 4
                → coordinate projection
                → soft vertex assignment
                → deformed template anchoring
                → geometric attention (distance-based)
                → CM-validated value projection
        → residual blend
    → modulated [B, 77, 768]
    → Frozen SD1.5 UNet (859M params, fp16)
    → v_predicted

All original SD1.5 weights load unmodified. Geometric parameters are purely additive.

---

2. Simplex Foundations

2.1 Regular Simplex Template

A k-simplex is the k-dimensional generalization of a triangle. For k=4 (pentachoron), we embed n = k+1 = 5 vertices in ℝ^{d_e} (d_e = 32).

The regular template T ∈ ℝ^{(k+1) × d_e} has unit edge lengths:

$$d_{ij}^2=\mathrm{\parallel }{\mathbf{T}}_i-{\mathbf{T}}_j{\mathrm{\parallel }}^2=1\mathrm{\forall }i\mathrm{\ne }jd\_\{ij\}^2\; =\; \backslash |\backslash mathbf\{T\}\_i\; -\; \backslash mathbf\{T\}\_j\backslash |^2\; =\; 1\; \backslash quad\; \backslash forall\backslash ;\; i\; \backslash neq\; j$$

Constructed via geovocab2.shapes.factory.SimplexFactory(k=4, embed_dim=32, method="regular").

2.2 Edge Count

For n = k+1 vertices:

$$\mathrm{\mid }\mathcal{E}\mathrm{\mid }=\left(\genfrac{}{}{0px}{}{n}{2}\right)=\frac{n(n-1)}{2}|\backslash mathcal\{E\}|\; =\; \backslash binom\{n\}\{2\}\; =\; \backslash frac\{n(n-1)\}\{2\}$$

For k=4: |E| = 5·4/2 = 10 edges.

2.3 Stability Constraint

Empirically validated: the embedding dimension must satisfy:

$$\frac{d_e}{k}\ge 8\backslash frac\{d\_e\}\{k\}\; \backslash geq\; 8$$

For k=4, d_e=32: ratio = 8 ✓. Below this ratio, CM determinants become numerically unstable.

2.4 Pairwise Squared Distances (Vectorized)

Given vertices V ∈ ℝ^{n × d_e}:

$$D_{ij}^2=\sum _{m=1}^{d_e}(V_{im}-V_{jm}{)}^{2}D\_\{ij\}^2\; =\; \backslash sum\_\{m=1\}^\{d\_e\}\; (V\_\{im\}\; -\; V\_\{jm\})^2$$

Computed as: diff = V.unsqueeze(1) - V.unsqueeze(0), D² = (diff²).sum(dim=-1). Upper triangle extracted via torch.triu_indices(n, n, offset=1).

---

3. Cayley-Menger Determinant

3.1 Bordered Distance Matrix

For n = k+1 points with pairwise squared distances d²_{ij}, the CM determinant is the determinant of the (n+1)×(n+1) bordered matrix:

Structure:

• cm[0, 0] = 0 (corner)
• cm[0, 1:] = cm[1:, 0] = 1 (border row/column)
• cm[i+1, j+1] = d²_{ij} for i ≠ j (squared distances)
• cm[i+1, i+1] = 0 (diagonal)

Computed via torch.linalg.det(cm) in float32 (half precision not supported).

3.2 Simplex Validity Condition

A set of distances defines a valid k-simplex in Euclidean space iff:

$$(-1{)}^{k+1}\cdot \text{CM}(d^2)>0(-1)^\{k+1\}\; \backslash cdot\; \backslash text\{CM\}(d^2)\; >\; 0$$

For k=4: (-1)^5 · CM < 0, so CM must be negative.

3.3 Simplex Volume from CM

The squared volume of a k-simplex:

$${\text{vol}}^2({\mathrm{\Delta }}_k)=\frac{(-1{)}^{k+1}}{2^k\cdot (k!{)}^2}\cdot \text{CM}(d^2)\backslash text\{vol\}^2(\backslash Delta\_k)\; =\; \backslash frac\{(-1)^\{k+1\}\}\{2^k\; \backslash cdot\; (k!)^2\}\; \backslash cdot\; \backslash text\{CM\}(d^2)$$

For k=4: denom = 2^4 · (4!)^2 = 16 · 576 = 9216.

$${\text{vol}}^2({\mathrm{\Delta }}_4)=\frac{-1}{9216}\cdot \text{CM}(d^2)\backslash text\{vol\}^2(\backslash Delta\_4)\; =\; \backslash frac\{-1\}\{9216\}\; \backslash cdot\; \backslash text\{CM\}(d^2)$$

---

4. Deformed Template

4.1 Per-Layer Deformation

Each attention layer maintains learnable offsets ΔV ∈ ℝ^{(k+1) × d_e}, initialized as N(0, 0.01):

$${\mathbf{V}}_{\text{deformed}}=\mathbf{T}+\delta \cdot \mathrm{\Delta }\mathbf{V}\backslash mathbf\{V\}\_\{\backslash text\{deformed\}\}\; =\; \backslash mathbf\{T\}\; +\; \backslash delta\; \backslash cdot\; \backslash Delta\backslash mathbf\{V\}$$

where the deformation scale δ is clamped:

$$\delta =\text{clamp}({\delta }_{\text{raw}},0.05,0.5)\backslash delta\; =\; \backslash text\{clamp\}(\backslash delta\_\{\backslash text\{raw\}\},\backslash ;\; 0.05,\backslash ;\; 0.5)$$

Stability zone: δ ∈ [0.15, 0.35] empirically optimal. The clamp bounds allow the optimizer to explore slightly beyond.

4.2 Timestep-Conditioned Deformation (Phase 3)

A small MLP maps the normalized timestep t ∈ [0,1] to per-layer deformation factors:

$$\mathbf{f}(t)=\sigma \text{\hspace{0.17em}⁣}(W_2\cdot \text{SiLU}(W_1\cdot t))\in [0,1{]}^L\backslash mathbf\{f\}(t)\; =\; \backslash sigma\backslash !\backslash left(W\_2\; \backslash cdot\; \backslash text\{SiLU\}(W\_1\; \backslash cdot\; t)\backslash right)\; \backslash in\; [0,\; 1]^\{L\}$$

where W₁ ∈ ℝ^{64×1}, W₂ ∈ ℝ^{L×64}, σ = sigmoid.

The effective deformation scale for layer ℓ at timestep t:

$${\delta }_{\mathrm{\ell }}^{\text{eff}}(t)=\text{clamp}\text{\hspace{0.17em}⁣}({\delta }_{\mathrm{\ell }}\cdot (0.5+f_{\mathrm{\ell }}(t)),0.05,0.5)\backslash delta\_\backslash ell^\{\backslash text\{eff\}\}(t)\; =\; \backslash text\{clamp\}\backslash !\backslash Big(\backslash delta\_\backslash ell\; \backslash cdot\; \backslash big(0.5\; +\; f\_\backslash ell(t)\backslash big),\backslash ;\; 0.05,\backslash ;\; 0.5\backslash Big)$$

The multiplier (0.5 + f) ∈ [0.5, 1.5] can both reduce and amplify the base deformation.

Critical implementation detail: Effective scales are computed transiently — the underlying parameter δ_ℓ is never mutated during forward passes. An earlier version used .data.copy_() which caused training drift.

---

5. Token-to-Simplex Mapping

5.1 Coordinate Projection

Each CLIP token h_i ∈ ℝ^{768} is projected into simplex coordinate space:

$${\mathbf{c}}_i=W_{\text{coord}}\cdot \text{LayerNorm}({\mathbf{h}}_i)\in {\mathbb{R}}^{d_e}\backslash mathbf\{c\}\_i\; =\; W\_\{\backslash text\{coord\}\}\; \backslash cdot\; \backslash text\{LayerNorm\}(\backslash mathbf\{h\}\_i)\; \backslash quad\; \backslash in\; \backslash mathbb\{R\}^\{d\_e\}$$

where W_coord ∈ ℝ^{d_e × 768}.

5.2 Soft Vertex Assignment

Tokens are softly routed to simplex vertices:

$${\mathbf{w}}_i=\text{softmax}\text{\hspace{0.17em}⁣}(W_{\text{vertex}}\cdot \text{LayerNorm}({\mathbf{h}}_i))\in {\mathbb{R}}^{k+1}\backslash mathbf\{w\}\_i\; =\; \backslash text\{softmax\}\backslash !\backslash left(W\_\{\backslash text\{vertex\}\}\; \backslash cdot\; \backslash text\{LayerNorm\}(\backslash mathbf\{h\}\_i)\backslash right)\; \backslash in\; \backslash mathbb\{R\}^\{k+1\}$$

where W_vertex ∈ ℝ^{(k+1) × 768}.

5.3 Template Anchoring

The geometric contribution for each token blends the deformed template vertices:

$${\mathbf{a}}_i=\sum _{j=0}^k{w}_{ij}\cdot {\mathbf{V}}_{\text{deformed},j}={\mathbf{w}}_i^T\cdot {\mathbf{V}}_{\text{deformed}}\backslash mathbf\{a\}\_i\; =\; \backslash sum\_\{j=0\}^\{k\}\; w\_\{ij\}\; \backslash cdot\; \backslash mathbf\{V\}\_\{\backslash text\{deformed\},j\}\; =\; \backslash mathbf\{w\}\_i^T\; \backslash cdot\; \backslash mathbf\{V\}\_\{\backslash text\{deformed\}\}$$

Batched: template_contribution = vertex_weights @ deformed # (B, 77, d_e)

5.4 Final Coordinates

Combined projection + anchor:

$${\widehat{\mathbf{c}}}_i={\mathbf{c}}_i+{\mathbf{a}}_i\backslash hat\{\backslash mathbf\{c\}\}\_i\; =\; \backslash mathbf\{c\}\_i\; +\; \backslash mathbf\{a\}\_i$$

5.5 Unit Sphere Normalization

L2-normalized to bound pairwise squared distances to [0, 4]:

$${\tilde{\mathbf{c}}}_i=\frac{{\widehat{\mathbf{c}}}_i}{\mathrm{\parallel }{\widehat{\mathbf{c}}}_i{\mathrm{\parallel }}_2}\backslash tilde\{\backslash mathbf\{c\}\}\_i\; =\; \backslash frac\{\backslash hat\{\backslash mathbf\{c\}\}\_i\}\{\backslash |\backslash hat\{\backslash mathbf\{c\}\}\_i\backslash |\_2\}$$

This keeps the CM determinant well-conditioned regardless of initialization scale.

---

6. Geometric Attention

6.1 Distance-Based Logits

Unlike standard dot-product attention, KSimplex attention derives from pairwise distances in simplex space:

$$d_{ij}^2=\mathrm{\parallel }{\tilde{\mathbf{c}}}_i-{\tilde{\mathbf{c}}}_j{\mathrm{\parallel }}^{2}d\_\{ij\}^2\; =\; \backslash |\backslash tilde\{\backslash mathbf\{c\}\}\_i\; -\; \backslash tilde\{\backslash mathbf\{c\}\}\_j\backslash |^2$$

$${\text{logits}}_{ij}=-\frac{d_{ij}^2}{\sqrt{d_e}}\backslash text\{logits\}\_\{ij\}\; =\; -\backslash frac\{d\_\{ij\}^2\}\{\backslash sqrt\{d\_e\}\}$$

The 1/√d_e scaling mirrors standard attention's 1/√d_k. Tokens that are geometrically close in simplex space attend more strongly to each other.

6.2 Attention Weights

$${\alpha }_{ij}={\text{softmax}}_j({\text{logits}}_{ij})\backslash alpha\_\{ij\}\; =\; \backslash text\{softmax\}\_j(\backslash text\{logits\}\_\{ij\})$$

Geometric attention is shared across all H heads. Each head has independent value projections.

6.3 Multi-Head Value Projection

$${\mathbf{V}}^{(h)}=\text{LayerNorm}(\mathbf{H})\cdot W_V,\text{reshaped to }(B,T,H,d_h)\backslash mathbf\{V\}^\{(h)\}\; =\; \backslash text\{LayerNorm\}(\backslash mathbf\{H\})\; \backslash cdot\; W\_V,\; \backslash quad\; \backslash text\{reshaped\; to\; \}\; (B,\; T,\; H,\; d\_h)$$

$${\text{out}}^{(h)}=\alpha \cdot {\mathbf{V}}^{(h)}\in {\mathbb{R}}^{B\times T\times d_h}\backslash text\{out\}^\{(h)\}\; =\; \backslash alpha\; \backslash cdot\; \backslash mathbf\{V\}^\{(h)\}\; \backslash quad\; \backslash in\; \backslash mathbb\{R\}^\{B\; \backslash times\; T\; \backslash times\; d\_h\}$$

$$\text{output}=\text{Concat}({\text{out}}^{(1)},\dots ,{\text{out}}^{(H)})\cdot W_O\backslash text\{output\}\; =\; \backslash text\{Concat\}(\backslash text\{out\}^\{(1)\},\; \backslash ldots,\; \backslash text\{out\}^\{(H)\})\; \backslash cdot\; W\_O$$

where d_h = 768/8 = 96.

6.4 Residual Connection

Each layer adds to the residual stream:

$${\mathbf{H}}^{(\mathrm{\ell }+1)}={\mathbf{H}}^{(\mathrm{\ell })}+{\text{output}}^{(\mathrm{\ell })}\backslash mathbf\{H\}^\{(\backslash ell+1)\}\; =\; \backslash mathbf\{H\}^\{(\backslash ell)\}\; +\; \backslash text\{output\}^\{(\backslash ell)\}$$

---

7. Entropy Sharpening (Emergent Property)

7.1 Per-Layer Attention Entropy

$${\mathcal{H}}_{\mathrm{\ell }}=-\sum _i\sum _j{\alpha }_{ij}^{(\mathrm{\ell })}\log \text{\hspace{0.17em}⁣}({\alpha }_{ij}^{(\mathrm{\ell })}+{10}^{-10})\backslash mathcal\{H\}\_\backslash ell\; =\; -\backslash sum\_\{i\}\; \backslash sum\_\{j\}\; \backslash alpha\_\{ij\}^\{(\backslash ell)\}\; \backslash log\backslash !\backslash left(\backslash alpha\_\{ij\}^\{(\backslash ell)\}\; +\; 10^\{-10\}\backslash right)$$

averaged over batch and token dimensions.

7.2 Monotonic Decrease

Validated across classification, language modeling, and now diffusion:

$${\mathcal{H}}_1>{\mathcal{H}}_2>{\mathcal{H}}_3>{\mathcal{H}}_4\backslash mathcal\{H\}\_1\; >\; \backslash mathcal\{H\}\_2\; >\; \backslash mathcal\{H\}\_3\; >\; \backslash mathcal\{H\}\_4$$

This coarse-to-fine sharpening emerges naturally from stacked geometric constraints — no explicit entropy loss is applied. Early layers attend broadly (soft geometry), later layers attend sharply (crystallized structure).

---

8. Residual Blending

8.1 Learnable Blend

The geo_prior output blends with original CLIP conditioning:

$${\mathbf{E}}_{\text{out}}=(1-\beta )\cdot {\mathbf{E}}_{\text{clip}}+\beta \cdot {\mathbf{E}}_{\text{geo}}\backslash mathbf\{E\}\_\{\backslash text\{out\}\}\; =\; (1\; -\; \backslash beta)\; \backslash cdot\; \backslash mathbf\{E\}\_\{\backslash text\{clip\}\}\; +\; \backslash beta\; \backslash cdot\; \backslash mathbf\{E\}\_\{\backslash text\{geo\}\}$$

$$\beta =\sigma ({\beta }_{\text{logit}})\backslash beta\; =\; \backslash sigma(\backslash beta\_\{\backslash text\{logit\}\})$$

Initialized: β_logit = 0.0 → β = 0.5 (equal mix). Sigmoid bounds β ∈ (0, 1).

8.2 Timestep-Conditioned Blend (Optional)

$$\beta (t)=\sigma \text{\hspace{0.17em}⁣}(f_{\beta }(t))\backslash beta(t)\; =\; \backslash sigma\backslash !\backslash left(f\_\backslash beta(t)\backslash right)$$

$$f_{\beta }:{\mathbb{R}}^1\stackrel{W_1\in {\mathbb{R}}^{64\times 1}}{\to }\text{SiLU}\stackrel{W_2\in {\mathbb{R}}^{1\times 64}}{\to }{\mathbb{R}}^{1}f\_\backslash beta:\; \backslash mathbb\{R\}^1\; \backslash xrightarrow\{W\_1\; \backslash in\; \backslash mathbb\{R\}^\{64\; \backslash times\; 1\}\}\; \backslash text\{SiLU\}\; \backslash xrightarrow\{W\_2\; \backslash in\; \backslash mathbb\{R\}^\{1\; \backslash times\; 64\}\}\; \backslash mathbb\{R\}^1$$

Allows the network to learn different blend strengths at different noise levels.

---

9. Geometric Loss Functions

9.1 CM Validity Loss

Penalizes invalid simplex configurations with a hinge loss:

$${\mathcal{L}}_{\text{CM}}=\frac{1}{L}\sum _{\mathrm{\ell }=1}^L\text{ReLU}\text{\hspace{0.17em}⁣}(-(-1{)}^{k+1}\cdot \text{CM}(d_{\mathrm{\ell }}^2)+ϵ)\backslash mathcal\{L\}\_\{\backslash text\{CM\}\}\; =\; \backslash frac\{1\}\{L\}\; \backslash sum\_\{\backslash ell=1\}^\{L\}\; \backslash text\{ReLU\}\backslash !\backslash left(-(-1)^\{k+1\}\; \backslash cdot\; \backslash text\{CM\}(d\_\backslash ell^2)\; +\; \backslash epsilon\backslash right)$$

where ε = 10⁻⁶ is a small margin preventing zero-volume degeneracies. Distances are sampled from the first k+1 tokens of each layer (deterministic sampling for stability).

9.2 Volume Consistency Loss

Prevents all layers from collapsing to identical geometry:

$${\mathcal{L}}_{\text{vol}}=-\text{std}\text{\hspace{0.17em}⁣}(\log \mathrm{\mid }{\text{vol}}_1^2\mathrm{\mid },\log \mathrm{\mid }{\text{vol}}_2^2\mathrm{\mid },\dots ,\log \mathrm{\mid }{\text{vol}}_L^2\mathrm{\mid })\backslash mathcal\{L\}\_\{\backslash text\{vol\}\}\; =\; -\backslash text\{std\}\backslash !\backslash left(\backslash log|\backslash text\{vol\}^2\_1|,\backslash ;\; \backslash log|\backslash text\{vol\}^2\_2|,\backslash ;\; \backslash ldots,\backslash ;\; \backslash log|\backslash text\{vol\}^2\_L|\backslash right)$$

The negative standard deviation of log-volumes rewards spread: each layer should carve a distinct geometric structure, not duplicate.

9.3 Combined Geometric Loss

$${\mathcal{L}}_{\text{geo}}={\lambda }_{\text{CM}}\cdot {\mathcal{L}}_{\text{CM}}+{\lambda }_{\text{vol}}\cdot {\mathcal{L}}_{\text{vol}}\backslash mathcal\{L\}\_\{\backslash text\{geo\}\}\; =\; \backslash lambda\_\{\backslash text\{CM\}\}\; \backslash cdot\; \backslash mathcal\{L\}\_\{\backslash text\{CM\}\}\; +\; \backslash lambda\_\{\backslash text\{vol\}\}\; \backslash cdot\; \backslash mathcal\{L\}\_\{\backslash text\{vol\}\}$$

Defaults: λ_CM = 0.01, λ_vol = 0.005.

9.4 Geometric Loss Warmup

Linear ramp over W steps prevents geometric regularization from dominating early training:

$${\lambda }_{\text{geo}}^{\text{eff}}(s)={\lambda }_{\text{geo}}\cdot \min \text{\hspace{0.17em}⁣}(\frac{s}{W},1)\backslash lambda\_\{\backslash text\{geo\}\}^\{\backslash text\{eff\}\}(s)\; =\; \backslash lambda\_\{\backslash text\{geo\}\}\; \backslash cdot\; \backslash min\backslash !\backslash left(\backslash frac\{s\}\{W\},\backslash ;\; 1\backslash right)$$

Default: W = 200 steps.

---

10. Rectified Flow Matching

10.1 Linear Interpolation

Rectified flow defines a straight path between clean data x₀ and Gaussian noise ε:

$${\mathbf{x}}_t=(1-t){\mathbf{x}}_0+t\mathit{ϵ},t\in [0,1]\backslash mathbf\{x\}\_t\; =\; (1\; -\; t)\backslash ,\backslash mathbf\{x\}\_0\; +\; t\backslash ,\backslash boldsymbol\{\backslash epsilon\},\; \backslash quad\; t\; \backslash in\; [0,\; 1]$$

10.2 Velocity Target

The derivative along this path is constant:

$$\mathbf{v}=\frac{d{\mathbf{x}}_t}{dt}=\mathit{ϵ}-{\mathbf{x}}_0\backslash mathbf\{v\}\; =\; \backslash frac\{d\backslash mathbf\{x\}\_t\}\{dt\}\; =\; \backslash boldsymbol\{\backslash epsilon\}\; -\; \backslash mathbf\{x\}\_0$$

10.3 Shifted Schedule

Shift parameter s > 1 biases timestep distribution toward higher noise:

$$t_s=\frac{s\cdot t}{1+(s-1)\cdot t}t\_s\; =\; \backslash frac\{s\; \backslash cdot\; t\}\{1\; +\; (s\; -\; 1)\; \backslash cdot\; t\}$$

Applied identically during training (timestep sampling) and inference (sigma schedule).

10.4 Logit-Normal Timestep Sampling

Biases samples toward mid-range timesteps where learning signal is strongest:

$$u\sim \mathcal{N}(\mu ,{\sigma }^2),t=\sigma (u)=\frac{1}{1+e^{-u}}u\; \backslash sim\; \backslash mathcal\{N\}(\backslash mu,\; \backslash sigma^2),\; \backslash quad\; t\; =\; \backslash sigma(u)\; =\; \backslash frac\{1\}\{1\; +\; e^\{-u\}\}$$

Default: μ = 0, σ = 1. Then scaled to [t_min, t_max] and shifted.

10.5 Continuous Float Timesteps

Timesteps passed to the UNet are not quantized:

$$\tau =t\times 1000.0\text{(float, not .long())}\backslash tau\; =\; t\; \backslash times\; 1000.0\; \backslash quad\; \backslash text\{(float,\; not\; .long())\}$$

The raw continuous t is additionally passed to the geo_prior as t_continuous to avoid quantization artifacts in the deformation schedule.

---

11. Min-SNR Loss Weighting

11.1 Signal-to-Noise Ratio

For the flow matching parameterization where σ = t:

$$\text{SNR}(t)=\frac{(1-t{)}^2}{t^2+{10}^{-8}}\backslash text\{SNR\}(t)\; =\; \backslash frac\{(1\; -\; t)^2\}\{t^2\; +\; 10^\{-8\}\}$$

11.2 Min-SNR-γ Clamping

Prevents high-noise timesteps (large t, low SNR) from dominating gradients:

$$w_{\text{snr}}(t)=\frac{\min (\text{SNR}(t),\gamma )}{\text{SNR}(t)}w\_\{\backslash text\{snr\}\}(t)\; =\; \backslash frac\{\backslash min(\backslash text\{SNR\}(t),\backslash ;\; \backslash gamma)\}\{\backslash text\{SNR\}(t)\}$$

11.3 Velocity Prediction Adjustment

An additional factor accounts for the velocity parameterization:

$$w_{\text{vel}}(t)=\frac{w_{\text{snr}}(t)}{\text{SNR}(t)+1}w\_\{\backslash text\{vel\}\}(t)\; =\; \backslash frac\{w\_\{\backslash text\{snr\}\}(t)\}\{\backslash text\{SNR\}(t)\; +\; 1\}$$

11.4 Weighted Task Loss

$${\mathcal{L}}_{\text{task}}=\frac{1}{B}\sum _{b=1}^B{w}_{\text{vel}}(t_b)\cdot \mathrm{\parallel }{\mathbf{v}}_{\theta }({\mathbf{x}}_{t_b},t_b)-{\mathbf{v}}_b{\mathrm{\parallel }}^2\backslash mathcal\{L\}\_\{\backslash text\{task\}\}\; =\; \backslash frac\{1\}\{B\}\; \backslash sum\_\{b=1\}^\{B\}\; w\_\{\backslash text\{vel\}\}(t\_b)\; \backslash cdot\; \backslash |\backslash mathbf\{v\}\_\backslash theta(\backslash mathbf\{x\}\_\{t\_b\},\; t\_b)\; -\; \backslash mathbf\{v\}\_b\backslash |^2$$

Default: γ = 5.0.

---

12. Classifier-Free Guidance

12.1 CFG Dropout (Training)

Encoder hidden states are zeroed with probability p_drop to train the unconditional path:

Default: p_drop = 0.1.

12.2 CFG at Inference

$${\mathbf{v}}_{\text{guided}}={\mathbf{v}}_{\text{uncond}}+s_{\text{cfg}}\cdot ({\mathbf{v}}_{\text{cond}}-{\mathbf{v}}_{\text{uncond}})\backslash mathbf\{v\}\_\{\backslash text\{guided\}\}\; =\; \backslash mathbf\{v\}\_\{\backslash text\{uncond\}\}\; +\; s\_\{\backslash text\{cfg\}\}\; \backslash cdot\; (\backslash mathbf\{v\}\_\{\backslash text\{cond\}\}\; -\; \backslash mathbf\{v\}\_\{\backslash text\{uncond\}\})$$

Implemented via batched forward pass: latent_input = cat([x, x]), enc_input = cat([E_cond, E_uncond]), then chunk the output.

---

13. Total Training Loss

$${\mathcal{L}}_{\text{total}}={\mathcal{L}}_{\text{task}}+{\lambda }_{\text{geo}}^{\text{eff}}(s)\cdot {\mathcal{L}}_{\text{geo}}\backslash mathcal\{L\}\_\{\backslash text\{total\}\}\; =\; \backslash mathcal\{L\}\_\{\backslash text\{task\}\}\; +\; \backslash lambda\_\{\backslash text\{geo\}\}^\{\backslash text\{eff\}\}(s)\; \backslash cdot\; \backslash mathcal\{L\}\_\{\backslash text\{geo\}\}$$

Expanding fully:

$${\mathcal{L}}_{\text{total}}=\underset{\text{Min-SNR flow matching}}{\underbrace{\frac{1}{B}\sum _b{w}_{\text{vel}}(t_b)\mathrm{\parallel }{\mathbf{v}}_{\theta }-{\mathbf{v}}_b{\mathrm{\parallel }}^2}}+\underset{\text{warmup}}{\underbrace{{\lambda }_{\text{geo}}\cdot \min \text{\hspace{0.17em}⁣}(\frac{s}{W},1)}}\cdot [\underset{\text{validity}}{\underbrace{{\lambda }_{\text{CM}}\cdot {\mathcal{L}}_{\text{CM}}}}+\underset{\text{diversity}}{\underbrace{{\lambda }_{\text{vol}}\cdot {\mathcal{L}}_{\text{vol}}}}]\backslash mathcal\{L\}\_\{\backslash text\{total\}\}\; =\; \backslash underbrace\{\backslash frac\{1\}\{B\}\backslash sum\_b\; w\_\{\backslash text\{vel\}\}(t\_b)\; \backslash |\backslash mathbf\{v\}\_\backslash theta\; -\; \backslash mathbf\{v\}\_b\backslash |^2\}\_\{\backslash text\{Min-SNR\; flow\; matching\}\}\; +\; \backslash underbrace\{\backslash lambda\_\{\backslash text\{geo\}\}\; \backslash cdot\; \backslash min\backslash !\backslash left(\backslash frac\{s\}\{W\},\; 1\backslash right)\}\_\{\backslash text\{warmup\}\}\; \backslash cdot\; \backslash left[\backslash underbrace\{\backslash lambda\_\{\backslash text\{CM\}\}\; \backslash cdot\; \backslash mathcal\{L\}\_\{\backslash text\{CM\}\}\}\_\{\backslash text\{validity\}\}\; +\; \backslash underbrace\{\backslash lambda\_\{\backslash text\{vol\}\}\; \backslash cdot\; \backslash mathcal\{L\}\_\{\backslash text\{vol\}\}\}\_\{\backslash text\{diversity\}\}\backslash right]$$

Only the geo_prior parameters (4.8M) receive gradients. The UNet backbone (859M) and CLIP encoder are frozen.

---

14. Euler ODE Sampling

14.1 Sigma Schedule

Linear spacing from 1→0 with shift applied:

$${\tau }_i=1-\frac{i}{N},{\sigma }_i=\frac{s\cdot {\tau }_i}{1+(s-1)\cdot {\tau }_i}\backslash tau\_i\; =\; 1\; -\; \backslash frac\{i\}\{N\},\; \backslash quad\; \backslash sigma\_i\; =\; \backslash frac\{s\; \backslash cdot\; \backslash tau\_i\}\{1\; +\; (s-1)\; \backslash cdot\; \backslash tau\_i\}$$

producing N+1 values from σ₀ ≈ 1 to σ_N = 0.

14.2 Euler Integration Step

$${\mathbf{x}}_{i+1}={\mathbf{x}}_i+({\sigma }_{i+1}-{\sigma }_i)\cdot {\mathbf{v}}_{\theta }({\mathbf{x}}_i,{\sigma }_i)\backslash mathbf\{x\}\_\{i+1\}\; =\; \backslash mathbf\{x\}\_i\; +\; (\backslash sigma\_\{i+1\}\; -\; \backslash sigma\_i)\; \backslash cdot\; \backslash mathbf\{v\}\_\backslash theta(\backslash mathbf\{x\}\_i,\; \backslash sigma\_i)$$

Note: dt = σ_{i+1} - σ_i is negative (moving from noise toward clean).

---

15. Learning Rate Schedule

15.1 Linear Warmup

15.2 Cosine Decay

$$p=\frac{s-W_{\text{warmup}}}{\max (S_{\text{total}}-W_{\text{warmup}},1)}p\; =\; \backslash frac\{s\; -\; W\_\{\backslash text\{warmup\}\}\}\{\backslash max(S\_\{\backslash text\{total\}\}\; -\; W\_\{\backslash text\{warmup\}\},\backslash ;\; 1)\}$$

$$\eta (s)={\eta }_{\min }+({\eta }_0-{\eta }_{\min })\cdot \frac{1+\cos (\pi \cdot p)}{2}\backslash eta(s)\; =\; \backslash eta\_\{\backslash min\}\; +\; (\backslash eta\_0\; -\; \backslash eta\_\{\backslash min\})\; \backslash cdot\; \backslash frac\{1\; +\; \backslash cos(\backslash pi\; \backslash cdot\; p)\}\{2\}$$

Defaults: η₀ = 10⁻⁴, η_min = 10⁻⁶, W_warmup = 100.

---

16. VAE Latent Space

16.1 Encoding

$$\mathbf{z}={\text{VAE}}_{\text{enc}}(\mathbf{x})\cdot s_f\backslash mathbf\{z\}\; =\; \backslash text\{VAE\}\_\{\backslash text\{enc\}\}(\backslash mathbf\{x\})\; \backslash cdot\; s\_f$$

where s_f = 0.18215 (SD1.5 scaling factor).

16.2 Decoding

$$\mathbf{x}={\text{VAE}}_{\text{dec}}\text{\hspace{0.17em}⁣}\left(\frac{\mathbf{z}}{s_f}\right)\backslash mathbf\{x\}\; =\; \backslash text\{VAE\}\_\{\backslash text\{dec\}\}\backslash !\backslash left(\backslash frac\{\backslash mathbf\{z\}\}\{s\_f\}\backslash right)$$

Latent shape: (B, 4, H/8, W/8). For 512×512 images: (B, 4, 64, 64).

---

17. Mixed Precision Strategy

The geo_prior runs inside torch.amp.autocast("cuda", enabled=False) with explicit .float() casts. On first inference call, geo_prior auto-casts itself to fp32 (one-time, 4.8M params ≈ 19MB).

---

18. Complete Hyperparameter Table

---

19. Empirical Validation

Prior KSimplex Research

Diffusion Results (1 epoch, 10k ImageNet-synthetic)

• Trainable params: 4,845,725 (0.56% of model)
• Training time: ~7 min on L4 (24GB)
• Batch size: 6
• Fragmented anatomy (split bodies, duplicated heads) → coherent spatial composition
• Circular geometry (bowls, wheels) maintains curvature
• Novel prompts (unseen during training) produce equally coherent outputs
• Lune base capabilities fully preserved — purely additive improvement

---

20. Formula Index

Quick reference of all 42 formulas by section:

---

21. Usage

from sd15_trainer_geo.pipeline import load_pipeline, load_geo_from_hub
from sd15_trainer_geo.generate import generate, show_images

pipe = load_pipeline(device="cuda", dtype=torch.float16)
pipe.unet.load_pretrained(
    "AbstractPhil/tinyflux-experts",
    subfolder="", filename="sd15-flow-lune-unet.safetensors"
)
load_geo_from_hub(pipe, "AbstractPhil/sd15-rectified-geometric-matching")

out = generate(pipe, ["a tabby cat on a windowsill"], shift=2.5, seed=42)
show_images(out)

---

Citation

@misc{abstractphil2026ksimplex,
  title   = {KSimplex Geometric Prior for Stable Diffusion},
  author  = {AbstractPhil},
  year    = {2026},
  url     = {https://huggingface.co/AbstractPhil/sd15-rectified-geometric-matching}
}

License: MIT — Build on it freely.