Prerequisite Notes for the ReLU Privileged-Basis Problem

This document records:

1. the concepts introduced,
2. the questions attempted,
3. my original solutions,
4. what was correct,
5. what was incorrect or incomplete,
6. the corrected derivations.

Question 1 — Covariance and second moments

• Let be a random vector with mean . We know that, the covariance of (not necessarily independent) two random variables is expressed as: . We can use this notion of covariance of two random variables to construct the notion of covariance of a vector with itself: or, in matrix form, equivalently:

• The diagonal entries are called the coordinate variances

• The Euclidean norm is expressed as

• Trivial Identity:

• Inner and Outer producs: For vectors

  • The standard Euclidean inner product is equivalently expressed as .
  • The Outer product is

• for vectors

• Using the above trivial identity, one can derive another useful identity -

• Useful -

• Linearity of Expectation - For integrable random variables and constants

• If and are independent and the relevant expectations exist, then . Independence is sufficient, but not necessary.

• The second moment of a random variable is simply the expectation

• For two random variables it is simple to see that

• Using the above, it can be seen that independence implies zero covariance - just expand the terms inside and write them elementwise for the entire probability space (use summation or integrals)

0.1 PLE performance gains?

It is not immediately obvious HOW PLE improves performance. Why even bother? What problem do they solve at this model scale? All I’ve been able to find so far on this matter are sloppy medium posts and I don’t trust them. I spent a good while trying to figure it out… honestly, I still don’t know for sure, I think it would make for a great mech-interp paper. Let’s try to work this out anyway!

My rough understanding of transformers is that attention approximates computation and MLP approximates memory. Looking at the architecture diagram, we can see two things:

1. Residual stream gating the PLE lookup - A simple hypothesis I agree with² is that the layer-wise lookup allows the model to store and recall a great deal of information as the lookup is precomputed for every forward pass, instead of going through matmuls in the MLP, thereby increasing capacity without commensurate increases in VRAM occupancy.

2. Mix-LayerNorm (pre + post) - The transformer blocks use post-norm as well as pre-norm. The disadvantage of Post-norm is that it keeps the L2 norm of the hidden state activations approximately constant across layers as there is no identity path, which means signals originating from early layers keep shrinking as they propagate forward, leading to vanishing influence on the output (and, correspondingly, vanishing gradients). This is a well known problem with post-norm, which is why most LLMs these days use pre-norm as it avoids training instability issues with increase in depth reference - blog by Ziming Liu. Since gemma uses both, it might still suffer from the described post-norm limitation. I have a suspicion that if post-norm were not used in the Gemma architecture, these per-layer embeddings would correspondingly not also have been necessary (simply allow information from first layer token embeddings to flow through the residual stream and reach every layer normally).

1 Implementing PLE - from scratch (it works!)

Out of curiosity, I cooked up a small-scale transformer baseline on the shakespeare corpus, modeled as NTP on raw byte pairs (yes, I didn’t use a tokenizer, literally next-byte prediction). Here are some results:

  - prenorm=False, postnorm=False
      - ple_dim=0: 1.601218
      - ple_dim=8: 1.571632
      - ple_dim=32: 1.550698
      - ple_dim=128: 1.622619

  - prenorm=False, postnorm=True
      - ple_dim=0: 1.581143
      - ple_dim=8: 1.595611
      - ple_dim=32: 1.579468
      - ple_dim=128: 1.559271

  - prenorm=True, postnorm=False
      - ple_dim=0: 1.541992
      - ple_dim=8: 1.546094
      - ple_dim=32: 1.537176
      - ple_dim=128: 1.529346

  - prenorm=True, postnorm=True
      - ple_dim=0: 1.541877
      - ple_dim=8: 1.544192
      - ple_dim=32: 1.534553
      - ple_dim=128: 1.528164

  Averaged across all four norm settings
  - ple_dim=0: 1.566558
  - ple_dim=8: 1.564382
  - ple_dim=32: 1.550474
  - ple_dim=128: 1.559850

Ignoring the weird error bars on the rightmost chart for x=8, we can see that the average performance over validation generally improves with increase in PLE dimension which tells us that PLE does indeed improve performance on this simple toy task.

