Skip to content
BranchesTags

Repository files navigation

Autoencoder Lab (AE + GM-VAE)

Autoencoder Lab is a local-first, artifact-driven web app for training and exploring:

  • Supervised Autoencoder (ae)
  • Gaussian Mixture VAE (gmvae)

The backend computes training/projections/neighbors and writes run artifacts to disk. The frontend renders those artifacts (plus live SSE updates) without recomputing ML outputs.

Current functionality

  • Run creation from UI with configurable AE/GM-VAE hyperparameters
  • Live run monitoring via SSE (run.*, train.*, artifact.created, projection.progress)
  • Epoch playback for metrics and reconstructions
  • Latent Explorer route: /runs/[id]/latent
  • Projection methods:
    • PCA 2D (default per epoch)
    • PCA 3D (per epoch + latest; generated during training for new runs, backfillable for older runs)
    • UMAP 2D (on-demand enable/backfill + future epochs)
    • UMAP 3D (on-demand/backfill + future epochs, same UMAP settings with n_components=3)
  • 2D and 3D latent rendering with epoch/method/dimension controls
  • Category legend toggles (show/hide labels) in 2D and 3D
  • Selectable points with dataset sample preview (image + dataset index + label)
  • Per-run persisted point-size controls in 2D and 3D for:
    • sample points
    • centroids
    • GM-VAE cluster centers
  • Latent Explorer UMAP controls (same behavior as run page):
    • UMAP quality selector (Standard 1024, High 2048)
    • compute UMAP directly from latent explorer when UMAP artifacts are missing
  • Sampling panel (GM-VAE, artifact-driven):
    • posterior-local sampling for selected point (q(z|x_i) decode grids)
    • component-conditional prior sampling (p(z|c) decode grids)
    • deterministic seed by default, optional randomized seed
    • cached artifact reuse and sample debug readout (path, seed, compute_ms)
  • Neighbors / Extent panel (PAIR-style):
    • Focus modes: Selected, Random 25, Random 100, All
    • k slider (2..50)
    • extent slider (0..100%)
    • edge modes: Gated by extent (PAIR) and Always show top-k
    • optional radii and edge toggles
    • debug panel with scale/radius/alpha/monotonicity counters
  • 2D overlay modes in Latent Explorer:
    • Off
    • Density (GM-VAE only): Posterior/Prior source, opacity control, density debug stats
    • Distortion: metric-ratio heatmap (mean d_proj / mean d_latent over kNN), opacity control, distortion debug stats
    • Supplemental legends in 2D:
      • centroid and cluster-center marker legend (shown only when those layers are enabled)
      • density/distortion gradient legend (shown only for the active overlay mode)
    • UMAP density guidance note shown in UI to explain nonlinear embedding distortion
  • Mixture Health panel (GM-VAE):
    • reads empirical prior stats from artifacts
    • shows Neff, Neff/K, pmax, entropy, and collapse-status badge
    • student interpretation line (Healthy / Moderate collapse / Severe collapse)
  • On-demand backfill for legacy runs:
    • UMAP projections
    • PCA3 projections
    • density artifacts (pca or umap)
    • distortion artifacts (pca or umap)
    • neighbors artifacts (pca or umap)
  • Artifact browser for each run
  • UI-state bridge and Copy context for AI on run detail, latent explorer, and network explorer
  • Embedded tutor panel with deterministic fallback, optional LLM responses, session history, and feedback (replacement in progress; see below)
  • Run cancellation and resumable/reloadable run history
  • Insights artifacts and UI panel on run detail
  • GM-VAE anti-collapse training options in run config:
    • categorical KL annealing (categorical_kl_anneal*)
    • entropy bonus schedule (entropy_bonus_*)
    • emitted metrics include cat_kl_weight, entropy_bonus_weight, qc_entropy_mean, qc_neff, qc_pmax, qc_argmax_counts
  • Network Architecture / Feature Map Explorer (/runs/[id]/network, artifact-driven):
    • architecture diagram and node details from network/architecture*.json
    • single-image forward-pass artifacts by epoch/index/layer (network/forward/...)
    • channel inspector artifacts (network/inspect/...) with:
      • channel map, heat overlay, fixed overlay fallback
      • optional (x,y) probe and approximate receptive-field mapping back to input (rf)
      • RF spotlight (darkened outside region), hover zoom crop, and probe marker on channel map
      • single-convolution-step explainer for conv1-like layers (C_in=1, 3x3) with patch/kernel/product math and model consistency check
    • latent perturb + decode artifacts (network/perturb/...) for decodable vectors:
      • AE: z
      • GM-VAE: mu, z_sample
    • GM-VAE forward reconstruction sample controls:
      • lock sample / resample (sample_id)
      • forward sampling metadata artifact (forward_meta.json)
    • kernel artifacts for selected conv channels (network/kernels/...):
      • conv1 direct 3x3 kernel display (C_in=1)
      • deeper conv aggregated kernel (L2 over C_in) + top input-channel kernel slices
    • top activating examples for encoder spatial layer+channel (network/topk/...), with clickable thumbnails that set panel dataset index
      • score metric is backend-selected (max by default in UI; mean supported by API)
    • compare mode:
      • synchronized layer/channel/probe/opacity controls across both panels
      • shared channel-level panels (kernel summary + top activating examples) shown once in full width

