Проект: Нейросеть для распознавания LaTeX-формул (Image-to-LaTeX)

Это учебный проект, демонстрирующий создание и обучение нейросетевой модели для преобразования изображений математических формул в соответствующую им разметку LaTeX. Проект прошел через несколько итераций, включая классическую архитектуру Encoder-Decoder с Attention и современную архитектуру на основе Трансформера.

Основная цель проекта — не только получить работающую модель, но и изучить на практике ключевые концепции, проблемы и техники, возникающие при решении подобных задач Computer Vision и NLP.

---

Содержание

1. Архитектура модели
2. Данные
3. Ключевые этапы и изученные концепции
4. Как запустить
5. Результаты и выводы
6. Возможные пути улучшения

---

Архитектура модели

Финальная версия модели построена на архитектуре Encoder-Decoder с использованием Трансформера.

Encoder (Кодировщик)

• Задача: Извлечь из входного изображения набор визуальных признаков.
• Реализация:
  1. Используется предобученная сверточная нейросеть ResNet34 в качестве "хребта" для извлечения признаков (трансферное обучение).
  2. Верхние слои ResNet, отвечающие за классификацию, удалены.
  3. Добавлен Conv2d 1x1 слой для приведения размерности признаков к d_model, совместимой с Трансформером.

Decoder (Декодировщик)

• Задача: Сгенерировать последовательность LaTeX-токенов на основе признаков от энкодера.
• Реализация:
  1. Используется декодер Трансформера, состоящий из нескольких слоев nn.TransformerDecoderLayer.
  2. Positional Encoding добавляется к эмбеддингам токенов, чтобы модель знала о порядке в последовательности.
  3. Masked Self-Attention позволяет модели понимать контекст уже сгенерированной части формулы.
  4. Cross-Attention связывает текстовый контекст с визуальными признаками от энкодера.

---

Данные

• Источник: im2latex-100k — публичный датасет, содержащий ~100,000 пар "изображение-формула", сгенерированных из реальных научных статей.
• Предварительная обработка:
  1. Токенизация: LaTeX-строки разбиваются на отдельные токены (команды, символы).
  2. Построение словаря: Создается словарь всех уникальных токенов.
  3. Материализация данных: Для ускорения обучения все изображения и формулы заранее обрабатываются и кэшируются в виде тензоров с помощью torch.save(). Это позволяет избежать "бутылочного горлышка" при загрузке данных на лету.

---

Ключевые этапы и изученные концепции

В ходе проекта были изучены и решены следующие практические задачи:

• Отладка производительности:

  • Диагностика "узких мест" (bottlenecks) с помощью ручного профилирования и torch.profiler.
  • Решение проблем с медленной загрузкой данных: использование num_workers, RAM-диска и, наконец, полной предварительной обработки данных.

• Оптимизация обучения:

  • Применение смешанной точности (AMP) с torch.amp.autocast и GradScaler для ускорения вычислений на GPU.
  • Использование torch.compile() для JIT-компиляции модели и получения дополнительного прироста скорости.

• Улучшение процесса обучения:

  • Реализация Ранней Остановки (Early Stopping) для предотвращения переобучения и экономии времени.
  • Применение Файн-тюнинга (Fine-tuning) с уменьшенным learning_rate для "дожатия" качества модели после выхода на плато.

• Оценка модели:

  • Использование нескольких метрик: Valid Loss (для принятия решений), BLEU (для оценки схожести) и Exact Match (для оценки абсолютной точности).
  • Реализация "редкой" оценки дорогих метрик для ускорения цикла валидации.

• Улучшение инференса:

  • Сравнение "жадного" поиска и Beam Search. Анализ проблем (зацикливание, "галлюцинации") при "жадной" генерации.
  • Реализация пакетной генерации (predict_batch) для ускорения валидации.

---

Как запустить

Проект реализован в виде Jupyter Notebook, предназначенного для среды Google Colab.

1. Среда: Убедитесь, что выбран аппаратный ускоритель GPU (желательно A100 для высокой производительности).
2. Подготовка (Ячейки 1-4):
   • Запустите ячейки для установки зависимостей, определения конфигурации и скачивания данных.
   • При первом запуске будет выполнен долгий процесс предварительной обработки и кэширования данных на ваш Google Drive. Это может занять до 40 минут.
   • При всех последующих запусках данные будут быстро загружаться из кэша.
3. Обучение (Ячейка 7):
   • Запустите ячейку основного цикла обучения. Прогресс будет сохраняться в папку im2latex_checkpoints на вашем Google Drive.
4. Тестирование (Ячейки 8 и 9):
   • После завершения обучения запустите ячейки для тестирования лучшей сохраненной модели на примерах из тестового набора или на ваших собственных изображениях.

---

Результаты и выводы

• Финальная модель на основе Трансформера демонстрирует стабильное обучение, а Valid Loss неуклонно снижается на протяжении ~30 эпох.
• Модель отлично изучила синтаксис LaTeX, генерируя правдоподобные формулы.
• Основная сложность заключается в установлении точной связи между изображением и текстом, что требует длительного обучения.
• Оптимизация производительности, особенно предварительная обработка данных и torch.compile, критически важна для эффективного использования мощных GPU.

---

Возможные пути улучшения

• Реализовать Beam Search для декодера Трансформера, чтобы значительно улучшить качество генерируемых формул.
• Применить аугментацию данных для изображений, чтобы повысить устойчивость модели.
• Использовать планировщик скорости обучения (Learning Rate Scheduler) для автоматизации процесса файн-тюнинга.
• Исследовать более продвинутые архитектуры энкодера (например, Swin Transformer) или реализации внимания (например, Flash Attention).