Home https://shonczinner.github.io/ Shon Czinner's blog about AI, statistics, quantitative finance, economics, and more quarto-1.9.37 Mon, 04 May 2026 00:00:00 GMT The 90-year-old idea behind JEPA models: Canonical Correlation Analysis (CCA) Shon Czinner https://shonczinner.github.io/posts/embedding-prediction/

Introduction

Concepts of correlation and regression may be applied not only to ordinary one-dimensional variates but also to variates of two or more dimensions.

This is the first sentence from the paper “Relations Between Two Sets of Variates” (Hotelling 1936) by statistician and economist Harold Hotelling. This paper introduced Canonical Correlation Analysis (CCA). In modern terminology, “CCA is used to find a common signal among two large matrices” (Bykhovskaya and Gorin 2025).

In JEPA, the objective is the same except the second data matrix happens to be simply a different view of the same data in the first dataset (e.g. via data augmentation or spatial or temporal proximity). One of the recent papers to acknowledge a connection states, “JEPA-based models implicitly perform a non-linear generalization of Canonical Correlation Analysis”. (Huang 2026)

CCA’s connection to JEPA is relevant to Schmidhuber’s debate on who invented JEPA, which is directed at Yann LeCun. Personally, I think Hotelling deserves the credit for the idea of maximizing correlation in embedding space.

Of course, the CCA model has many differences from JEPA.

For one, CCA does not enforce a shared encoder. But the biggest difference is that CCA is linear. Non-linear neural variants of CCA have been researched with the earliest usage of the term “Deep CCA” being (Andrew et al. 2013).

Connecting JEPA models back to its CCA roots is genuinely useful. Another Deep CCA paper (Benton et al. 2017) relaxed the assumption of two sets of variables to an arbitrary number based on a generalization of CCA proposed in 1961 (Horst 1961). Conceivably, JEPAs could be expanded to handle more than two views as well.

CCA vs. JEPA Overview

CCA

Suppose we have zero-mean matrices and .

Let and and so that and .

CCA solves the following maximization problem,

This maximizes the trace of the cross-correlation matrix, while constraining embedding vectors to unit variance and zero covariance.

Similar to the equivalence between maximizing variance and minimizing prediction error in solving PCA, we have a relationship between the trace of the cross-correlation matrix and embedding prediction error,

And due to the whitening constraints,

So maximizing the trace of the cross-correlation under the whitening constraints is equivalent to minimizing the MSE of the embedding representations. Therefore we can write CCA as,

JEPA

Adopting the previous notation, JEPA is constrained to as a result of the joint-embedding. In JEPA, we have the encoder , and predictor .

Let , .

Then we solve,

Note the similarity in the objective function but the lack of whitening constraints. The lack of whitening constraints results in representational and dimensional collapse. For example, a trivial solution to the above problem is .

As discussed in my previous blog post SIGReg (Balestriero and LeCun 2025) fixes this problem. What does it do? It encourages the embeddings and to have an isotropic (i.e. unit variance, uncorrelated) Gaussian distribution. As a result it encourages,

Conclusion

As I mentioned in the introduction, Schmidhuber has debated who invented JEPA and said this about LeCun,

Dr. LeCun’s heavily promoted Joint Embedding Predictive Architecture (JEPA) is the heart of his new company. However, the core ideas are not original to LeCun. Instead, JEPA is essentially identical to our 1992 Predictability Maximization system.

Schmidhuber references Yann LeCun’s response,

JEPA is merely a name for a general concept. The question is, and has always been, how do you make it work (particularly how do you prevent it from collapsing), and how do you make it work at scale with SOTA results on non-toy problems. That’s the hard part. Ideas are a dime a dozen. Making them work is what the community will give you credit for.

Do I agree with LeCun? Yes and no.

Yes, because of course you will get credit for making things work, and ideas are indeed arguably “a dime a dozen”.

No, because the thread of citations is important for progress. If important citations are missed, whether intentionally or not, the correct thing to do is just add them. We’re all only the better for doing so. The connection that JEPA models have to CCA is informative.

My opinion is that JEPA/Predictability Maximization models are architectural enhancements layered on top of CCA. Non-linearity is an enhancement.

Ultimately, these models all have the same objective function introduced by CCA: find the transformations that result in maximal correlation between sets of multidimensional data.

References

Andrew, Galen, Raman Arora, Jeff Bilmes, and Karen Livescu. 2013. “Deep Canonical Correlation Analysis.” International Conference on Machine Learning, 1247–55. https://proceedings.mlr.press/v28/andrew13.html.
Balestriero, Randall, and Yann LeCun. 2025. LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics. https://arxiv.org/abs/2511.08544.
Benton, Adrian, Huda Khayrallah, Biman Gujral, Dee Ann Reisinger, Sheng Zhang, and Raman Arora. 2017. Deep Generalized Canonical Correlation Analysis. https://arxiv.org/abs/1702.02519.
Bykhovskaya, Anna, and Vadim Gorin. 2025. Canonical Correlation Analysis: Review. https://arxiv.org/abs/2411.15625.
Horst, Paul. 1961. Generalized Canonical Correlations and Their Application to Experimental Data. Journal of clinical psychology.
Hotelling, Harold. 1936. “Relations Between Two Sets of Variates.” Biometrika 28 (3/4): 321–77. http://www.jstor.org/stable/2333955.
Huang, Yongchao. 2026. VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models. https://arxiv.org/abs/2601.14354.
]]> ai jepa https://shonczinner.github.io/posts/embedding-prediction/ Mon, 04 May 2026 00:00:00 GMT A Small JEPA Word Embedding Model Shon Czinner https://shonczinner.github.io/posts/small-jepa-language-model/ After my prior blog post about SIGReg, I figured I’d train a small Joint-Embedding Predictive Architecture (JEPA) model to demonstrate it.