Architecture

  • Backend: FastAPI (backend.app.main:app) + PyTorch
  • Frontend: Next.js App Router + TypeScript
  • Transport: SSE for live events; filesystem artifacts for deterministic reloads
  • Storage root: runs/<run_id>/...

Canonical contracts live in:

  • AGENTS.md (operating rules)
  • SPEC.md (architecture/artifact semantics)
  • SPEC-mcp-tutor.md (implemented bridge/copy contracts and planned MCP replacement)
  • PLAN-mcp-tutor.md (implementation progress and remaining acceptance gates)
  • http://localhost:8000/openapi.json (authoritative REST API schema)

External AI context and MCP replacement status

The UI-state bridge, context aggregation, and clipboard workflow are implemented. The MCP server is not implemented: there is no /mcp endpoint or client configuration yet. The embedded tutor remains available while its replacement is built and evaluated.

To ask a question in your own AI client:

  1. Open a run's detail, latent, or network page and select the view you want to discuss.
  2. In Context for your own AI, click Copy context for AI.
  3. Paste the text into your AI conversation with your question. Attach screenshots when asking about visual output; the bundle suggests which panels to capture.
  4. Copy again after changing the view when you need to share updated context.

The bundle includes the local page's tracked selections and displayed artifact paths, run highlights, the last five metric rows, latest mixture health when available, and help filtered to the page and model type. It does not attach image bytes or establish a live connection to the AI client. No backend LLM key is needed for this workflow. Clipboard and context-fetch errors appear in the card; live-sync errors are shown separately and do not block copying local state.

The bridge records instrumented UI state, not a screen capture or every UI setting. Its snapshots live in backend memory, are lost on restart, and include age information on reads. Run one backend worker; multiple workers do not share this state. The run-list and run-creation pages do not publish snapshots.

Remaining work is the MCP transport/tools and client setup, comparative evaluation, gated removal of the embedded tutor, and final documentation updates. The planned client acceptance gate covers Claude Desktop and Claude Code; ChatGPT compatibility has not been validated. See the implementation plan.

Repo layout

  • backend/ FastAPI routes, run manager, training runner, projections/backfills, artifacts, tests
  • frontend/ Next.js UI (/, /runs, /runs/[id], /runs/[id]/latent, /runs/[id]/network)
  • backend/app/models/ PyTorch model definitions
  • runs/ generated artifacts (gitignored)

Run artifact structure

Each run is persisted under:

runs/<run_id>/
  config.json
  events.jsonl
  summary.json
  metrics/
  checkpoints/
  samples/
  projections/
  neighbors/
  density/
  distortion/
  mixture/
  insights/
  network/
  tutor/sessions/