# score for google/gemma-3n-E2B-it, google's official published slice of Gemma3n-E4B-it
{
    "en_text_recognition": 0.6096878412118829,
    "en_text_detection": 0.23172620479653033,
    "en_text_spotting": 0.0,
    "en_element_parsing": 0.04,
    "en_mathematical_calculation": 0.3240026170865559,
    "en_visual_text_understanding": 0.7103983997859018,
    "en_knowledge_reasoning": 0.5074402603391522,
    "English Overall Score": 0.34617933188857475
}

# score for capriciousmarauder/test-submodel-3, which is a slice of Gemma3n-E4B-it
{
    "en_text_recognition": 0.5459304480487074,
    "en_text_detection": 0.21590682007667475,
    "en_text_spotting": 0.0,
    "en_element_parsing": 0.0275,
    "en_mathematical_calculation": 0.2693079603436732,
    "en_visual_text_understanding": 0.676684722674171,
    "en_knowledge_reasoning": 0.4669721878768469,
    "English Overall Score": 0.3146145912885819
}

# score for capriciousmarauder/test-submodel-4, which is a slice of Gemma3n-E4B-it
{
    "en_text_recognition": 0.34489594363685844,
    "en_text_detection": 0.14555663691294746,
    "en_text_spotting": 0.0,
    "en_element_parsing": 0.02,
    "en_mathematical_calculation": 0.20092149837002324,
    "en_visual_text_understanding": 0.5540588220515938,
    "en_knowledge_reasoning": 0.44284559706420745,
    "English Overall Score": 0.2440397854336615
}

Pitfalls

Training individual schemas has a negative transfer to performance on the whole task, as it is theorized that the main properties of schemas lie in the connections between them (which are not developed subsequently).

What is a Transformer?

Transformers are a sequence modelling and transduction architecture heavily leveraging self-attention in order to eliminate the need for convolutions or gating. Advantages over RNNs and LSTMs include:

1. Non recurrent computation - Computation for a layer at a given depth is completely parallelized, while in prior architectures the computation at a given position and depth required computing for all the depths before them.
2. Global dependency modelling - Prior architectures featured a growth in computation required to relate information between positions with larger distance between them. Transformers do away with this limitation by modelling dependencies in a global fashion.
3. Permutation invariance - The transformer architecture fundamentally looks at its inputs as a set of elements¹, as opposed to a sequence. This makes it incredibly powerful and widely applicable.

Transformer Architecture

The transformer architecture features an encoder-decoder setup in the fashion of a typical neural machine translation task. The encoder maps a sequence of symbol representations to a sequence of continuous representations . Given , the decoder then generates the output sequence of symbols in an auto-regressive fashion.

Let us introduce the building blocks of the transformer architecture. ### 1. Scaled dot-product attention The output is calculated as a weighted sum of values of dimension . The weights are softmaxed scores, and the scores are obtained by computing the dot-product of an input query with the key vector associated with that value, each being of dimension .

… which yields the scaled dot-product attention formula:

… which, in practice, is further vectorized:

$$O = \text{Attention}(Q, K, V) = \text{softmax}\bigg(\dfrac{Q\cdot K}{\sqrt{d_k}}\bigg)\cdot V$$

Footnotes

Why

The core of expert performance is an organised set of response strategies than can be employed flexibly to meet task demands.

Subjects trained under rotating emphasis manipulation performed significantly better than control players, and experienced a continued advantage increase after experiment termination.

What

The schema model of executive function observes the formation of schemas, individual and indivisible units of control that can be composed together (many, and often multiple at the same time together) in order to meet the demands of a complex task.

When performing a complicated task, human subjects often have little knowledge of, and do not spontaneously explore, alternative strategies. When led to do so, performance improves.

Performance under high-attention load conditions benefits most when subjects are trained on several, varying levels of subelement priorities.

How

Divide skills into core components, train individually, together and then experiments with their composition. This is aptly called Part-Whole training.

However, instead of training individual schemas in isolation of each other and in absence of the entire goal, only attempt to manipulate and assign the relative emphasis (priority) of selected subcomponents in a rotating manner.