The paper “LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels” (Maes et al. 2026) suggested significantly reducing the complexity of JEPA models by removing stop-gradients, and the exponential-moving-average encoder. This was in the context of world models and planning.

In this case, I’m applying JEPA to the task of creating word embeddings. Prior methodologies include Word2vec (Mikolov et al. 2013) which uses a log-linear model and negative sampling, MLP next-word prediction (Bengio et al. 2000), applying CCA to small context windows (Dhillon et al. 2011), and training an autoencoder on small context windows (Shao et al. 2025).

We’ll be training a linear JEPA model with SIGReg on a small shakespeare dataset to show that it learns some informative embeddings. In other words, we’ll train an encoder that turns two words into word embeddings, and train a linear predictor that predicts the second word embedding from the first. It would be easy to extend this methodology to non-linear encoders and use larger contexts than single words.

Overview

First we’ll take our dataset and convert it into tokens. For example,

["to", "be", ",", "or", "not", "to", "be"] -> [1, 2, 3, 4, 5, 1, 2]

Then we create the dataset where we have context/target pairs. So in this case that would look like,

Context 1: [1], Target 1: [2]
Context 2: [2], Target 2: [3]
...
Context 5: [5], Target 5: [1]
Context 6: [1], Target 6: [2]

Then we create the JEPA model which uses the same embedding for the context and target , and then has predictor predict the target from the context. More formally,

Code 
import pandas as pd
import torch
import matplotlib.pyplot as plt
import numpy as np
import requests
import re

Preparing The Data

The dataset is a text file containing some of Shakespeare’s plays.

Code 
txt_url = "https://www.gutenberg.org/cache/epub/100/pg100.txt"
response = requests.get(txt_url)
txt = response.text
txt[:100]
'\ufeffThe Project Gutenberg eBook of The Complete Works of William Shakespeare\r\n    \r\nThis eBook is for t'

To turn this into our dataset, we’ll convert everything to lower-case, split out punctuation, and then split on spaces to get our tokens. We’ll treat everything with frequency below 5 as an unknown token. The dataset consists of a single context word and the target word is simply the next word.

Code 
# Lowercase then put spaces around punctuation and \n and then split on spaces
tokens = re.findall(r"\w+|[^\w\s]", txt.lower(), re.UNICODE)
print(tokens[:10])  # tokens[:10]

min_freq = 25
vocab_freq = pd.Series(tokens).value_counts()
vocab = vocab_freq[vocab_freq >= min_freq].index.tolist()
print("Vocab size: ", len(vocab))
print("First 10 vocab tokens: ", vocab[:10])

# add <unk> token for out-of-vocab words
vocab.insert(0, "<unk>")
token_to_id = {token: idx for idx, token in enumerate(vocab)}
id_to_token = {idx: token for idx, token in enumerate(vocab)}
def encode(tokens):
    return [token_to_id.get(token, token_to_id["<unk>"]) for token in tokens]

def decode(token_ids):
    return [id_to_token.get(token_id, "<unk>") for token_id in token_ids]  

encoded = encode(tokens)
print(encoded[:15]) 

x0 = encoded[:-1]
x1 = encoded[1:]
print("Dataset size: ", len(x0))
pd.DataFrame({"x0": x0, "x1": x1}).head()
['\ufeff', 'the', 'project', 'gutenberg', 'ebook', 'of', 'the', 'complete', 'works', 'of']
Vocab size:  3227
First 10 vocab tokens:  [',', '.', 'the', 'and', '’', 'i', 'to', 'of', 'a', 'you']
[0, 3, 1071, 1098, 0, 8, 3, 0, 1523, 8, 1174, 0, 27, 0, 16]
Dataset size:  1262243
x0 x1
0 0 3
1 3 1071
2 1071 1098
3 1098 0
4 0 8

SIGReg

We use the same SIGReg code as in my prior blog post. This is what makes the embedding space a bit Gaussian, avoiding dimensional collapse, as you’ll see later in Figure 1 which plots the first two embedding dimensions against each other.

Code 
def SIGReg(x, num_slices=256, k=17):
    # x: (N, D) samples
    N, D = x.shape
    device = x.device

    # --- Projection directions ---
    A = torch.randn(D, num_slices, device=device)
    A /= A.norm(dim=0)  # normalize columns → unit directions

    # Project to 1D: shape → (N, num_slices)
    X_proj = x @ A

    # --- Integration points ---
    t = torch.linspace(-5, 5, k, device=device)  # (k,)
    phi_normal = torch.exp(-0.5 * t**2)          # (k,)
    weight = phi_normal                          # Gaussian window

    # Broadcast shapes: (N, M, 1) ⋅ (1, 1, k)
    X_t = X_proj.unsqueeze(-1) * t

    # Empirical characteristic function across samples
    ecf = torch.exp(-1j * X_t).mean(dim=0)  # (M, k)

    # Squared difference
    diff_sq = (ecf - phi_normal).abs()**2  # (M, k)

    # Weighted integration for all projections → shape (M,)
    per_direction_T = torch.trapz(diff_sq * weight, t, dim=1) * N

    # GLOBAL aggregation — MEAN instead of MAX
    T_global = per_direction_T.mean()

    return T_global

Making and Training The Embedding Model

Now we’re ready to train a model. The encoder in this case is just an Embedding module and the predictor is just a Linear module. We use MSE loss comparing the predicted next word embedding versus the actual next word embedding as the objective function.

Code 
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device: ", device)

x0 = torch.tensor(x0, dtype=torch.long).to(device)
x1 = torch.tensor(x1, dtype=torch.long).to(device)

embedding_dim = 10
encoder = torch.nn.Embedding(len(vocab), embedding_dim).to(device)
next_encoding_predictor = torch.nn.Linear(embedding_dim, embedding_dim).to(device)
sigreg_lambda = 0.01