Directories appear as their features are used. Metrics are appended to metrics/epoch_metrics.jsonl; tutor/sessions/ holds embedded-tutor transcripts. UI-state bridge snapshots are not persisted in this tree.

Projection artifact naming:

  • PCA 2D:
    • projections/latent_2d_epoch_<N>.json
    • projections/latent_2d_latest.json
  • PCA 3D:
    • projections/latent_3d_epoch_<N>_pca.json
    • projections/latent_3d_latest_pca.json
  • UMAP 2D:
    • projections/latent_2d_epoch_<N>_umap.json
    • projections/latent_2d_latest_umap.json
  • UMAP 3D:
    • projections/latent_3d_epoch_<N>_umap.json
    • projections/latent_3d_latest_umap.json

Neighbors artifact naming (k_max=50, latent-space Euclidean kNN):

  • PCA:
    • neighbors/knn_pca_epoch_<N>.json
    • neighbors/knn_pca_latest.json
  • UMAP:
    • neighbors/knn_umap_epoch_<N>.json
    • neighbors/knn_umap_latest.json

Density artifact naming (2D only):

  • PCA:
    • density/posterior_2d_epoch_<N>_pca.json
    • density/posterior_2d_latest_pca.json
    • density/prior_2d_epoch_<N>_pca.json
    • density/prior_2d_latest_pca.json
    • density/prior_empirical_2d_epoch_<N>_pca.json
    • density/prior_empirical_2d_latest_pca.json
  • UMAP:
    • density/posterior_2d_epoch_<N>_umap.json
    • density/posterior_2d_latest_umap.json
    • density/prior_2d_epoch_<N>_umap.json
    • density/prior_2d_latest_umap.json
    • density/prior_empirical_2d_epoch_<N>_umap.json
    • density/prior_empirical_2d_latest_umap.json

Sampling artifact naming (GM-VAE, on demand):

  • Posterior-local:
    • samples/posterior_local_epoch_<N>_idx_<i>_m_<M>.png
    • samples/posterior_local_latest_idx_<i>_m_<M>.png
    • samples/posterior_local_epoch_<N>_idx_<i>_m_<M>.json
  • Prior-component:
    • samples/prior_component_epoch_<N>_c_<c>_m_<M>.png
    • samples/prior_component_latest_c_<c>_m_<M>.png
    • samples/prior_component_epoch_<N>_c_<c>_m_<M>.json

Metric distortion artifact naming (2D only):

  • PCA:
    • distortion/metric_ratio_2d_epoch_<N>_pca.json
    • distortion/metric_ratio_2d_latest_pca.json
  • UMAP:
    • distortion/metric_ratio_2d_epoch_<N>_umap.json
    • distortion/metric_ratio_2d_latest_umap.json