TODO: make the sampling safe for zero division errors

TODO: add stopping condition (max_new_tokens, EOS token) for verifier and draft model as pointed out in both the model forward pass line comments

TODO: implement get_probs_from_logits

TODO: convert logits to probs, add temperature

TODO: add analysis of temperature to acceptance rate

TODO: add SFT simple for both

TODO: add RL with acceptance rate

def remove_last_t_from_kvcache(kvcache, t): for i, (k_layer, v_layer) in enumerate(kvcache): if k_layer is not None: kvcache[i] = ( k_layer[:, :, :-t, :], # remove last num_to_remove keys v_layer[:, :, :-t, :] # remove last num_to_remove values ) ““” 1. Inference loop: - inputs object with keys “input_ids” and “attention_mask” # TODO: code currently works for mask of shape [1, L] but won’t for a 3d [1, L, L] square mask. - past_key_values_draft : (DynamicCache(config=draft_model.config)) - past_key_values_verifier: (DynamicCache(config=draft_model.config)) - k = 10 - max_new_tokens = 100 ““”

initialize generated_ids for the recursion

generated_ids = inputs[“input_ids”]

infer draft model k times

start_position = 0 if past_key_values_draft[0][0] is None else past_key_values_draft[0].shape[-2] draft_cache_position = torch.arange(start_position, inputs.input_ids.shape[1], dtype=torch.int64, device=draft_model.device) draft_logits = torch.empty_like(inputs.input_ids[:, 0:0]) # tensor of shape [N, 0]

for draft_iter in range(k): outputs = draft_model(**inputs, cache_position=draft_cache_position, past_key_values=past_key_values_draft, use_cache=True) # TODO: add stopping condition here.

# sampling - greedy
next_token_ids = outputs.logits[:, -1].argmax(-1)
draft_logits = torch.concat([draft_logits, outputs.logits[:, -1:, :]], dim=-2)

generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
attention_mask = inputs["attention_mask"]
attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}
draft_cache_position = draft_cache_position[-1:] + 1 # add one more position for the next token to be decoded.

run verifier in parallel

inputs = {“input_ids”: generated_ids, “attention_mask”: inputs[“attention_mask”]} start_position = 0 if past_key_values_verifier[0][0] is None else past_key_values_verifier[0].shape[-2] verifier_cache_position = torch.arange(start_position, inputs.input_ids.shape[1], dtype=torch.int64, device=verifier_model.device) outputs = verifier_model(**inputs, cache_position=verifier_cache_position, past_key_values=past_key_values_verifier, use_cache=True) # TODO: add stopping condition here.

determine the number of accepted guesses n

verifier_logits = outputs.logits[:, -(draft_logits.shape[-2]+1):-1, :] verifier_probs = verifier_logits.softmax(-1) draft_probs = draft_logits.softmax(-1) prob_ratios = verifier_probs / draft_probs # TODO: make the sampling safe for zero division errors? prob_ratios_per_token = torch.gather(prob_ratios, dim=-2, index=generated_ids.unsqueeze(dim=-2))

destroy the last (k-n) key-values from draft kvcache

destroy the last (k-n) key-values from verifier kvcache

remove_last_t_from_kvcache(past_key_values_verifier, k-n)

adjust the final token from verifier if needed

if n < k: extra_token_logits = norm * (max(0, draft_logits[:, -1-(k-n), :] - verifier_logits[:, -1-(k-n), :])) else: extra_token_logits = verifier_logits[:, -1:, :] extra_token_logits = extra_token_logits.unsqueeze(1)

sample one extra token from adjusted distribution

sampling - greedy

extra_token = extra_token_logits.argmax(-1)

add the generated ids to inputs object to close the recursion, and adjust attention mask. rebuild attention mask entirely?

generated_ids = torch.concat([generated_ids, extra_token], dim=-1) attention_mask = attention_mask.new_ones((attention_mask.shape[0], generated_ids.shape[-1])) inputs = {“inputs”: generated_ids, “attention_mask”: attention_mask}

if q(x) <= p(x): q(x) if q(x) > p(x): - sample (p(x) / q(x)) * ( q(x) ) + (1 - p(x) / q(x)) * ( norm(max(0, p(x) - q(x))) )

