Skip to content

English · Español

Lab 02 — Walkthrough del selective-scan de Mamba (lab de lectura)

Objetivo: leer mamba-minimal de punta a punta. Anotar la lógica del selective-scan. Escribir un resumen de 1 página conectando las matemáticas (theory/03) con el código.

Tiempo estimado: 2–3 horas.

Prerrequisito: theory/03-state-space-models.md entendido. Fase 25 (interna de PyTorch) hecha. Borja sabe leer cómodamente una implementación de referencia de transformer.

Lo que produces

Un directorio experiments/36-mamba-walkthrough/ que contenga:

  • mamba-sha.txt — el SHA del repo mamba-minimal leído.
  • walkthrough.md — ~1 página de lectura anotada del núcleo del selective-scan (selective_scan_ref o equivalente).
  • state-update-diagram.mmd — diagrama mermaid de un paso del selective scan, anotado con shapes.
  • grammar-tutor-applicability.md — veredicto corto: ¿ayudaría Mamba al tutor de gramática?

TODOs

Bloque A — clonar mamba-minimal

git clone https://github.com/johnma2006/mamba-minimal /tmp/mamba-minimal
cd /tmp/mamba-minimal
git rev-parse HEAD > /home/overdrive/claude/lynx-cortex/experiments/36-mamba-walkthrough/mamba-sha.txt

mamba-minimal es una reimplementación deliberadamente pedagógica (~300 LOC) — mucho más fácil de leer que el repo oficial de Mamba que usa kernels CUDA. Lee la versión educativa.

Bloque B — recorrer el archivo

El archivo interesante es model.py. Céntrate en estas funciones / clases:

  • MambaBlock — el bloque de construcción (análogo a un bloque transformer).
  • selective_scan (o selective_scan_ref según la versión) — la recurrencia núcleo.

En walkthrough.md, escribe anotaciones que cubran:

  1. El paso de discretización. ¿Dónde está en el código la transición continuo-a-discreto (computar \(\bar{A}, \bar{B}\) a partir de \(A, B, \Delta\))? Cita números de línea.
  2. La selectividad. ¿Qué líneas hacen \(B, C, \Delta\) dependientes del input? (En contraste con S4, donde son parámetros fijos.)
  3. La actualización de estado. Traza \(h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t\) — encuentra la(s) línea(s) correspondiente(s).
  4. La proyección de salida. ¿Dónde está \(y_t = C_t h_t\)?
  5. La convolución. Mamba usa una conv 1D como preprocesador. ¿Por qué? (Pista: da suavizado de corto alcance antes del SSM.)

Cada anotación: 2-3 frases + cita de línea. ~5-8 anotaciones en total.

Bloque C — el diagrama de shapes

Dibuja el diagrama mermaid para un paso del selective-scan:

flowchart LR
    x[x_t : (B, D)] --> deltaP[Linear -> Δ_t : (B, D)]
    x --> Bproj[Linear -> B_t : (B, N)]
    x --> Cproj[Linear -> C_t : (B, N)]
    A_param[A : (D, N) param] -.discretize.-> Abar[Ā_t : (B, D, N)]
    deltaP --> Abar
    Bproj --> Bbar[B̄_t : (B, D, N)]
    deltaP --> Bbar
    h_prev[h_{t-1} : (B, D, N)] --> update[h_t = Ā_t · h_{t-1} + B̄_t · x_t]
    Abar --> update
    Bbar --> update
    update --> h_t[h_t : (B, D, N)]
    h_t --> output[y_t = C_t · h_t : (B, D)]
    Cproj --> output

Entrega como state-update-diagram.mmd. Anota el diagrama añadiendo una referencia "dónde en el código" al lado de cada caja (p. ej., "Δ projection: line 142").

Bloque D — la aplicabilidad al tutor de gramática

Escribe grammar-tutor-applicability.md (~200 palabras):

  • ¿Ayudaría Mamba al tutor de gramática? (Spoiler de theory/03: no.)
  • ¿Qué específicamente sobre la tarea del tutor de gramática hace a la attention estrictamente mejor que Mamba? (Respuesta: el acuerdo sujeto-verbo-tiempo requiere lookup preciso de un token pasado específico; Mamba comprime el pasado en un estado, la attention lee directamente.)
  • ¿Cuándo cogerías Mamba? (Respuesta: contexto muy largo, donde la KV cache se vuelve inviable.)
  • ¿Qué pasa con un enfoque híbrido (estilo Jamba)? ¿Podría ayudar una sola capa de attention + múltiples capas de Mamba? (Pista: probablemente no, a nuestro contexto máximo de 32 tokens. El cómputo de la attention a esta escala es despreciable.)

Restricciones

  • Sin entrenamiento de Mamba. Este es un lab de lectura. Levantar inferencia de Mamba está bien si quieres sentir la velocidad; entrenar está fuera de alcance (y es lento en CPU).
  • Sin copiar fuente. Cita rangos de líneas, no código completo. El walkthrough debe ser un resumen, no una transcripción.
  • Solo diagramas mermaid. Sin PNGs de herramientas externas — mantenlo diff-able.
  • Solo CPU, cero gasto en la nube.

Condiciones de parada

Has terminado cuando:

  1. experiments/36-mamba-walkthrough/{mamba-sha.txt, walkthrough.md, state-update-diagram.mmd, grammar-tutor-applicability.md} existen todos.
  2. walkthrough.md tiene ≥5 puntos anotados con citas de línea.
  3. El diagrama mermaid está entregado y renderiza correctamente.
  4. La nota de aplicabilidad al tutor de gramática responde a las cuatro preguntas.
  5. Puedes explicar, de memoria, "qué hace a Mamba 'selectivo'" en una frase.

Pista de último recurso

Si mamba-minimal ha cambiado respecto a la versión contra la que se escribió este lab: pinea la versión con el SHA en el momento del clone. Si los nombres de función han cambiado (selective_scan vs selective_scan_ref, etc.), usa el que corresponda al archivo actual.

Si el paso de discretización (Bloque B punto 1) es opaco: la fórmula está en theory/03-state-space-models.md §"Discretization". Compara la matemática con el código línea a línea. El código está haciendo exactamente la matemática, en ops de PyTorch.

Cuándo consultar solutions/

Tras commitear. La solución vive en solutions/02-mamba-walkthrough-ref.md — escrita al abrir la fase con la versión actual de mamba-minimal pineada. La referencia es un conjunto de elecciones de anotación con rangos de línea; las elecciones de Borja pueden diferir — la comparación es "¿qué se me escapó?", no "¿coincidí exactamente?".

Siguiente lab: lab/03-speculative-survey.md.