The trained models, as programs of the flow signature

The models of SemiInv2, with the same stages, weights and inits. A stage is written with a builder (B): an operation applied to positions returns the positions of its results, so a stage is a function from the positions of its weights and inputs to the positions it returns. Sorts are compared by value when the program is checked (Prog.WF), so the builder needs no typing of its own.

A model is a list of stages over one weight context, the stages' weights side by side (Stage.at). Its program (modelProg) runs the stages' bodies one after another in one builder, each on what the previous one passes on, and is typed by inference as a whole; or (callProg) it calls each stage, defined on its own weights in a tower of definitions (Linked). A stage may factor out its first fac outputs: they are returned by the model, and in stageProgs every later stage takes them as extra (unused) weights. Nothing here is trusted: the printer (RoleFlow.Flow.Print) checks what it prints.

The builder

structure RoleFlow.Flow.BSt (Sg : Sig FS) : Type

The builder's state: the next free position and the lines so far.

• n : ℕ
• lines : Array (Line Sg)

@[reducible, inline]

abbrev RoleFlow.Flow.BG (Sg : Sig FS) (α : Type) : Type

A builder of lines of Sg.

Equations

• RoleFlow.Flow.BG Sg = StateM (RoleFlow.Flow.BSt Sg)

@[reducible, inline]

abbrev RoleFlow.Flow.B (α : Type) : Type

A builder of lines of the flow signature.

Equations

• RoleFlow.Flow.B = RoleFlow.Flow.BG RoleFlow.Flow.flowSig

def RoleFlow.Flow.B.line {Sg : Sig FS} {as bs : List FS} (o : Sg.Op as bs) (xs : List ℕ) : BG Sg (List ℕ)

ys = o(xs), returning the positions ys.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.B.op {as bs : List FS} (o : FOp as bs) (xs : List ℕ) : B (List ℕ)

ys = o(xs) for an operation of the flow signature.

Equations

• RoleFlow.Flow.B.op o xs = RoleFlow.Flow.B.line o xs

def RoleFlow.Flow.B.o0 {b : FS} (o : FOp [] [b]) : B ℕ

Equations

• RoleFlow.Flow.B.o0 o = do let __do_lift ← RoleFlow.Flow.B.op o [] pure (__do_lift.headD 0)

def RoleFlow.Flow.B.o1 {a b : FS} (o : FOp [a] [b]) (x : ℕ) : B ℕ

Equations

• RoleFlow.Flow.B.o1 o x = do let __do_lift ← RoleFlow.Flow.B.op o [x] pure (__do_lift.headD 0)

def RoleFlow.Flow.B.o2 {a₁ a₂ b : FS} (o : FOp [a₁, a₂] [b]) (x y : ℕ) : B ℕ

Equations

• RoleFlow.Flow.B.o2 o x y = do let __do_lift ← RoleFlow.Flow.B.op o [x, y] pure (__do_lift.headD 0)

def RoleFlow.Flow.B.o3 {a₁ a₂ a₃ b : FS} (o : FOp [a₁, a₂, a₃] [b]) (x y z : ℕ) : B ℕ

Equations

• RoleFlow.Flow.B.o3 o x y z = do let __do_lift ← RoleFlow.Flow.B.op o [x, y, z] pure (__do_lift.headD 0)

def RoleFlow.Flow.B.o12 {a b₁ b₂ : FS} (o : FOp [a] [b₁, b₂]) (x : ℕ) : B (ℕ × ℕ)

Equations

• RoleFlow.Flow.B.o12 o x = do let ys ← RoleFlow.Flow.B.op o [x] pure (ys.getD 0 0, ys.getD 1 0)

def RoleFlow.Flow.B.actnorm (n x s b : ℕ) : B ℕ

An actnorm on vec n: x ↦ x ⊙ exp s + b.

Equations

• RoleFlow.Flow.B.actnorm n x s b = do let __do_lift ← RoleFlow.Flow.B.o2 (RoleFlow.Flow.FOp.scaleexp n) x s RoleFlow.Flow.B.o2 (RoleFlow.Flow.FOp.add n) __do_lift b