desired probability of selecting x: p(x) actual selection: q(x)

distribution that is being sampled from is:

note: alpha is the normalization constant multiplied with the random variable: alpha * (max(0, p(x) - q(x))) to get a probability distribution.

1. q(x) + alpha

let us suppose there are two groups of events for the decoding of an N-vocabulary.

A = {xk : p(xi) >= q(xi) for i in range(N)} R = {xk : p(xi) < q(xi) for i in range(N)}

now, say that the probability that x ~ q(x) is in A is pA = sum( q(xi) * (1 if p(xi) >= q(xi) else 0) ) and, say that the probability that x ~ p(x) is in R is pR = sum( q(xi) * (1 if p(xi) < q(xi) else 0) )

now, we need to show that the final sampling results in the final token sampling probability being equal to p(x).

probability of a token xt (from set A) being sampled = pA * ( q(xt) / pA ) + pR * rejection probability * probabilty that xt is resampled from adjusted dist. = pA * (q(xt) / pA) + pR * rejection probability * norm * max(0, p(xt) - q(xt)) = q(xt) + pR * rejection probability * norm * (p(xt) - q(xt))

now, we need to compute norm and rejection prob. note: rejection probability in the above is conditioned on sample being from B. so lets compute it.

note: P is not p. P is probability of any event function while p is the verifier distribution on tokens. rejection prob. = P(sampled token is rejected | sampled token is from R) = [P(rejected and from R)] / [P(from R)] = [ sum( q(xi) * (1 - p(xi) / q(xi)) if p(xi) < q(xi) else 0) ] / pR = [ sum(q(xi) - p(xi) if p(xi) < q(xi) else 0) ] / pR

-> probability of a token xt from set A being sampled = q(xt) + [ sum(q(xi) - p(xi) if p(xi) < q(xi) else 0) ] * norm * (p(xt) - q(xt))

so, the normalizing constant norm comes out to be: norm = 1 / [ sum(q(xi) - p(xi) if p(xi) < q(xi) else 0) ]

and, we have:

-> probability of a token xt from set A being sampled = q(xt) + (p(xt) - q(xt)) = p(xt)

similarly, let us compute: probability of a token xt (from set R) being sampled = pR * ( q(xt) / pR) * acceptance probability = q(xt) * ( p(xt) / q(xt) ) = p(xt)

hence, we have proven the correctness of the sampling scheme and also found the normalizing constant.

1 -Algebra

A sigma algebra on a set is a nonempty set of subsets of that is closed under complement, countable union, and countable intersection. In other words:

$$\begin{aligned} &\mathcal{A} \subseteq \mathcal{P}(X) \text{ is a } \sigma\text{-algebra on }X\text{ if:} \\ (a) \text{ } &\phi, X \in \mathcal{A}& \\ (b) \text{ } & A \in \mathcal{A} \implies A^c \coloneqq X \setminus A \in \mathcal{A} \\ (c) \text{ } & A_i \in \mathcal{A}, i \in \mathbb{N} \implies \bigcup^\infty_{i=1} A_i \in \mathcal{A}\end{aligned}$$

is also called a -measurable set. In other words, a measure (which is a kind of generalised notion of volume) can be assigned to a -algebra. This notion is formally established later in Measure

$$\begin{aligned} &\text{A map } \mu \colon \mathcal{A} \rightarrow [0, \infty) \cup\{\infty\} \text{ is called a measure if:} \\ (a) \space & \mu(\phi) =  0\\ (b) \space & \mu\biggl(\bigcup_{i=1}^n A_i \biggr) =  \sum_{i=1}^{n}\sigma(A_i), \space\space\space\space\space\space\space\space \text{with }A_i \cap A_j = \phi, i \neq j \text{ }\forall \text{ countable sequences } A_i \in \mathcal{A}\end{aligned}$$

2 Grading

Award 2 points for a correct and explicit answer, 1 for the correct method with a minor error, and 0 for a missing or materially incorrect answer.

2 Unit 1 — Foundations of vector spaces

Videos: 1, 2, 4, 5.

Vectors in (R^n) are tuples of scalars. A line through a point (p) with direction (v) is

A vector space is closed under vector addition and scalar multiplication and satisfies the vector-space axioms. A subspace (UV) must contain (0) and satisfy

Representative problem: determine whether the solution set of (x+2y-z=0) is a subspace of (R^3), and parameterize it.

3 Unit 2 — Subspaces, inner products, and geometry

Videos: 6, 9, 10.

The Euclidean inner product and norm are

The angle relation is

In (R^3), (uv) is perpendicular to both inputs and has magnitude equal to the parallelogram area:

Representative problem: compute the projection of (u) onto nonzero (v):

4 Unit 3 — Matrices and linear maps

Videos: 11, 13, 14, 16, 17, 19, 20.

A linear map satisfies

For a matrix (A=[a_1  a_n]),

Composition is matrix multiplication:

The matrix of a linear map is determined by the images of the standard basis vectors. Special matrices include diagonal, triangular, symmetric, and identity matrices.

Representative problem: find the matrix of (T(x,y)=(2x-y,x+3y)), then compute (T(1,2)).

5 Unit 4 — Independence, bases, coordinates, and dimension

Videos: 22, 23, 25, 26, 28.

Vectors (v_1,,v_k) are linearly independent when

A basis is an independent spanning set. Coordinates relative to a basis (B=(b_1,,b_n)) satisfy

The exchange lemma explains why all bases of a finite-dimensional space have the same number of vectors: the dimension.

Representative problem: express ((5,1)) in the basis ({(1,1),(2,-1)}).

6 Unit 5 — Inverses and transposes

Videos: 29, 31, 32.

An inverse satisfies

For a (2) matrix,

The transpose is defined by

) and verify it by multiplication.

7 Unit 6 — Range, kernel, rank-nullity, and solving systems

Videos: 34, 35, 37, 38, 40, 42.

The kernel and range are

The rank-nullity theorem is

Row operations preserve the solution set of a system when applied to its augmented matrix. In row-echelon form, pivot variables are determined and free variables parameterize the remaining solutions.

Representative problem: row-reduce an augmented matrix and state whether the system has no solution, one solution, or infinitely many solutions.

8 Unit 7 — Determinants

Videos: 43, 44, 46, 47, 48, 50, 51.

For a (2) matrix,

The determinant is multilinear and alternating, and satisfies

For an (nn) matrix, Laplace expansion along row (i) is

Geometrically, (|(A)|) is volume scaling. A square matrix is invertible exactly when its determinant is nonzero.

) by expansion or elimination.

9 Unit 8 — Eigenvalues and eigenvectors

Videos: 53, 54, 55.

An eigenvector is a nonzero vector satisfying

Eigenvalues solve the characteristic equation

The algebraic multiplicity of an eigenvalue is its multiplicity as a root of (p_A). Eigenvectors for () form the eigenspace

).

10 Quiz

Attempt this closed-notes. Each question is worth 2 points.

1. Parameterize the line through ((1,2)) parallel to ((3,-1)).
2. Prove that the solution set of a homogeneous linear system is a subspace.
3. Compute ((1,2,-1),(2,0,3)) and both vector norms.
4. Compute ((1,2,0)(2,-1,1)).
5. Find the matrix of (T(x,y)=(2x-y,x+3y)).
6. Explain why the columns of a matrix determine its matrix-vector product.
7. Test whether ((1,2)) and ((2,4)) are linearly independent.
8. Find the coordinates of ((5,1)) in the basis ({(1,1),(2,-1)}).
9. Find the inverse of ( ).
10. State and interpret the rank-nullity theorem.
11. Describe how free variables appear in a row-echelon solution.
12. Compute ( ).
13. State the determinant criterion for invertibility.
14. Find the characteristic polynomial of ( ).
15. Find the eigenspaces of that matrix.
16. Explain why eigenvectors belonging to distinct eigenvalues are linearly independent.

2 Grading

Questions 1–10 are computational and worth 2 points each. Questions 11–12 are conceptual and worth 2 points each. Award 1 point for the correct method with an arithmetic or notation error, and 2 for a correct, explicit answer.

2 Unit 1 — Vectors, combinations, and bases

A vector is an ordered list of scalars:

The essential operations are

A linear combination has the form

A basis is a linearly independent spanning set. Coordinates tell us which linear combination of the basis vectors produces a vector.

3 Unit 2 — Linear transformations and matrices

A transformation (T) is linear when

Every linear transformation can be written as

The columns of (A) are the images of the standard basis vectors. Matrix multiplication is composition:

For a (2) matrix,

The matrix is invertible exactly when its determinant is nonzero.

4 Unit 3 — Systems, dot products, and cross products

A linear system is represented by

The dot product is

Perpendicular vectors have dot product zero.

In three dimensions,

The cross product is perpendicular to both inputs, and

5 Unit 4 — Determinants, basis changes, and eigenvalues

Important determinant identities are

Eigenvectors satisfy

Eigenvalues solve the characteristic equation

If (A) has enough independent eigenvectors, then

For a (2) matrix,

6 Unit 5 — Abstract vector spaces

Vectors need not be arrows: polynomials, matrices, functions, and sequences can also form vector spaces. For the polynomial basis

the polynomial (p(x)=4-2x+3x^2) has coordinate vector

The same ideas—span, independence, basis, coordinates, and linear transformations—apply in every vector space.

7 Quiz

Attempt these closed-notes first. Each question is worth 2 points.

1. Find (a,b) such that (a[1,2]^T+b[2,-1]T=[7,1]^T).
2. State the two conditions a set must satisfy to be a basis.
3. Find the matrix for (T(x,y)=(x+2y,3x-y)).
4. Compute
5. Solve (x+y=5) and (2x-y=1).
6. Are ([1,2,3]^T) and ([2,-1,0]^T) perpendicular?
7. Compute ([1,0,1]^TT).
8. Find the determinant of ([]).
9. Find the eigenvalues of ([]).
10. Express (2+5x-x^2) in the basis ({1+x,x,x^2}).
11. Explain why matrix multiplication represents composition of transformations.
12. Explain why the determinant being zero prevents invertibility.

2 Grading rule

Each main question is worth (2) points.

• (0): absent, materially false, circular, or confuses object types;
• (1): correct central idea with a missing definition, assumption, or proof step;
• (2): correct, explicit, and type-consistent.

Equivalent formulations are acceptable.

3 Part I — Measurability and information

4 Part II — Conditional expectation

5 Part III — (L^2) geometry and MMSE

6 Bridge problems

2 Primary reading

Use:

Nima Moshayedi, Lectures on Probability Theory, University of Zurich

https://arxiv.org/pdf/2010.16280

The page numbers below refer to the printed page numbers in the notes.

3 Optional terse reference

For a compact formal reference:

Matija Vidmar, Probability with Measure

https://arxiv.org/pdf/2604.01685

Relevant sections:

• §1.1, measurable sets;
• §1.3, measurable maps and generated sigma-fields;
• §3.3, conditioning.

Do not use this as the first exposition unless you prefer definition-dense notes.

4 How to use the quiz

Do three passes.

1. Closed-notes pass: attempt every question without reading the answer key.
2. Open-notes repair pass: repair incomplete definitions and proofs.
3. Compression pass: rewrite each starred proof in at most ten lines without omitting assumptions.

Each ordinary question is worth (2) points:

• (0): absent, materially false, or circular;
• (1): right idea but missing a required definition, assumption, or proof step;
• (2): correct and explicit.

A starred question is load-bearing. A zero on any starred question means the corresponding level is not yet secure, regardless of total score.

5 Quiz

6 Part I — Measurability and information

7 Part II — Conditional expectation

8 Part III — (L^2) geometry and MMSE

9 Bridge problems toward the privileged-basis proof

These are not prerequisites. They test whether you can now begin the actual problem.

10 Readiness criterion

There are (32) main questions, for (64) points.

You are ready to begin the privileged-basis proof when:

• total score is at least (52/64);
• no starred question receives (0);
• Questions 27–30 can be reproduced without notes;
• you can explain in one paragraph why the problem is fundamentally about the conditional spread of ReLU fibres under the law of (Z).

The answer key is a separate file. Do not open it until completing the closed-notes pass.