n_epochs = 10
batch_size = 2048
optimizer = torch.optim.Adam(list(encoder.parameters()) + list(next_encoding_predictor.parameters()), lr=0.01)
loss_fn = torch.nn.MSELoss()
for epoch in range(n_epochs):
    total_loss = 0
    for i in range(0, len(x0), batch_size):
        x0_batch = x0[i:i+batch_size]
        x1_batch = x1[i:i+batch_size]
        
        x0_embedded = encoder(x0_batch)
        x1_embedded = encoder(x1_batch)
        
        x1_predicted = next_encoding_predictor(x0_embedded)
        
        loss = loss_fn(x1_predicted, x1_embedded)
        sigreg_loss = SIGReg(x0_embedded)
        loss += sigreg_lambda*sigreg_loss
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    if (epoch+1) % (n_epochs // 10) == 0:
        print(f"Epoch {epoch+1}/{n_epochs}, Loss: {total_loss/len(x0)}")
Using device:  cuda
Epoch 1/10, Loss: 0.00048023610461530306
Epoch 2/10, Loss: 0.0004523864227085027
Epoch 3/10, Loss: 0.0004391360668865251
Epoch 4/10, Loss: 0.0004303184247303124
Epoch 5/10, Loss: 0.00042514875212748543
Epoch 6/10, Loss: 0.00042165358612022394
Epoch 7/10, Loss: 0.0004195104785958759
Epoch 8/10, Loss: 0.0004182606700457514
Epoch 9/10, Loss: 0.0004172710015018738
Epoch 10/10, Loss: 0.00041632126264852835

Visualize

Code 
def visualize_embeddings(encoder, vocab, token_to_id, max_tokens=200):
    """
    Visualize token embeddings.
    
    If embedding_dim == 2 → plot directly.
    If embedding_dim > 2 → plot first 2 dimensions
    
    max_tokens limits how many tokens to plot for readability.
    """

    # optionally limit tokens for clarity
    tokens = vocab[:max_tokens]
    indices = [token_to_id[t] for t in tokens]
    emb_subset = encoder(torch.tensor(indices).to(next(encoder.parameters()).device)).detach().cpu().numpy()
    

    # plot
    plt.figure(figsize=(10, 8))
    plt.scatter(emb_subset[:, 0], emb_subset[:, 1])
    
    # annotate tokens
    for i, token in enumerate(tokens):
        plt.annotate(token, (emb_subset[i, 0], emb_subset[i, 1]), fontsize=8)
    
    plt.title("Token Embeddings Visualization")
    plt.xlabel("Dim 1")
    plt.ylabel("Dim 2")
    plt.grid()
    plt.show()

visualize_embeddings(encoder, vocab, token_to_id)
Figure 1

The visualization above shows only the first two dimensions of the embedding space. We can see several clusters including character names, royal titles, and tokens that follow apostrophes in words like ne’er, ’tis and o’er. This shows that the JEPA model is learning informative embeddings.

Further Embedding Investigation

We can also observe what words are closest in embedding space.

Code 
def neighbour_table_l2(words, n=5):
    device = next(encoder.parameters()).device
    encoder.eval()

    vocab_ids = torch.arange(len(vocab)).to(device)

    with torch.no_grad():
        vocab_emb = encoder(vocab_ids)

    rows = []

    word_order = {w: i for i, w in enumerate(words)}

    for word in words:
        if word not in token_to_id:
            continue

        word_id = token_to_id[word]

        with torch.no_grad():
            query_emb = encoder(torch.tensor([word_id]).to(device))

            x_sq = (query_emb ** 2).sum(dim=1, keepdim=True)
            v_sq = (vocab_emb ** 2).sum(dim=1).unsqueeze(0)
            cross = torch.matmul(query_emb, vocab_emb.T)

            distances = (x_sq + v_sq - 2 * cross).squeeze(0)

            distances[word_id] = float("inf")

            top_ids = torch.topk(-distances, n).indices.tolist()

        for rank, i in enumerate(top_ids, 1):
            rows.append({
                "query": word,
                "query_order": word_order[word],
                "rank": rank,
                "token": id_to_token[i],
                "l2_distance": distances[i].item()
            })

    df = pd.DataFrame(rows)

    # enforce deterministic ordering for display
    df = df.sort_values(["query_order", "rank"])

    pivot_tokens = (
        df.pivot(index="query", columns="rank", values="token")
        .reindex(words)
    ).T

    display(pivot_tokens)

    return df


# ---- run ----
words = ["young", "king", "romeo"]
df = neighbour_table_l2(words, n=5)
query young king romeo
rank
1 old friar malcolm
2 delicate chief hamlet
3 civil tamora wolsey
4 honourable taking lucius
5 troubled perfect viola

It’s encouraging that “king” is near other professions, young is near other adjectives (including its opposite - old), and “romeo” is near other names.

Future Directions

As I mentioned earlier, it would be easy to extend this methodology to non-linear (e.g. MLP, CNN, RNN, Transformer) models and use larger contexts and targets than single words. It’s also possible to play around with other hyperparameters like SIGReg regularizer coefficient, embedding dimension, and hidden dimension and try larger datasets.

Compared to many prior methods for getting word embeddings, this does appear to be less complicated than things like negative sampling (e.g. word2vec).

There’s also recent work on approximations to SIGReg that are likely more computationally efficient with very little downside (Akbar 2026).

References

Akbar, Habibullah. 2026. Weak-SIGReg: Covariance Regularization for Stable Deep Learning. https://arxiv.org/abs/2603.05924.
Bengio, Yoshua, Réjean Ducharme, and Pascal Vincent. 2000. “A Neural Probabilistic Language Model.” In Advances in Neural Information Processing Systems, edited by T. Leen, T. Dietterich, and V. Tresp, vol. 13. MIT Press. https://proceedings.neurips.cc/paper_files/paper/2000/file/728f206c2a01bf572b5940d7d9a8fa4c-Paper.pdf.
Dhillon, Paramveer, Dean P Foster, and Lyle Ungar. 2011. “Multi-View Learning of Word Embeddings via CCA.” In Advances in Neural Information Processing Systems, edited by J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K. Weinberger, vol. 24. Curran Associates, Inc. https://proceedings.neurips.cc/paper_files/paper/2011/file/6c4b761a28b734fe93831e3fb400ce87-Paper.pdf.
Maes, Lucas, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero. 2026. LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. https://arxiv.org/abs/2603.19312.
Mikolov, Tomas, Kai Chen, Greg Corrado, and Jeffrey Dean. 2013. Efficient Estimation of Word Representations in Vector Space. https://arxiv.org/abs/1301.3781.
Shao, Chenze, Darren Li, Fandong Meng, and Jie Zhou. 2025. Continuous Autoregressive Language Models. https://arxiv.org/abs/2510.27688.
]]> ai code jepa https://shonczinner.github.io/posts/small-jepa-language-model/ Thu, 30 Apr 2026 00:00:00 GMT Sketched Isotropic Gaussian Regularization (SIGReg) Explained Shon Czinner https://shonczinner.github.io/posts/sigreg-sketched-isotropic-gaussian-regularization/ The paper “LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics”(Balestriero and LeCun 2025) proposed an interesting method of regularizing latent spaces that I thought I’d explore.

This is a regularization technique to make latent embeddings have isotropic Gaussian distribution - each dimension is encouraged to be uncorrelated and independent, and gaussian distributed. They argue that gaussian latent embeddings are optimal because they are unbiased and lower variance for downstream tasks.

They make latent embeddings have Gaussian distribution using one-dimensional tests of normality in random directions via characteristic functions.

One-Dimension

ECDF Tests

First the paper motivates determining how Gaussian a distribution is with its empirical (i.e. observed) cumulative distribution function (ecdf). Lets look at this in one-dimension.

First we generate two “one sample” univariate datasets - one uniform and one standard gaussian. In this case the target distribution is the standard gaussian. Then we just compare the ecdfs with the theoretical gaussian cdf.

Code 
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
import torch
Code 
# Sample size
n = 100

# Generate samples
X_uniform = np.random.rand(n)*6 - 3  # uniform(-3,3)
X_gaussian = np.random.randn(n)      # Standard Normal

# Define points for theoretical gaussian CDF
x = np.linspace(-3, 3, 1000)
gaussian_cdf = stats.norm.cdf(x)

# ECDF function
def ecdf(data):
    x_sorted = np.sort(data)
    y = np.arange(1, len(x_sorted)+1) / len(x_sorted)
    return x_sorted, y

# Compute ECDFs
x_ecdf_uniform, y_ecdf_uniform = ecdf(X_uniform)
x_ecdf_gaussian, y_ecdf_gaussian = ecdf(X_gaussian)

# Plotting using a loop but keeping samples separate
plt.figure(figsize=(10, 4))

plots = [
    ("Uniform Sample ECDF", x_ecdf_uniform, y_ecdf_uniform),
    ("Gaussian Sample ECDF", x_ecdf_gaussian, y_ecdf_gaussian)
]

for i, (title, x_ecdf, y_ecdf) in enumerate(plots, start=1):
    plt.subplot(1, 2, i)
    plt.plot(x_ecdf, y_ecdf, marker='.', linestyle='none', label=r'ECDF $F_n(x)$')
    plt.plot(x, gaussian_cdf, color='red', label=r'Gaussian CDF $F(x)$')
    plt.title(f'{title} vs Gaussian CDF')
    plt.xlabel('x')
    plt.ylabel('CDF')
    plt.legend()
    plt.grid(True)

plt.tight_layout()
plt.show()
Figure 1

As you’d expect, in Figure 1 the ecdf of the uniform distribution sample deviates from the gaussian cdf having large gaps, but the gaussian sample does not.

We can quantify the “gap” or how close the empirical cdf is to the desired theoretical with the following formula,

where is the empirical cdf, is the theoretical cdf, and is a weighting function. With this is known the cramer-von-mise test statistic and is known as the anderson darling test statistic. We can compute this in scipy.

With some algebraic manipulation there are actually closed-form formulas for these test statistics. However, using the ecdf is still computationally expensive, and non-differentiable because it requires sorting the datasets. We calculate the cramer von mise and anderson darling test statistics below.

Code 
# Tests
cvm_1 = stats.cramervonmises(X_uniform, cdf=stats.norm.cdf)
cvm_2 = stats.cramervonmises(X_gaussian, cdf=stats.norm.cdf)

ad_1 = stats.anderson(X_uniform, dist='norm', method='interpolate')
ad_2 = stats.anderson(X_gaussian, dist='norm', method='interpolate')

alpha = 0.05  # significance level


def cvm_report(name, cvm):
    decision = "REJECT H0 (not Gaussian)" if cvm.pvalue < alpha else "FAIL TO REJECT H0 (Gaussian-consistent)"
    print(f"{name}:")
    print(f"  statistic (T) = {cvm.statistic:.4f}")
    print(f"  p-value   = {cvm.pvalue:.4g}")
    print(f"  decision  = {decision}\n")


def ad_report(name, ad):
    decision = "REJECT H0 (not Gaussian)" if ad.pvalue < alpha else "FAIL TO REJECT H0 (Gaussian-consistent)"
    print(f"{name}:")
    print(f"  statistic (T) = {ad.statistic:.4f}")
    print(f"  p-value   = {ad.pvalue:.4g}")
    print(f"  decision  = {decision}\n")


print("Cramér–von Mises Test Results")
print("-----------------------------")
cvm_report("Uniform sample", cvm_1)
cvm_report("Gaussian sample", cvm_2)


print("Anderson–Darling Test Results")
print("-----------------------------")
ad_report("Uniform sample", ad_1)
ad_report("Gaussian sample", ad_2)
Cramér–von Mises Test Results
-----------------------------
Uniform sample:
  statistic (T) = 1.7954
  p-value   = 3.264e-05
  decision  = REJECT H0 (not Gaussian)

Gaussian sample:
  statistic (T) = 0.3512
  p-value   = 0.09743
  decision  = FAIL TO REJECT H0 (Gaussian-consistent)

Anderson–Darling Test Results
-----------------------------
Uniform sample:
  statistic (T) = 1.6875
  p-value   = 0.01
  decision  = REJECT H0 (not Gaussian)

Gaussian sample:
  statistic (T) = 0.2015
  p-value   = 0.15
  decision  = FAIL TO REJECT H0 (Gaussian-consistent)

As expected, the gaussian hypothesis is rejected for the uniform sample and not rejected for the gaussian sample.

Note: scipy.stats.anderson with dist set to “norm” tests for normality not standard normality and it does so by subtracting the sample mean and dividing by the sample standard deviation of your dataset.

Characteristic Functions

Although a distribution is uniquely defined by its cdf, it is also uniquely defined by its characteristic function. The characteristic function (CF) of a random variable is defined as,

where is the imaginary unit. For a sample we have empirical characteristic function (ECF),

Reminiscent of how the gap in cdf was defined, for CFs we can quantify the gap using,

which is known as the Epps-Pulley test. The weight function is typically which is also the theoretical CF for the Gaussian distribution.

The advantage of the Epps-Pulley test is that it’s differentiable and continuous i.e. if I move one data point a small amount the Epps-Pulley test statistic changes proportionally. In contrast, moving a data point would change the jumps in the ecdf which makes ecdf tests non-differentiable.

This integral can be estimated through the basic trapezoidal Riemann sum.

Code 
def eps_pulley(X, k, plot=False, title=""):
    t = np.linspace(-5, 5, k)

    # theoretical characteristic function of N(0,1)
    phi_gauss = np.exp(-0.5 * t**2)

    # empirical characteristic function
    X_t = np.outer(X, t)
    exp_itx = np.exp(1j * X_t)   # switched to +1j
    phi_emp = np.mean(exp_itx, axis=0)

    # integrated squared error
    err = np.abs(phi_emp - phi_gauss)**2
    T = np.trapezoid(err, t)

    if plot:
        plt.figure(figsize=(8, 5))

        # plot individual exp(itx) curves (real parts)
        for i in range(len(X)):
            plt.plot(
                t,
                exp_itx[i].real,
                color="lightgrey",
                alpha=0.3,
                linewidth=0.8
            )

        # empirical mean
        plt.plot(
            t,
            phi_emp.real,
            color="dimgray",
            linewidth=2,
            label=r"Empirical mean $Re[E[e^{itX}]]=Re[\hat\varphi_n(t)]$"
        )

        # theoretical Gaussian CF
        plt.plot(
            t,
            phi_gauss,
            color="black",
            linestyle="--",
            linewidth=2,
            label=r"Gaussian CF $e^{-t^2/2}=\varphi(t))$"
        )

        plt.xlabel("t")
        plt.ylabel("Real part")
        plt.title(title)
        plt.legend()
        plt.grid(alpha=0.3)
        plt.show()

    return T


k = 1000

T_gaussian = eps_pulley(
    X_gaussian, k,
    plot=True,
    title="Gaussian Sample"
)

T_uniform = eps_pulley(
    X_uniform, k,
    plot=True,
    title="Uniform Sample"
)

print("Epps–Pulley Test Results")
print("-------------------------")
print(f"Gaussian sample: T = {T_gaussian:.6f}")
print(f"Uniform sample:  T = {T_uniform:.6f}")

if T_gaussian < T_uniform:
    print("\nConclusion: Gaussian sample is closer to the standard normal reference distribution.")
else:
    print("\nConclusion: Uniform sample appears closer to the standard normal reference distribution.")

Epps–Pulley Test Results
-------------------------
Gaussian sample: T = 0.064450
Uniform sample:  T = 0.924244

Conclusion: Gaussian sample is closer to the standard normal reference distribution.

The lower T statistic for the gaussian sample indicates less “gap”. I haven’t looked into the limiting distribution of the Epps-Pulley Test Statistic that would be used to calculate p-values.

We can also test the robustness to the number of divisions in the trapezoidal Riemann sum and see it converges quite quickly.

Code 
ks = [x**2 for x in range(2,20,2)]
T_gaussian = [eps_pulley(X_gaussian, k) for k in ks]
T_uniform = [eps_pulley(X_uniform, k) for k in ks]

plots = [
    ("uniform Sample", T_uniform),
    ("gaussian Sample", T_gaussian)
]

for i, (title, T) in enumerate(plots, start=1):
    plt.plot(ks, T, marker='.', label='ECDF')
    plt.title(f'{title} Epps-Pulley Test')
    plt.xlabel('k')
    plt.ylabel('T')
    plt.legend()
    plt.grid(True)
    plt.show()

Gradient Descent

Now we know how to quantify the gap between data and a target distribution with Epps-Pulley. Lets make a neural network output data that fits the desired target distribution. First lets send 1D data through a random, untrained neural network and check its output distribution.

Code 
def plot_output_distribution(model, X, bins=20):
    """
    Plot model output histogram against standard normal PDF.
    
    Args:
        model: PyTorch model
        X: input tensor
        bins: histogram bins
    """
    model.eval()

    with torch.no_grad():
        Y = model(X).squeeze().cpu().numpy()

    plt.figure(figsize=(6,4))

    plt.hist(
        Y,
        bins=bins,
        density=True,
        alpha=0.6,
        label="Network outputs"
    )

    x = np.linspace(Y.min() - 1, Y.max() + 1, 500)

    plt.plot(
        x,
        stats.norm.pdf(x, loc=0, scale=1),
        'r--',
        linewidth=2,
        label=r'$\mathcal{N}(0,1)$ PDF'
    )

    plt.xlabel("Output value")
    plt.ylabel("Density")
    plt.title("Distribution of Network Outputs")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.show()


# -----------------------
# Example setup
# -----------------------
n = 1000
hidden = 16

X = torch.rand(n) * 20 - 10
X = X.unsqueeze(1)

X_train = X[:n//2]
X_test = X[n//2:]

model = torch.nn.Sequential(
    torch.nn.Linear(1, hidden),
    torch.nn.ReLU(),
    torch.nn.Linear(hidden, 1)
)

# Before training
plot_output_distribution(model, X_test)

We can see the output isn’t very Gaussian. Now lets train the model using the Epps-Pulley Regularizer,

Code 
def epps_pulley_1d(X, k):
    t = torch.linspace(-5, 5, k, device=X.device)

    # theoretical CF of N(0,1)
    phi_gaussian = torch.exp(-0.5 * t**2)

    # empirical CF
    X_t = X.unsqueeze(1) * t
    phi_emp = torch.mean(torch.exp(-1j * X_t), dim=0)

    # weighted squared error
    err = phi_gaussian * torch.abs(phi_emp - phi_gaussian)**2

    T = torch.trapz(err, t)

    return T

epochs = 2000

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = lambda x: epps_pulley_1d(x, k=17)

for epoch in range(epochs):
    optimizer.zero_grad()
    loss = loss_fn(model(X_train))
    loss.backward()
    optimizer.step()
    if (epoch+1) % (epochs//10) == 0:
      print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item()}")

plot_output_distribution(model, X_test)
Epoch 200/2000, Loss: 0.0006097470759414136
Epoch 400/2000, Loss: 0.0005753693985752761
Epoch 600/2000, Loss: 0.0005708417156711221
Epoch 800/2000, Loss: 0.0005661727045662701
Epoch 1000/2000, Loss: 0.0005581520381383598
Epoch 1200/2000, Loss: 0.0005494060460478067
Epoch 1400/2000, Loss: 0.0005422477843239903
Epoch 1600/2000, Loss: 0.0005335460300557315
Epoch 1800/2000, Loss: 0.0005124190938659012
Epoch 2000/2000, Loss: 0.0004806756041944027

Now we can see the trained model’s output is more Gaussian.

Multiple Dimensions

Multidimensional Gaussian Regularization

It’s not immediately clear how to apply this in higher dimensions without running into the curse of dimensionality. Naively approximating the integral in higher dimensions would require evaluating it at a number of points that scales exponentially with dimension. In the 1D example we used points. We want to avoid 2D requiring points, 3D requiring , and so on.

The paper’s idea is instead to sample directions in the -dimensional space, project the -dimensional points onto these directions, and then compute the corresponding univariate test statistics.

For a multivariate Gaussian random variable

any linear projection is also Gaussian. And when is a unit vector , the variance remains unchanged:

This means every projection of a standard isotropic Gaussian looks like:

The converse is what makes this useful: by the Cramér–Wold theorem, if every one-dimensional projection of a distribution is standard Gaussian, then the full multivariate distribution must itself be a standard(isotropic) multivariate Gaussian.

This gives us a scalable approximation strategy: instead of testing Gaussianity over the full -dimensional space, we test many random one-dimensional projections.

That’s the core idea behind SIGReg: approximate high-dimensional Gaussian regularization using randomized projections rather than exponentially expensive multidimensional integration.

Here I’ll demonstrate this in two dimensions. First, we’ll inspect scatter plots of 2D data, then choose a random direction on the 2D unit circle, project the points onto that direction, and examine the resulting histograms. The key observation is that these projections reduce the problem back to one dimension, where we can run the same Epps–Pulley test as before.

Code 
n = 400

# Generate samples
X_uniform_2d = np.random.rand(n, 2) * 6 - 3
X_gaussian_2d = np.random.randn(n, 2)

# Random unit direction
v = np.random.randn(2)
v = v / np.linalg.norm(v)

# Projections
proj_uniform = X_uniform_2d @ v
proj_gaussian = X_gaussian_2d @ v


# ----------------------------
# 1. Scatter plot with direction
# ----------------------------
plt.figure(figsize=(6,6))

plt.scatter(
    X_uniform_2d[:,0],
    X_uniform_2d[:,1],
    alpha=0.3,
    label="Uniform"
)

plt.scatter(
    X_gaussian_2d[:,0],
    X_gaussian_2d[:,1],
    alpha=0.3,
    label="Gaussian"
)

plt.arrow(
    0, 0,
    3*v[0], 3*v[1],
    width=0.05,
    color="black",
    length_includes_head=True
)

plt.text(
    3*v[0],
    3*v[1],
    "projection direction",
    fontsize=10
)

plt.axhline(0, color='gray', lw=0.5)
plt.axvline(0, color='gray', lw=0.5)
plt.legend()
plt.title("2D data with random projection direction")
plt.axis("equal")
plt.show()


# ----------------------------
# 2. Projection geometry comparison
# ----------------------------

datasets = [
    ("Uniform Projection Geometry", X_uniform_2d),
    ("Gaussian Projection Geometry", X_gaussian_2d)
]

for title, X_subset in datasets:
    plt.figure(figsize=(6,6))
    plt.scatter(X_subset[:,0], X_subset[:,1])

    # projection line
    t_line = np.linspace(-4, 4, 100)
    plt.plot(
        t_line*v[0],
        t_line*v[1],
        'k--',
        label='projection line'
    )

    # projection segments
    for x in X_subset:
        proj = (x @ v) * v
        plt.plot(
            [x[0], proj[0]],
            [x[1], proj[1]],
            'r-',
            alpha=0.4
        )

    plt.legend()
    plt.title(title)
    plt.axis("equal")
    plt.show()


# ----------------------------
# 3. Histogram comparison
# ----------------------------
x = np.linspace(-4, 4, 400)

plt.figure(figsize=(6,4))

plt.hist(
    proj_uniform,
    bins=30,
    density=True,
    alpha=0.6,
    label="Uniform projected"
)

plt.hist(
    proj_gaussian,
    bins=30,
    density=True,
    alpha=0.6,
    label="Gaussian projected"
)

plt.plot(
    x,
    stats.norm.pdf(x),
    'k--',
    lw=2,
    label='N(0,1) density'
)

plt.legend()
plt.title("1D projections along a random direction")
plt.show()

SIGReg

SIGReg functions the same as above except takes the mean of slices. See the code block below which is similar to the code in the paper.

Code 
def SIGReg(x, num_slices=256, k=17):
    # x: (N, D) samples
    N, D = x.shape
    device = x.device

    # --- Projection directions ---
    A = torch.randn(D, num_slices, device=device)
    A /= A.norm(dim=0)  # normalize columns → unit directions

    # Project to 1D: shape → (N, num_slices)
    X_proj = x @ A

    # --- Integration points ---
    t = torch.linspace(-5, 5, k, device=device)  # (k,)
    phi_normal = torch.exp(-0.5 * t**2)          # (k,)
    weight = phi_normal                          # Gaussian window

    # Broadcast shapes: (N, M, 1) ⋅ (1, 1, k)
    X_t = X_proj.unsqueeze(-1) * t

    # Empirical characteristic function across samples
    ecf = torch.exp(-1j * X_t).mean(dim=0)  # (M, k)

    # Squared difference
    diff_sq = (ecf - phi_normal).abs()**2  # (M, k)

    # Weighted integration for all projections → shape (M,)
    per_direction_T = torch.trapz(diff_sq * weight, t, dim=1) * N

    # GLOBAL aggregation — MEAN instead of MAX
    T_global = per_direction_T.mean()

    return T_global

Now lets make a network’s 10-dimensional output standard gaussian. We’ll plot a histogram in the direction of the first dimension, before and after training on the regularizer.

Code 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

n = 1000
hidden = 20
D = 10
X = torch.rand((n,D))*10-5   # uniform(0,6)

X_train = X[:n//2].to(device)
X_test = X[n//2:].to(device)

model = torch.nn.Sequential(
    torch.nn.Linear(D, hidden),
    torch.nn.ReLU(),
    torch.nn.Linear(hidden, D)
).to(device)

Y = model(X_test)

def plot_output_distribution(model, X, bins=20):
    """
    Plot model output histogram against standard normal PDF.
    
    Args:
        model: PyTorch model
        X: input tensor
        bins: histogram bins
    """
    model.eval()

    with torch.no_grad():
        Y = model(X).squeeze().cpu().numpy()[:,0]

    plt.figure(figsize=(6,4))

    plt.hist(
        Y,
        bins=bins,
        density=True,
        alpha=0.6,
        label="Network outputs"
    )

    x = np.linspace(Y.min() - 1, Y.max() + 1, 500)

    plt.plot(
        x,
        stats.norm.pdf(x, loc=0, scale=1),
        'r--',
        linewidth=2,
        label=r'$\mathcal{N}(0,1)$ PDF'
    )

    plt.xlabel("Output value")
    plt.ylabel("Density")
    plt.title("Distribution of Network Outputs")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.show()


plot_output_distribution(model, X_test)

Again, not very Gaussian. Lets train using SIGReg:

Code 
epochs = 1000

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = lambda x: SIGReg(x)

for epoch in range(epochs):
    optimizer.zero_grad()
    loss = loss_fn(model(X_train))
    loss.backward()
    optimizer.step()
    if (epoch+1) % (epochs//10) == 0:
      print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item()}")

Y=model(X_test)
plot_output_distribution(model, X_test)
Epoch 100/1000, Loss: 2.2469592094421387
Epoch 200/1000, Loss: 0.9632902145385742
Epoch 300/1000, Loss: 0.7503044605255127
Epoch 400/1000, Loss: 0.548032283782959
Epoch 500/1000, Loss: 0.42803165316581726
Epoch 600/1000, Loss: 0.4888719618320465
Epoch 700/1000, Loss: 0.4376848340034485
Epoch 800/1000, Loss: 0.37319087982177734
Epoch 900/1000, Loss: 0.3438933491706848
Epoch 1000/1000, Loss: 0.2580392360687256

And now after training, the output is more Gaussian. So adding this regularizer to your objective function encourages Gaussian latent embeddings.

In a future blog post, I’d like to train a small JEPA model using SIGReg as in “LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels”(Maes et al. 2026).

References

Balestriero, Randall, and Yann LeCun. 2025. LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics. https://arxiv.org/abs/2511.08544.
Maes, Lucas, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero. 2026. LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. https://arxiv.org/abs/2603.19312.
]]> ai code jepa https://shonczinner.github.io/posts/sigreg-sketched-isotropic-gaussian-regularization/ Sun, 26 Apr 2026 00:00:00 GMT Motivation To Write Shon Czinner https://shonczinner.github.io/posts/welcome/ I plan to write about AI, statistics, quantitative finance, economics, and more.

For my first post I’ve decided to write about why I’m writing.

  1. Writing things down provides a reference for when I forget things.
  2. Writing improves understanding.

If other people find the things I write about interesting, that’s a bonus.

Feel free to get in touch if you find any mistakes!

]]> meta https://shonczinner.github.io/posts/welcome/ Sat, 25 Apr 2026 00:00:00 GMT