Network explorer artifact naming (on demand):

  • Forward pass:
    • network/forward/epoch_<N>/idx_<i>/input.png
    • network/forward/epoch_<N>/idx_<i>/recon.png
    • network/forward/epoch_<N>/idx_<i>/layers.json
    • network/forward/epoch_<N>/idx_<i>/latent.json
    • network/forward/epoch_<N>/idx_<i>/forward_meta.json
  • Layer summary/image:
    • network/forward/epoch_<N>/idx_<i>/layer_<layer_id>/summary.json
    • network/forward/epoch_<N>/idx_<i>/layer_<layer_id>/grid.png (spatial)
    • network/forward/epoch_<N>/idx_<i>/layer_<layer_id>/vector.png (vector)
  • Channel inspector:
    • network/inspect/epoch_<N>/idx_<i>/layer_<layer_id>/ch_<c>/inspect.json
    • network/inspect/epoch_<N>/idx_<i>/layer_<layer_id>/ch_<c>/channel.png
    • network/inspect/epoch_<N>/idx_<i>/layer_<layer_id>/ch_<c>/heat_rgba.png
    • network/inspect/epoch_<N>/idx_<i>/layer_<layer_id>/ch_<c>/overlay.png
  • Perturb + decode:
    • network/perturb/epoch_<N>/idx_<i>/vector_<key>/dim_<d>/delta_<tag>/recon.png
    • network/perturb/epoch_<N>/idx_<i>/vector_<key>/dim_<d>/delta_<tag>/meta.json
  • Top activating examples:
    • network/topk/epoch_<N>/layer_<layer_id>/ch_<c>/metric_<metric>/subset_<n>/k_<k>/topk.json
    • network/topk/epoch_<N>/layer_<layer_id>/ch_<c>/metric_<metric>/subset_<n>/k_<k>/input_<dataset_idx>.png
  • Kernel artifacts:
    • network/kernels/epoch_<N>/layer_<layer_id>/out_<c>/kernel.json (conv1)
    • network/kernels/epoch_<N>/layer_<layer_id>/out_<c>/kernel.png (conv1 optional)
    • network/kernels/epoch_<N>/layer_<layer_id>/out_<c>/kernel_agg.json (C_in>1)
    • network/kernels/epoch_<N>/layer_<layer_id>/out_<c>/kernel_agg.png (optional)
    • network/kernels/epoch_<N>/layer_<layer_id>/out_<c>/top_in_channels.json (C_in>1)
    • network/kernels/epoch_<N>/layer_<layer_id>/out_<c>/kernel_in_<in_channel>.png (optional)
  • Single convolution step:
    • network/conv_step/epoch_<N>/idx_<i>/layer_<layer_id>/out_<c>/x_<x>_y_<y>/conv_step.json

Local development

Prerequisites

  • Python 3.13
  • uv
  • Node.js + npm

1) Start backend

From repo root:

uv sync
uv run uvicorn backend.app.main:app --reload --host 127.0.0.1 --port 8000

The supported deployment is one local user with one backend worker. Keep the backend bound to localhost; using a non-localhost bind exposes the API without authentication.

Backend environment (local dev):

  • Put local backend overrides in backend/.env (auto-loaded on backend startup via python-dotenv).
  • Keep secrets out of git; use backend/.env.example as the committed template.
  • Embedded-tutor keys commonly used (not needed for Copy context for AI):
    • TUTOR_LLM_ENABLED=false
    • TUTOR_PROVIDER=openai
    • TUTOR_MODEL=gpt-4o-mini
    • TUTOR_TEMPERATURE=0.2
    • TUTOR_MAX_TOKENS=500
    • OPENAI_API_KEY=...
    • OPENAI_BASE_URL=https://api.openai.com/v1

2) Start frontend

From frontend/:

npm install
npm run dev

Optional backend override:

NEXT_PUBLIC_BACKEND_URL=http://localhost:8000 npm run dev

For frontend-only local overrides, Next.js reads frontend/.env.local.

3) Open the app

  • Frontend: http://localhost:3000
  • OpenAPI docs: http://localhost:8000/docs
  • OpenAPI JSON: http://localhost:8000/openapi.json
  • Health: http://localhost:8000/api/health

Typical workflow

  1. Create a run from /.
  2. Monitor live training on /runs/<run_id>.
  3. Open /runs/<run_id>/latent for method/dimension/epoch exploration.
  4. If needed for older runs:
    • click Compute PCA 3D for this run
    • enable/recompute UMAP on run detail or latent explorer
    • in Neighbors panel click Compute <METHOD> neighbors for this run
    • in Projection controls set overlay to Density or Distortion; use compute buttons if artifacts are missing
  5. Use point selection + sample preview, Sampling panel, Neighbors/Extent controls, Mixture Health, and overlay modes for exploration.