def RoleFlow.Flow.B.affine (n x s t : ℕ) : B ℕ

The affine coupling step on vec n: x ↦ x ⊙ exp (tanh s) + t.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.he (n : ℕ) : String

How a weight starts: he:n is He-normal with fan-in n, zero:n is zero, small:n is normal with standard deviation 0.05 (n is the fan-in used when tests randomise it; 0 means a bias).

Equations

• RoleFlow.Flow.he n = toString "he:" ++ toString n

def RoleFlow.Flow.small (n : ℕ) : String

Equations

• RoleFlow.Flow.small n = toString "small:" ++ toString n

def RoleFlow.Flow.zero (n : ℕ) : String

Equations

• RoleFlow.Flow.zero n = toString "zero:" ++ toString n

def RoleFlow.Flow.bias : String

Equations

• RoleFlow.Flow.bias = "zero:0"

structure RoleFlow.Flow.Stage : Type

A stage: its weights with their inits, inputs and outputs, and its body from the positions of the weights and the inputs to the returned positions. The first fac outputs are factored out.

• name : String
• wts : List (FS × String)
• ins : List FS
• outs : List FS
• fac : ℕ
• body (w x : List ℕ) : B (List ℕ)

def RoleFlow.Flow.Stage.at (s : Stage) (Θ C : List FS) (off : ℕ) : Prog flowSig

The stage as a program with the weights Θ ++ C, its own at off in Θ.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Stage.prog (s : Stage) : Prog flowSig

The stage as a program on its own weights.

Equations

• s.prog = s.at (List.map (fun (x : RoleFlow.Flow.FS × String) => x.1) s.wts) [] 0

def RoleFlow.Flow.modelWts (stages : List Stage) : List FS

A model's weights: the stages' weights side by side.

Equations

• RoleFlow.Flow.modelWts stages = (List.map (fun (x : RoleFlow.Flow.Stage) => List.map (fun (x : RoleFlow.Flow.FS × String) => x.1) x.wts) stages).flatten

def RoleFlow.Flow.stageProgs (stages : List Stage) : List (Stage × Prog flowSig × List FS)

The stages of a model as programs over the model's weights: each with the weights Θ ++ C, C what the earlier stages factored out, and with that C.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.modelProgOf {Sg : Sig FS} (run : Stage → List ℕ → List ℕ → BG Sg (List ℕ)) (stages : List Stage) : Option (Prog Sg)

A model as one program, each stage run by run (from the positions of its weights and inputs) in one builder, on the outputs the previous stage does not factor out. It returns the factored-out outputs, then the rest of the last stage's.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.modelProg (stages : List Stage) : Option (Prog flowSig)

A model as one program: the stages' bodies one after another.

Equations

• RoleFlow.Flow.modelProg stages = RoleFlow.Flow.modelProgOf (fun (x : RoleFlow.Flow.Stage) => x.body) stages

def RoleFlow.Flow.callProg {Sg : Sig FS} (L : Linked flowSig Sg) (stages : List Stage) : Option (Prog Sg)

A model as one program of calls: each stage a call of the definition of its name in L (Linked.find), on its inputs, then its weights.

Equations

• One or more equations did not get rendered due to their size.

Stages shared by the models

def RoleFlow.Flow.scoreStage (n : ℕ) : Stage

score n: the last stage, so that the forward log-weight is the log-density.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.logitStage (n : ℕ) : Stage

Data in (0, 1)ⁿ, given in its logit chart, as a vector.

Equations

• One or more equations did not get rendered due to their size.

Two dimensions

def RoleFlow.Flow.couplingN (name : String) (m hid : ℕ) : Stage

An affine coupling on vec (m + m), after an actnorm: x₂ ↦ x₂ ⊙ exp (tanh s) + t with (s, t) an MLP of x₁ with two hidden layers. The output is (x₂', x₁).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.linmix (name : String) (m : ℕ) : Stage

A learned invertible linear map on vec (m + m) in block-LU form: a diagonal scale exp logS, then the shears a ↦ a + U b and b ↦ b + L a.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.augment (n k : ℕ) : Stage

Augmentation: the data x : vec n and a hidden standard-normal input u : gauss k, which sample turns into data and which is appended to x.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.toyN (n blocks hid : ℕ) (mix : Bool := false) : List Stage

A flow on ℝⁿ (n even): blocks couplings of width hid, each after a linmix if mix, then score.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.toy2d (blocks hid : ℕ) (mix : Bool) (augPairs : ℕ) : List Stage

A flow on the plane: toyN on the data augmented by 2 * augPairs hidden dimensions.

Equations

• RoleFlow.Flow.toy2d blocks hid mix augPairs = (if augPairs = 0 then [] else [RoleFlow.Flow.augment 2 (2 * augPairs)]) ++ RoleFlow.Flow.toyN (2 * (1 + augPairs)) blocks hid mix

Permutations of images

Permutations as steps (PermS): σ.trans τ (first σ) is τ ++ σ, since a gather reads τ's index through σ's.

def RoleFlow.Flow.squeezeC (cin f s : ℕ) (chan : PermS) (shift : ℕ) : PermS

Squeezing cin channels of (s f) × (s f) by f, then a rotation of the flat vector by shift. Output channel chan c', with c' = ci * f² + f dy + dx, at (r, q) reads channel ci at (f r + dy, f q + dx).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.squeeze (f s : ℕ) (chan : PermS) (shift : ℕ) : PermS

squeezeC with one input channel.

Equations

• RoleFlow.Flow.squeeze f s chan shift = RoleFlow.Flow.squeezeC 1 f s chan shift

def RoleFlow.Flow.chanSwap (a b : ℕ) : PermS

Reordering a × b channels as b × a.

Equations

• RoleFlow.Flow.chanSwap a b = [RoleFlow.Flow.PermStep.tr [b, a] [1, 0]]

def RoleFlow.Flow.pairings2 : List PermS

At f = 2: rows, columns, the checkerboard (swap 1 2, swap 1 3 of the four positions).

Equations

• RoleFlow.Flow.pairings2 = [[], [RoleFlow.Flow.PermStep.tr [2, 2] [1, 0]], [RoleFlow.Flow.PermStep.roll [4] 0 3, RoleFlow.Flow.PermStep.flip [4] 0]]

def RoleFlow.Flow.chanPerm (n : ℕ) (π : PermS) : PermS

A channel permutation of the channels of n pixels each: output channel a reads input channel π a.

Equations

• RoleFlow.Flow.chanPerm n π = List.map (RoleFlow.Flow.PermStep.embed 1 n) π

def RoleFlow.Flow.chanRot (k s : ℕ) : PermS

The rotation of k channels by s: channel a reads a + s.

Equations

• RoleFlow.Flow.chanRot k s = [RoleFlow.Flow.PermStep.roll [k] 0 ((k - s % k) % k)]

MNIST

def RoleFlow.Flow.glowNet (c h w hid : ℕ) (W : ℕ → ℕ) (a : ℕ) : B (ℕ × ℕ)

The coupling network of a Glow coupling on c channels of h × w (hid wide), reading a and the weights at w i₀, …: a 3 × 3 conv, relu, a 1 × 1 conv, relu, a 3 × 3 conv to (s, t).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.resNet (c h w hid : ℕ) (W : ℕ → ℕ) (a : ℕ) : B (ℕ × ℕ)

The deeper coupling network of glowResAt: u₁ = relu (conv₃ₓ₃ a), two residual layers uᵢ₊₁ = uᵢ + conv₃ₓ₃ (relu (conv₁ₓ₁ (relu uᵢ))), then (s, t) = conv₃ₓ₃ (relu u₃).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.glowNetW (c hid : ℕ) : List (FS × String)

The weights of glowNet.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.resNetW (c hid : ℕ) : List (FS × String)

The weights of resNet.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.coupleBlock (name : String) (c h w : ℕ) (sq : PermS) (netW : List (FS × String)) (net : (ℕ → ℕ) → ℕ → B (ℕ × ℕ)) : Stage

A Glow block at any resolution, after an actnorm: c + c channels of h × w, permuted by sq; the halves are coupled in turn (b' = b ⊙ exp sA(a) + tA(a), then a' = a ⊙ exp sB(b') + tB(b')), and sq⁻¹ restores the layout. net is the coupling network and netW its weights.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.splitPrior (c h w hid : ℕ) (sq : PermS) : Stage

Factoring out half of the dimensions, with Glow's learned prior: after sq, the first half x₁ passes on, and the second is scored under N (μ, σ²) with (μ, log σ) a conv net of x₁. The scored half is factored out, so it comes first: the stage returns (z, x₁).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.mnistCoupling (fine coarse hidF hidC : ℕ) : List Stage

The multiscale coupling flow for MNIST with fixed permutations (mnistCoupling of semiInv): logit 784, fine blocks on 4 × 14 × 14, splitPrior, coarse residual blocks on 8 × 7 × 7, score 392.

Equations

• One or more equations did not get rendered due to their size.

Glow

def RoleFlow.Flow.glowStepAt (name : String) (c h w hid : ℕ) (sq π : PermS) : Stage

One step of Glow on c + c channels of h × w, in the image layout (squeezed by sq and back): an actnorm, the invertible 1 × 1 convolution W = P L U (U, L by triu, L between two channel reversals, P the channel permutation π), and one affine coupling.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.mnistGlow (fine coarse hidF hidC : ℕ) : List Stage

Glow for MNIST, in two levels: logit 784, fine steps on 4 × 14 × 14, splitPrior, coarse steps on 8 × 7 × 7, score 392.

Equations

• One or more equations did not get rendered due to their size.

TARFlow on image patches

def RoleFlow.Flow.tfLayerW (h : ℕ) : List (FS × String)

The weights of one transformer layer: attention W_q, W_k, W_v, W_o, then the MLP W₁, b₁, W₂, b₂ of width 4 h.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.tarflowW (T d h L : ℕ) : List (FS × String)

The weights of tarflowBlock: the token embedding, one position embedding per token, the layers, the head to (α, μ), and (α₁, μ₁) for the first token.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.tarflowBlock (name : String) (T d h L : ℕ) (σ : PermS) : Stage

One TARFlow block on T ≥ 1 tokens of dimension d, in the token layout given by σ: zₙ = (xₙ − μₙ) ⊙ exp (−tanh αₙ), with (αₙ, μₙ) read off a causal transformer (width h, L pre-norm layers of single-head attention and a GELU MLP, a key-value cache per layer) at token n − 1, and learned (α₁, μ₁). σ⁻¹ restores the layout.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.patchify (f s : ℕ) : PermS

The image s f × s f cut into s × s patches of f × f, patch-major.

Equations

• RoleFlow.Flow.patchify f s = RoleFlow.Flow.squeeze f s [] 0 ++ RoleFlow.Flow.chanSwap (s * s) (f * f)

def RoleFlow.Flow.revTokens (T d : ℕ) : PermS

The token order reversed (T tokens of d).

Equations

• RoleFlow.Flow.revTokens T d = [RoleFlow.Flow.PermStep.flip [T, d] 0]

def RoleFlow.Flow.roll2 (n k : ℕ) : PermS

The image n × n rolled by k pixels down and right.

Equations

• RoleFlow.Flow.roll2 n k = [RoleFlow.Flow.PermStep.roll [n, n] 0 ((n - k % n) % n), RoleFlow.Flow.PermStep.roll [n, n] 1 ((n - k % n) % n)]

def RoleFlow.Flow.mnistTarflow (p blocks h L : ℕ) : Option (List Stage)

TARFlow for Fashion-MNIST: logit 784, then blocks TARFlow blocks on the (28 / p)² patches of p × p pixels, then score 784. Odd blocks reverse the token order; blocks 2, 3 (mod 4) roll the image by p / 2 pixels first.

Equations

• One or more equations did not get rendered due to their size.