Network Explorer workflow

  1. Open /runs/<run_id>/network.
  2. Choose epoch, layer, and dataset index in the Forward Pass / Feature Map Explorer.
  3. For GM-VAE, use Lock sample + Resample to control sampled-z reconstruction deterministically (sample_id).
  4. In Channel Inspector (spatial layers):
    • click activation tiles or channel map to inspect channel responses
    • click activation-map coordinates to probe (x,y) and view receptive-field overlay on input
    • for conv1-like layers, inspect Single convolution step to connect a probed activation to patch/kernel/product math
    • use Top activating examples to compute/show top-K dataset items for a selected encoder channel
  5. In latent/vector views, use Perturb + Decode to compare baseline (delta=0) vs perturbed reconstructions.

Interpreting overlays

  • Density and distortion overlays are visualization-space diagnostics.
  • In PCA space, geometry is linear and often easier to interpret globally.
  • In UMAP space, geometry is nonlinear:
    • UMAP is not a coordinate transform of latent probability space.
    • UMAP does not preserve volume or Gaussian structure.
    • UMAP primarily preserves local neighborhood relationships.
  • Practical guidance:
    • compare PCA vs UMAP overlays side-by-side
    • treat UMAP overlays as local-structure aids, not direct probability maps

API quick reference

Base prefix: /api

  • GET /health
  • POST /runs
  • GET /runs
  • GET /runs/{run_id}
  • POST /runs/{run_id}/cancel
  • GET /runs/{run_id}/events (SSE with Last-Event-ID replay)
  • GET /runs/{run_id}/artifacts
  • GET /runs/{run_id}/artifacts/{artifact_path}
  • GET /runs/{run_id}/dataset-samples/{sample_idx}
  • POST /runs/{run_id}/projections/umap (enable/backfill/recompute UMAP)
  • POST /runs/{run_id}/projections/pca3 (backfill PCA3)
  • POST /runs/{run_id}/neighbors/backfill?method=pca|umap
  • POST /runs/{run_id}/density/backfill?method=pca|umap
  • POST /runs/{run_id}/distortion/backfill?method=pca|umap
  • POST /runs/{run_id}/sample/posterior_local
  • POST /runs/{run_id}/sample/prior_component
  • POST /runs/{run_id}/network/forward
  • POST /runs/{run_id}/network/inspect
  • POST /runs/{run_id}/network/conv_step
  • POST /runs/{run_id}/network/kernel
  • POST /runs/{run_id}/network/perturb
  • POST /runs/{run_id}/network/topk
  • POST /runs/{run_id}/ui-state (publish a version-2 UI snapshot)
  • GET /runs/{run_id}/ui-state (latest stored snapshot for this run; 404 if absent)
  • GET /runs/{run_id}/ai-context (run/context aggregation; works without a snapshot)
  • GET /ui-state/current (last-active snapshot, or no_activity)
  • POST /tutor/ask (embedded tutor)
  • GET /tutor/sessions/{run_id}/{session_id}
  • POST /tutor/sessions/{run_id}/{session_id}/feedback

Troubleshooting

  • CORS / fetch errors:
    • ensure frontend is on http://localhost:3000 or http://127.0.0.1:3000
    • ensure backend is on http://localhost:8000
  • Neighbors artifact not ready yet for this method/epoch:
    • use the Neighbors panel backfill button or POST /runs/{id}/neighbors/backfill
  • Density artifact not ready / Distortion artifact not ready:
    • use the corresponding compute button in Latent Explorer
    • or call POST /runs/{id}/density/backfill / POST /runs/{id}/distortion/backfill
  • UMAP options unavailable:
    • enable UMAP from run detail or latent explorer (POST /runs/{id}/projections/umap)
  • 3D PCA unavailable for old runs:
    • trigger PCA3 backfill (POST /runs/{id}/projections/pca3)
  • command-not-found issues:
    • run uv sync for Python deps
    • run npm install in frontend/

Verification commands

Backend tests:

uv run pytest -q backend/tests

Frontend typecheck:

cd frontend && npx tsc --noEmit

Frontend tests (from frontend/):

npm test

About

A GUI autoencoder & GM-VAE tutorial / explorer, with an AI "tutor" feature as a guide

Topics

autoencodermachine-learningtutorialvisualization

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages